updated aggent

agent.py

@@ -1,4 +1,4 @@
-# agent.py — CPU-only version using Microsoft Phi-2
+# agent.py — CPU-ready, optimized for speed with TinyLlama
 import os
 import fitz
 import faiss
@@ -7,18 +7,27 @@ import sqlite3
 import numpy as np
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from sentence_transformers import SentenceTransformer
+from torch.quantization import quantize_dynamic
 
 class CodingAgent:
     def __init__(self):
-        # Load Microsoft Phi-2 (non-chat, CPU-friendly)
-        model_id = "microsoft/phi-2"
+        # Load TinyLlama (chat-optimized, CPU-compatible)
+        model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
         self.tokenizer = AutoTokenizer.from_pretrained(model_id)
         self.model = AutoModelForCausalLM.from_pretrained(model_id)
+
+        # Apply dynamic quantization to speed up inference on CPU
+        self.model = quantize_dynamic(
+            self.model, {torch.nn.Linear}, dtype=torch.qint8
+        )
+
+        # Text-generation pipeline, limit token output
         self.llm = pipeline(
             "text-generation",
             model=self.model,
             tokenizer=self.tokenizer,
-            max_new_tokens=512,
+            max_new_tokens=128,
+            do_sample=False,
             device=-1
         )
 
@@ -62,7 +71,7 @@ class CodingAgent:
         self.id_map.extend(range(len(self.docs)-len(chunks), len(self.docs)))
         return f"Added {len(chunks)} chunks."
 
-    def retrieve_context(self, query, top_k=3):
+    def retrieve_context(self, query, top_k=1):  # reduce to 1 chunk for speed
         if self.index.ntotal == 0:
             return ""
         query_emb = self.embed_chunks([query])[0]
@@ -77,9 +86,12 @@ class CodingAgent:
             return f"[From memory] {result[0]}"
 
         context = self.retrieve_context(query)
-        prompt = f"Q: {query}\nContext:\n{context}\nA:"
-        result = self.llm(prompt)[0]['generated_text'].split("A:")[-1].strip()
+        prompt = f"### User:\n{query}\n\n### Context:\n{context}\n\n### Assistant:"
+        # Generate with limited tokens
+        out = self.llm(prompt)[0]['generated_text']
+        result = out.split("### Assistant:")[-1].strip()
 
+        # Save to memory
         self.conn.execute("INSERT INTO memory (query, response) VALUES (?, ?)", (query, result))
         self.conn.commit()
         return result