<s> response issue

service/llm_service.py

@@ -9,11 +9,13 @@ class LLMService:
         # torch.set_num_threads(...)
         # torch.set_num_interop_threads(...)
 
+        # Tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.model_name,
             use_fast=True
         )
 
+        # Model in FP32 on CPU
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
             torch_dtype=torch.float32
@@ -23,30 +25,37 @@ class LLMService:
 
         print("LLM loaded | dtype:", next(self.model.parameters()).dtype)
 
-    def generate(self, prompt: str) -> str:
+    def generate(self, user_query: str) -> str:
+        # Wrap user input with role tokens for TinyLlama
+        prompt = f"<|user|>{user_query}<|assistant|>"
+
         inputs = self.tokenizer(
             prompt,
             return_tensors="pt",
             truncation=True,
-            max_length=640      # ⬅️ important
+            max_length=640  # maintain context window
         )
 
         with torch.no_grad():
             output = self.model.generate(
                 **inputs,
-                max_new_tokens=96,     # ⬅️ enough for helpdesk
-                do_sample=False,
-                eos_token_id=self.tokenizer.eos_token_id
+                max_new_tokens=120,        # slightly higher for complete answer
+                do_sample=False,           # deterministic + faster
+                eos_token_id=self.tokenizer.eos_token_id,
+                use_cache=True
             )
 
-        text = self.tokenizer.decode(
-            output[0],
-            skip_special_tokens=False
-        )
+        # Decode and remove special tokens
+        text = self.tokenizer.decode(output[0], skip_special_tokens=True)
 
         return self._clean(text)
 
     def _clean(self, text: str) -> str:
+        """
+        Maintains your previous cleaning logic:
+        - Extract after <|assistant|>
+        - Stop at <|system|> or <|user|>
+        """
         if "<|assistant|>" in text:
             text = text.split("<|assistant|>")[-1]