Update knowledge_engine.py

knowledge_engine.py

@@ -1,20 +1,10 @@
 import os
 import pickle
-from typing import List, Dict, Any
+from typing import Dict, Any
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor
 
 from config import Config
-
-# Get token from environment variable
-hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-
-if not hf_token:
-    raise ValueError("HUGGINGFACEHUB_API_TOKEN not found in environment variables. Please set it in your Space secrets.")
-
-os.environ["HUGGINGFACEHUB_API_TOKEN"] = hf_token
-
-# Core ML/AI libraries
 from langchain_community.document_loaders import TextLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
@@ -22,33 +12,30 @@ from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
 from langchain.retrievers import BM25Retriever
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from transformers import AutoTokenizer, pipeline
-from langchain_community.llms import HuggingFacePipeline
+from langchain_community.llms import HuggingFaceHub
 
 class KnowledgeManager:
     def __init__(self):
         Config.setup_dirs()
         self.embeddings = self._init_embeddings()
         self.vector_db, self.bm25_retriever = self._init_retrievers()
-        self.qa_chain = self._create_moe_qa_chain()
+        self.qa_chain = self._create_qa_chain()
 
     def _init_embeddings(self):
         print("[i] Using Hugging Face embeddings")
         return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
     def _init_llm(self):
-        print("[i] Using Hugging Face LLM")
-        model_id = "tiiuae/falcon-7b-instruct"
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        pipe = pipeline(
-            "text-generation",
-            model=model_id,
-            tokenizer=tokenizer,
-            temperature=0.1,
-            max_new_tokens=512,
-            device_map="auto"
+        print("[i] Using HuggingFaceHub with Mistral-7B")
+        return HuggingFaceHub(
+            repo_id="mistralai/Mistral-7B-Instruct-v0.1",
+            huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),
+            model_kwargs={
+                "temperature": 0.1,
+                "max_new_tokens": 512,
+                "do_sample": True
+            }
         )
-        return HuggingFacePipeline(pipeline=pipe)
 
     def _init_retrievers(self):
         faiss_index_path = Config.VECTOR_STORE_PATH / "index.faiss"
@@ -124,7 +111,7 @@ class KnowledgeManager:
 
         return vector_results + bm25_results
 
-    def _create_moe_qa_chain(self):
+    def _create_qa_chain(self):
         if not self.vector_db or not self.bm25_retriever:
             return None