Update knowledge_engine.py

knowledge_engine.py

@@ -3,7 +3,6 @@ import pickle
 from typing import List, Dict, Any
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor
-from sentence_transformers import SentenceTransformer  
 
 from config import Config
 
@@ -23,7 +22,8 @@ from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
 from langchain.retrievers import BM25Retriever
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFaceEndpoint
+from transformers import AutoTokenizer, pipeline
+from langchain_community.llms import HuggingFacePipeline
 
 class KnowledgeManager:
     def __init__(self):
@@ -38,14 +38,17 @@ class KnowledgeManager:
 
     def _init_llm(self):
         print("[i] Using Hugging Face LLM")
-        return HuggingFaceEndpoint(
-            endpoint_url="https://api-inference.huggingface.co/models/tiiuae/falcon-7b-instruct",
-            task="text-generation",
-            huggingfacehub_api_token=hf_token,
+        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,
-            do_sample=True
+            device_map="auto"
         )
+        return HuggingFacePipeline(pipeline=pipe)
 
     def _init_retrievers(self):
         faiss_index_path = Config.VECTOR_STORE_PATH / "index.faiss"