update

rag.py

@@ -1,63 +1,12 @@
 from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import Chroma
-from langchain.document_loaders import TextLoader, DirectoryLoader
-import os
-import re
-from sentence_transformers.cross_encoder import CrossEncoder
-import numpy as np
-from langchain.schema.retriever import BaseRetriever, Document
-from typing import List
-from langchain.callbacks.manager import CallbackManagerForRetrieverRun
-from langchain.vectorstores import VectorStore
-from llm import URALLM
 from langchain.prompts import PromptTemplate
+from utility import load_data, process_data, CustomRetriever
 
-# Get role for  passage document
-def get_role(document):
-    """
-    Get role for student.
-    """
-    # Tìm kiếm các từ khóa liên quan đến vai trò học viên trong document.
-    keywords = [
-        "sinh viên",
-        "đại học",
-        "học viên",
-        "thạc sĩ",
-        "nghiên cứu sinh",
-        "tiến sĩ",
-    ]
-    role = []
-    for keyword in keywords:
-        if keyword in document.metadata['source'].lower():
-            role.append(keyword)
-    return ", ".join(role)
-
-def processing_data(data_path):
-    folders = os.listdir(data_path)
-
-    dir_loaders = []
-
-    # Add the documents to the project
-    for folder in folders:
-        dir_loader = DirectoryLoader((os.path.join(data_path, folder)), loader_cls=TextLoader)
-        dir_loaders.append(dir_loader)
-
-    # Load the text files.
-    loaded_documents = []
-    for dir_loader in dir_loaders:
-        loaded_documents.append(dir_loader.load())
-
-    data = []
-    for i in range(len(loaded_documents)):
-        for j in range(len(loaded_documents[i])):
-            data.append(loaded_documents[i][j])
-
-    # Final data prepare for vector database
-    for document in data:
-        role = get_role(document)
-        document.metadata['role'] = role
-    
-    return data
+
+data1 = load_data('raw_data/sv')
+data2 = load_data('raw_data/thacsi')
+data3 = load_data('raw_data/tiensi')
+data = data1 + data2 + data3
 
 # Embedding model
 embedding = HuggingFaceEmbeddings(
@@ -65,96 +14,200 @@ embedding = HuggingFaceEmbeddings(
     model_kwargs={"device": "cpu"}
 )
 
-# embedding = HuggingFaceEmbeddings(
-#     model_name="sentence-transformers/all-MiniLM-L6-v2",
-#     model_kwargs={"device": "cpu"}
-# )
-
-# Vector database
-data_path = 'raw_data'
-persist_directory = 'vector_db'
-vectordb = Chroma.from_documents(
-    documents=processing_data(data_path),
-    embedding=embedding,
-    persist_directory=persist_directory
-)
+# The splitter to use to create smaller chunks
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 
-class CustomRetriever(BaseRetriever):
-    vectorstores:Chroma
-    retriever:vectordb.as_retriever()
+child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=400)
 
-    def _get_relevant_documents(
-        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
-    ) -> List[Document]:
-        # Use your existing retriever to get the documents
-        documents = self.retriever.get_relevant_documents(query, callbacks=run_manager.get_child())
+#####################################################################
 
-        # Get page content
-        docs_content = []
-        for i in range(len(documents)):
-          docs_content.append(documents[i].page_content)
+vectorstore1, retriever1 = process_data(data, child_text_splitter, embedding, "data")
+vectorstore2, retriever2 = process_data(data2, child_text_splitter, embedding, "data2")
+vectorstore3, retriever3 = process_data(data3, child_text_splitter, embedding, "data3")
 
-        model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
+##############################################################################
 
-        # So we create the respective sentence combinations
-        sentence_combinations = [[query, document] for document in docs_content]
+ANYSCALE_API_BASE = "credential-1711634141163"
+ANYSCALE_API_KEY = "esecret_chitz7splr5ut6vfvqpn72itd3"
+ANYSCALE_MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"
+# ANYSCALE_MODEL_NAME = "meta-llama/Llama-3-8b-chat-hf"
+# ANYSCALE_MODEL_NAME = "google/gemma-7b-it"
+# ANYSCALE_MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"
+# ANYSCALE_MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
-        # Compute the similarity scores for these combinations
-        similarity_scores = model.predict(sentence_combinations)
+import os
 
