Switch to Llama-2-7b-chat-hf model with proper chat formatting

app.py

@@ -85,8 +85,8 @@ class TextilindoAI:
     
     def __init__(self):
         self.api_key = os.getenv('HUGGINGFACE_API_KEY')
-        # Use a more accessible model for free tier
-        self.model = os.getenv('DEFAULT_MODEL', 'microsoft/DialoGPT-small')
+        # Use Llama model for better performance
+        self.model = os.getenv('DEFAULT_MODEL', 'meta-llama/Llama-2-7b-chat-hf')
         self.system_prompt = self.load_system_prompt()
         
         if not self.api_key:
@@ -144,8 +144,11 @@ Minimum purchase is 1 roll (67-70 yards)."""
             return self.get_mock_response(user_message)
         
         try:
-            # For DialoGPT, use a simpler prompt format
-            if "dialogpt" in self.model.lower():
+            # For Llama models, use the proper chat format
+            if "llama" in self.model.lower():
+                # Llama 2 chat format
+                prompt = f"<s>[INST] <<SYS>>\n{self.system_prompt}\n<</SYS>>\n\n{user_message} [/INST]"
+            elif "dialogpt" in self.model.lower():
                 # DialoGPT works better with conversation format
                 prompt = f"User: {user_message}\nAssistant:"
             else:
@@ -154,21 +157,40 @@ Minimum purchase is 1 roll (67-70 yards)."""
             
             logger.info(f"Generating response for prompt: {prompt[:100]}...")
             
-            # Generate response
-            response = self.client.text_generation(
-                prompt,
-                max_new_tokens=200,  # Reduced for better performance
-                temperature=0.7,
-                top_p=0.9,
-                top_k=40,
-                repetition_penalty=1.1,
-                stop_sequences=["<|end|>", "<|user|>", "User:", "Assistant:"]
-            )
+            # Generate response with model-specific parameters
+            if "llama" in self.model.lower():
+                response = self.client.text_generation(
+                    prompt,
+                    max_new_tokens=200,
+                    temperature=0.7,
+                    top_p=0.9,
+                    top_k=40,
+                    repetition_penalty=1.1,
+                    stop_sequences=["</s>", "[INST]", "User:", "Assistant:"]
+                )
+            else:
+                response = self.client.text_generation(
+                    prompt,
+                    max_new_tokens=200,
+                    temperature=0.7,
+                    top_p=0.9,
+                    top_k=40,
+                    repetition_penalty=1.1,
+                    stop_sequences=["<|end|>", "<|user|>", "User:", "Assistant:"]
+                )
             
             logger.info(f"Raw AI response: {response[:200]}...")
             
-            # Clean up the response
-            if "Assistant:" in response:
+            # Clean up the response based on model type
+            if "llama" in self.model.lower():
+                # For Llama models, extract content after [/INST]
+                if "[/INST]" in response:
+                    assistant_response = response.split("[/INST]")[-1].strip()
+                else:
+                    assistant_response = response.strip()
+                # Remove Llama-specific tokens
+                assistant_response = assistant_response.replace("<s>", "").replace("</s>", "").replace("[INST]", "").replace("[/INST]", "").strip()
+            elif "Assistant:" in response:
                 assistant_response = response.split("Assistant:")[-1].strip()
             elif "<|assistant|>" in response:
                 assistant_response = response.split("<|assistant|>")[-1].strip()