Update utils.py

utils.py

@@ -1,37 +1,108 @@
-import gradio as gr
+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
+from langchain.llms import OpenAI
+from langchain_openai import ChatOpenAI
+from langchain.chains import RetrievalQA
 import os
-from langchain.retrievers import EnsembleRetriever
-from utils import *
-import requests
-from pyvi import ViTokenizer, ViPosTagger
-import time
-from transformers import AutoTokenizer, AutoModelForQuestionAnswering
-import torch
 
-retriever = load_the_embedding_retrieve(is_ready=False, k=3)
-bm25_retriever = load_the_bm25_retrieve(k=3)
+def split_with_source(text, source):
+    splitter = CharacterTextSplitter(
+        separator = "\n",
+        chunk_size = 256,
+        chunk_overlap  = 0,
+        length_function = len,
+        add_start_index = True,
+    )
+    documents = splitter.create_documents([text])
+    print(documents)
+    for doc in documents:
+        doc.metadata["source"] = source
+        # print(doc.metadata)
 
-ensemble_retriever = EnsembleRetriever(
-    retrievers=[bm25_retriever, retriever], weights=[0.5, 0.5]
-)
+    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
 
-def greet2(quote):
+    # 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)
 
-    qa_chain = get_qachain(retriever=ensemble_retriever)
+    # Đế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:
+            # Xử lý bằng text_spliter
+            # 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
+
+            ##Xử lý mỗi khi xuống dòng
+            # for line in file:
+            #     # Loại bỏ khoảng trắng thừa và ký tự xuống dòng ở đầu và cuối mỗi dòng
+            #     line = line.strip()
+            #     documents.append(Document(page_content=line, metadata={"source": i}))
+    print(documents)
+    return documents
+
+def load_the_embedding_retrieve(is_ready = False, k = 3, model= 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'):
+    embeddings = HuggingFaceEmbeddings(model_name=model)
+    if is_ready:
+        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()
+        print(type(documents))
+        retriever = Chroma.from_documents(documents, embeddings).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
+
+def get_qachain(llm_name = "gpt-3.5-turbo-0125", chain_type = "stuff", retriever = None, return_source_documents = True):
+    llm = ChatOpenAI(temperature=0,
+                     model_name=llm_name)
+    return RetrievalQA.from_chain_type(llm=llm,
+                                  chain_type=chain_type,
+                                  retriever=retriever,
+                                  return_source_documents=return_source_documents)
+
+def process_llm_response(llm_response):
+    print(llm_response['result'])
+    print('\n\nSources:')
+    for source in llm_response["source_documents"]:
+        print(source.metadata['source'])
 
-    prompt = os.environ['PROMPT']
 
-    qa_chain.combine_documents_chain.llm_chain.prompt.messages[0].prompt.template = prompt
 
-    llm_response = qa_chain(quote)
 
-    return llm_response['result']
 
 
-if __name__ == "__main__":
-    quote = "Địa chỉ nhà trường?"
 
-    iface = gr.Interface(fn=greet2, inputs="text", outputs="text")
-    iface.launch()