-        # Sort the scores in decreasing order
-        sim_scores_argsort = reversed(np.argsort(similarity_scores))
+os.environ["ANYSCALE_API_BASE"] = ANYSCALE_API_BASE
+os.environ["ANYSCALE_API_KEY"] = ANYSCALE_API_KEY
 
-        # Store the rerank document in new list
-        docs = []
-        for idx in sim_scores_argsort:
-          docs.append(documents[idx])
+from langchain.chains import LLMChain
+from langchain_community.llms import Anyscale
+from langchain_core.prompts import PromptTemplate
+from langchain_community.chat_models import ChatAnyscale
 
-        docs_top_4 = docs[0:4]
+# llm = Anyscale(model_name=ANYSCALE_MODEL_NAME)
+llm= ChatAnyscale(model_name=ANYSCALE_MODEL_NAME, temperature=0)
 
-        return docs_top_4
-    
-llm = URALLM()
-custom_retriever = CustomRetriever(vectorstores = vectordb,retriever = vectordb.as_retriever(search_kwargs={"k": 50}))
+#####################################################################
 
-# Build prompt
-template = """[INST] <<SYS>>
+from langchain_openai.llms.azure import AzureOpenAI
+llm_openai = AzureOpenAI(
+    deployment_name="gpt-35-turbo-instruct",
+    # deployment_name="gpt-35-turbo-16k",
+    api_key = 'c90c0e7fb1894a898c56123580a6ee3e',
+    api_version = "2023-09-15-preview",
+    azure_endpoint = "https://bkchatbot.openai.azure.com/",
+    temperature=0.0,
+    max_tokens=500
+)
+
+##########################################################################
+
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
 
-Bạn là một chatbot hỗ trợ các quy định học vụ của trường Đại học Bách Khoa - ĐHQG TP.HCM.
-Bạn sử dụng văn bản được cung cấp để trả lời câu hỏi cho người dùng.
-Không sử dụng bất kỳ thông tin nào khác ngoài văn bản đã cho.
-Trả lời đầy đủ và ngắn gọn nhất có thể.
-Không đề cập tên riêng trong câu trả lời.
-Không chứa các ký tự: "-/", "]", "/", "-" trong câu trả lời.
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+# Build prompt
+from langchain.prompts import PromptTemplate
+template ="""
+Trả lời câu hỏi dựa trên những quy định được cung cấp, tổng hợp thông tin và đưa ra câu trả lời ngắn gọn và đầy đủ cuối cùng.
+Không cần ghi chú và trích dẫn nguồn thông tin đã tham khảo trong câu trả lời.
+Câu trả lời nên bắt đầu bằng: "Theo quy định của Trường ĐH Bách Khoa Tp.HCM, ..."
+Nếu trong quy văn bản không có thông tin cho câu trả lời, vui lòng thông báo: "Xin lỗi, tôi không có thông tin cho câu hỏi này!"
 
-<</SYS>>
+Quy định: {context}
 
-Văn bản: {context}
 Câu hỏi: {question}
+
 Câu trả lời:
