speed improvement

service/knowledge_base.py

@@ -0,0 +1,25 @@
+
+
+from service.data_loader_service import CSVDataLoader
+from service.embedded_service import EmbeddingService
+from service.vector_store_service import VectorStoreService
+
+
+class KnowledgeBase:
+    def __init__(self, csv_path: str):
+        self.embedder = EmbeddingService()
+
+        loader = CSVDataLoader(csv_path)
+        qa_pairs = loader.load_qa_pairs()
+
+        self.documents = [
+            f"Question: {p['question']}\nAnswer: {p['answer']}"
+            for p in qa_pairs
+        ]
+
+        self.embeddings = self.embedder.embed(self.documents)
+
+        self.vector_store = VectorStoreService(
+            self.embeddings,
+            self.documents
+        )

service/llm_service.py

@@ -1,24 +1,25 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-
 class LLMService:
     def __init__(self):
         self.model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-        # Tokenizer
+        # 🔥 Limit CPU threads (CRITICAL)
+        torch.set_num_threads(4)          # try 2–6 depending on CPU
+        torch.set_num_interop_threads(1)
+
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.model_name,
             use_fast=True
         )
 
-        # Load model in FP32 on CPU
         model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
             torch_dtype=torch.float32
         )
 
-        # 🔥 CPU INT8 dynamic quantization
+        # 🔥 INT8 dynamic quantization
         self.model = torch.quantization.quantize_dynamic(
             model,
             {torch.nn.Linear},
@@ -27,22 +28,22 @@ class LLMService:
 
         self.model.eval()
 
-        # Optional sanity check
-        print("LLM loaded with dtype:", next(self.model.parameters()).dtype)
+        print("LLM loaded | dtype:", next(self.model.parameters()).dtype)
 
     def generate(self, prompt: str) -> str:
         inputs = self.tokenizer(
             prompt,
             return_tensors="pt",
             truncation=True,
-            max_length=1024
+            max_length=768      # ⬅️ smaller context = faster attention
         )
 
         with torch.no_grad():
             output = self.model.generate(
                 **inputs,
-                max_new_tokens=120,      # ⬅️ faster + enough
-                do_sample=False,         # ⬅️ HUGE speed win
+                max_new_tokens=120,
+                do_sample=False,
+                use_cache=True,     # ⬅️ IMPORTANT
                 eos_token_id=self.tokenizer.eos_token_id
             )
 
@@ -54,12 +55,10 @@ class LLMService:
         return self._clean(text)
 
     def _clean(self, text: str) -> str:
-        # Extract content AFTER <|assistant|>
         if "<|assistant|>" in text:
             text = text.split("<|assistant|>")[-1]
 
-        # Stop if model continues roles
-        for stop in ["<|system|>", "<|user|>"]:
+        for stop in ("<|system|>", "<|user|>"):
             if stop in text:
                 text = text.split(stop)[0]
 

service/vector_store_service.py

@@ -2,11 +2,13 @@ import numpy as np
 
 class VectorStoreService:
     def __init__(self, embeddings, documents):
-        self.embeddings = np.array(embeddings)
+        self.embeddings = np.array(embeddings, dtype="float32")
         self.documents = documents
 
-    def search(self, query_embedding, top_k: int = 3):
-        query = np.array(query_embedding)
+    def search(self, query_embedding, top_k=3):
+        query = np.array(query_embedding, dtype="float32")
+
         scores = np.dot(self.embeddings, query)
-        top_idx = scores.argsort()[-top_k:][::-1]
-        return [self.documents[i] for i in top_idx]
+        top_indices = scores.argsort()[-top_k:][::-1]
+
+        return [self.documents[i] for i in top_indices]