Update knowledge_engine.py

knowledge_engine.py

@@ -1,61 +1,59 @@
 import os
-from pathlib import Path
-from langchain.document_loaders import TextLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.chains import RetrievalQA
-from langchain.llms import HuggingFaceHub
+from langchain.llms import HuggingFacePipeline
+from transformers import pipeline
 
 class KnowledgeManager:
-    def __init__(self, knowledge_dir="."):  # root dir by default
-        self.knowledge_dir = Path(knowledge_dir)
-        self.documents = []
-        self.embeddings = None
-        self.vectorstore = None
-        self.retriever = None
-        self.llm = None
+    def __init__(self, root_dir="."):
+        self.root_dir = root_dir
+        self.docsearch = None
         self.qa_chain = None
+        self.llm = None
 
-        self._load_documents()
-        if self.documents:
-            self._initialize_embeddings()
-            self._initialize_vectorstore()
-            self._initialize_llm()
-            self._initialize_qa_chain()
-
-    def _load_documents(self):
-        if not self.knowledge_dir.exists():
-            raise FileNotFoundError(f"Directory {self.knowledge_dir} does not exist.")
-        
-        files = list(self.knowledge_dir.glob("*.txt"))
-        if not files:
-            raise FileNotFoundError(f"No .txt files found in {self.knowledge_dir}. Please upload your knowledge base files in root.")
-        
-        for file in files:
-            loader = TextLoader(str(file))
-            self.documents.extend(loader.load())
-        
-        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-        self.documents = splitter.split_documents(self.documents)
+        self._initialize_llm()
+        self._initialize_embeddings()
+        self._load_knowledge_base()
+
+    def _initialize_llm(self):
+        local_pipe = pipeline("text2text-generation", model="google/flan-t5-small", max_length=256)
+        self.llm = HuggingFacePipeline(pipeline=local_pipe)
 
     def _initialize_embeddings(self):
-        self.embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        self.embeddings = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-large")
 
-    def _initialize_vectorstore(self):
-        self.vectorstore = FAISS.from_documents(self.documents, self.embeddings)
-        self.retriever = self.vectorstore.as_retriever()
+    def _load_knowledge_base(self):
+        # Find all .txt files in root directory
+        txt_files = [f for f in os.listdir(self.root_dir) if f.endswith(".txt")]
 
-    def _initialize_llm(self):
-        self.llm = HuggingFaceHub(repo_id="google/flan-t5-small", model_kwargs={"temperature":0, "max_length":256})
+        if not txt_files:
+            raise FileNotFoundError("No .txt files found in root directory.")
 
-    def _initialize_qa_chain(self):
-        self.qa_chain = RetrievalQA.from_chain_type(llm=self.llm, chain_type="stuff", retriever=self.retriever)
+        all_texts = []
+        for filename in txt_files:
+            path = os.path.join(self.root_dir, filename)
+            with open(path, "r", encoding="utf-8") as f:
+                content = f.read()
+                all_texts.append(content)
+
+        full_text = "\n\n".join(all_texts)
+
+        from langchain.text_splitter import RecursiveCharacterTextSplitter
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+        docs = text_splitter.create_documents([full_text])
+
+        self.docsearch = FAISS.from_documents(docs, self.embeddings)
+
+        self.qa_chain = RetrievalQA.from_chain_type(
+            llm=self.llm,
+            chain_type="stuff",
+            retriever=self.docsearch.as_retriever(),
+            return_source_documents=True,
+        )
 
     def ask(self, query):
         if not self.qa_chain:
-            return "Knowledge base not initialized properly."
-        return self.qa_chain.run(query)
-
-    def get_knowledge_summary(self):
-        return f"Loaded {len(self.documents)} document chunks from {self.knowledge_dir}"
+            raise ValueError("Knowledge base not initialized.")
+        result = self.qa_chain(query)
+        return result['result']