utils.py

from langchain_community.document_loaders import TextLoader
from langchain_community.docstore.document import Document
from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.retrievers import BM25Retriever

import os

def split_with_source(text, source):
    splitter = CharacterTextSplitter(
        separator = "\n",
        chunk_size = 256,
        chunk_overlap  = 72,
        length_function = len,
        add_start_index = True,
    )
    documents = splitter.create_documents([text])
    for doc in documents:
        doc.metadata["source"] = source
        # print(doc.metadata)

    return documents


def count_files_in_folder(folder_path):
    # Kiểm tra xem đường dẫn thư mục có tồn tại không
    if not os.path.isdir(folder_path):
        print("Đường dẫn không hợp lệ.")
        return None

    # Sử dụng os.listdir() để lấy danh sách các tập tin và thư mục trong thư mục
    files = os.listdir(folder_path)

    # Đếm số lượng tập tin trong danh sách
    file_count = len(files)

    return file_count

def get_document_from_raw_text():
    documents = [Document(page_content="", metadata={'source': 0})]
    files = os.listdir(os.path.join(os.getcwd(), "raw_data"))
    # print(files)
    for i in files:
        file_path = i
        with open(os.path.join(os.path.join(os.getcwd(), "raw_data"),file_path), 'r', encoding="utf-8") as file:
            # Tiền xử lý văn bản
            content = file.read().replace('\n\n', "\n")
            # content = ''.join(content.split('.'))
            new_doc = content
            texts = split_with_source(new_doc, i)
            documents = documents + texts

    return documents

def load_the_embedding_retrieve(is_ready = False, k = 3, model= 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'):
    if is_ready:
        embeddings = HuggingFaceEmbeddings(model_name=model)
        retriever = Chroma(persist_directory=os.path.join(os.getcwd(), "Data"), embedding_function=embeddings).as_retriever(
            search_kwargs={"k": k}
        )
    else:

        documents = get_document_from_raw_text()

        retriever = Chroma.from_documents(documents, embedding=model).as_retriever(
            search_kwargs={"k": k}
        )

    return retriever

def load_the_bm25_retrieve(k = 3):
    documents = get_document_from_raw_text()
    bm25_retriever = BM25Retriever.from_documents(documents)
    bm25_retriever.k = k

    return bm25_retriever