Update utils/llm_analyzer.py

utils/llm_analyzer.py

@@ -14,19 +14,28 @@ class LLMAnalyzer:
     
     def __init__(self):
         """Initialize Gemma model"""
+        import os
         print("Loading Gemma model...")
-        self.tokenizer = AutoTokenizer.from_pretrained(GEMMA_MODEL)
+        
+        # Get token from environment
+        hf_token = os.environ.get("HF_TOKEN", None)
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            GEMMA_MODEL,
+            token=hf_token
+        )
         self.model = AutoModelForCausalLM.from_pretrained(
             GEMMA_MODEL,
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto"
+            device_map="auto",
+            token=hf_token
         )
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Gemma loaded on {self.device}!")
     
     def generate_response(self, prompt: str, max_length: int = LLM_MAX_LENGTH) -> str:
         """
-        Generate response from Gemma
+        Generate response from LLM
         
         Args:
             prompt: Input prompt
@@ -35,8 +44,12 @@ class LLMAnalyzer:
         Returns:
             Generated text
         """
-        # Format prompt for Gemma
-        formatted_prompt = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
+        # Format prompt (works for both Gemma and Zephyr)
+        if "gemma" in GEMMA_MODEL.lower():
+            formatted_prompt = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
+        else:
+            # Zephyr format
+            formatted_prompt = f"<|user|>\n{prompt}</s>\n<|assistant|>\n"
         
         inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.device)
         
@@ -55,6 +68,12 @@ class LLMAnalyzer:
         # Extract only the model's response
         if "<start_of_turn>model" in response:
             response = response.split("<start_of_turn>model")[-1].strip()
+        elif "<|assistant|>" in response:
+            response = response.split("<|assistant|>")[-1].strip()
+        
+        # Remove the original prompt if still present
+        if prompt in response:
+            response = response.replace(prompt, "").strip()
         
         return response