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 pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 class KnowledgeManager:
@@ -12,23 +12,38 @@ 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):
-        # Load local text2text model using HuggingFace pipeline (FLAN-T5 small)
-        local_pipe = pipeline("text2text-generation", model="google/flan-t5-small", max_length=500)
-        self.llm = HuggingFacePipeline(pipeline=local_pipe)
+        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)
 
     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:
@@ -38,18 +53,16 @@ class KnowledgeManager:
         for filename in txt_files:
             path = os.path.join(self.root_dir, filename)
             with open(path, "r", encoding="utf-8") as f:
-                all_texts.append(f.read())
+                content = f.read()
+                all_texts.append(content)
 
         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",
@@ -61,4 +74,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"]