-[/INST]"""
-QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"],template=template,)
-
-# Run chain
-from langchain.chains import RetrievalQA
-
-qa_chain = RetrievalQA.from_chain_type(llm,
-                                       verbose=False,
-                                       # retriever=vectordb.as_retriever(),
-                                       retriever=custom_retriever,
-                                       return_source_documents=True,
-                                       chain_type_kwargs={"prompt": QA_CHAIN_PROMPT})
-
-def remove_special_characters(text):
-    text = text.replace('].', '')
-    text = text.replace('/.', '')
-    text = text.replace('/.-', '')
-    text = text.replace('-', '')
-    return text
+"""
+QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"],template=template)
+
+#############################################################################
+
+from langchain_core.runnables import RunnableParallel
+
+rag_chain_from_docs = (
+    RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))
+    | QA_CHAIN_PROMPT
+    | llm
+    | StrOutputParser()
+)
+
+###############################################################################
+
+from langchain.prompts import ChatPromptTemplate
+
+# Multi Query: Different Perspectives
+template = """
+### Hãy tạo ra thêm các truy vấn tìm kiếm tương đương ngữ nghĩa với một câu hỏi ban đầu.
+Kết quả hiển thị dạng list gồm câu hỏi ban đầu và 2 câu hỏi thay thế.
+
+### Câu hỏi ban đầu: {question}
+### Kết quả:
+
+"""
+prompt_perspectives = ChatPromptTemplate.from_template(template)
+
+from langchain_core.output_parsers import StrOutputParser
+# from langchain_openai import ChatOpenAI
+
+generate_queries = (
+    prompt_perspectives
+    | llm_openai
+    | StrOutputParser()
+    | (lambda x: x.split("\n"))
+)
+
+#########################################################################################
+
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+
+# initialize the bm25 retriever and chroma retriever
+bm25_retriever1 = BM25Retriever.from_documents(data, k=25)
+ensemble_retriever1 = EnsembleRetriever(retrievers=[bm25_retriever1, retriever1], weights=[0.5, 0.5])
+
+bm25_retriever2 = BM25Retriever.from_documents(data2, k=25)
+ensemble_retriever2 = EnsembleRetriever(retrievers=[bm25_retriever2, retriever2], weights=[0.5, 0.5])
+
+bm25_retriever3 = BM25Retriever.from_documents(data3, k=25)
+ensemble_retriever3 = EnsembleRetriever(retrievers=[bm25_retriever3, retriever3], weights=[0.5, 0.5])
+
+#########################################################################################
+
+custom_retriever1 = CustomRetriever(retriever = ensemble_retriever1)
+custom_retriever2 = CustomRetriever(retriever = ensemble_retriever2)
+custom_retriever3 = CustomRetriever(retriever = ensemble_retriever3)
+
+multiq_chain1 = generate_queries | custom_retriever1
+multiq_chain2 = generate_queries | custom_retriever2
+multiq_chain3 = generate_queries | custom_retriever3
+
+rag_chain_with_source1 = RunnableParallel(
+    {"context": multiq_chain1, "question": RunnablePassthrough()}
+).assign(answer=rag_chain_from_docs)
+
+rag_chain_with_source2 = RunnableParallel(
+    {"context": multiq_chain2 , "question": RunnablePassthrough()}
+).assign(answer=rag_chain_from_docs)
+
+rag_chain_with_source3 = RunnableParallel(
+    {"context": multiq_chain3, "question": RunnablePassthrough()}
+).assign(answer=rag_chain_from_docs)
+
+############################################################################################
+
+from flashtext import KeywordProcessor
+keyword_processor = KeywordProcessor()
+# keyword_processor.add_keyword(<unclean name>, <standardised name>)
+keyword_processor.add_keyword('thạc sĩ')
+keyword_processor.add_keyword('học viên')
+keyword_processor.add_keyword('nghiên cứu sinh')
+keyword_processor.add_keyword('tiến sĩ')
+
+################################################################################
+
+rag_chain = [rag_chain_with_source1, rag_chain_with_source2, rag_chain_with_source3]
+
+###################################################################################
 
 def rag(question: str) -> str:
-    # call QA chain
-    response = qa_chain({"query": question})
 
-    return remove_special_characters(response["result"])
+    keywords_found = keyword_processor.extract_keywords(question)
+    if 'thạc sĩ' in keywords_found or 'học viên' in keywords_found:
+      response = rag_chain[1].invoke(question)
+    elif 'nghiên cứu sinh' in keywords_found or 'tiến sĩ' in keywords_found:
+      response = rag_chain[2].invoke(question)
+    else:
+      response = rag_chain[0].invoke(question)
+    
+    return response['answer']
+
+###################################################################################
+
+
+# # Run chain
+# from langchain.chains import RetrievalQA
+
+# qa_chain = RetrievalQA.from_chain_type(llm,
+#                                        verbose=False,
+#                                        # retriever=vectordb.as_retriever(),
+#                                        retriever=custom_retriever,
+#                                        return_source_documents=True,
+#                                        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT})
+
+# def remove_special_characters(text):
+#     text = text.replace('].', '')
+#     text = text.replace('/.', '')
+#     text = text.replace('/.-', '')
+#     text = text.replace('-', '')
+#     return text
+
+# def rag(question: str) -> str:
+#     # call QA chain
+#     response = qa_chain({"query": question})
+
+#     return remove_special_characters(response["result"])