Update knowledge_engine.py

knowledge_engine.py

@@ -3,7 +3,7 @@ from langchain.vectorstores import FAISS
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
 from langchain.llms import HuggingFacePipeline
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import pipeline
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 class KnowledgeManager:
@@ -12,38 +12,23 @@ class KnowledgeManager:
         self.docsearch = None
         self.qa_chain = None
         self.llm = None
+        self.embeddings = None
 
         self._initialize_llm()
         self._initialize_embeddings()
         self._load_knowledge_base()
 
     def _initialize_llm(self):
-        model_id = "tiiuae/falcon-7b-instruct"
-
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            trust_remote_code=True,
-            torch_dtype="auto",  # Will use float16 on GPU, float32 on CPU
-            device_map="auto"
-        )
-
-        falcon_pipeline = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.95,
-            repetition_penalty=1.1
-        )
-
-        self.llm = HuggingFacePipeline(pipeline=falcon_pipeline)
+        # Load local text2text model using HuggingFace pipeline (FLAN-T5 small)
+        local_pipe = pipeline("text2text-generation", model="google/flan-t5-small", max_length=1024)
+        self.llm = HuggingFacePipeline(pipeline=local_pipe)
 
     def _initialize_embeddings(self):
+        # Use general-purpose sentence transformer
         self.embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
     def _load_knowledge_base(self):
+        # Automatically find all .txt files in the root directory
         txt_files = [f for f in os.listdir(self.root_dir) if f.endswith(".txt")]
 
         if not txt_files:
@@ -53,16 +38,18 @@ class KnowledgeManager:
         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)
+                all_texts.append(f.read())
 
         full_text = "\n\n".join(all_texts)
 
+        # Split text into chunks for embedding
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
         docs = text_splitter.create_documents([full_text])
 
+        # Create FAISS vector store
         self.docsearch = FAISS.from_documents(docs, self.embeddings)
 
+        # Build the QA chain
         self.qa_chain = RetrievalQA.from_chain_type(
             llm=self.llm,
             chain_type="stuff",
@@ -74,4 +61,4 @@ class KnowledgeManager:
         if not self.qa_chain:
             raise ValueError("Knowledge base not initialized.")
         result = self.qa_chain(query)
-        return result["result"]
+        return result['result']