Switched to TheBloke/Mistral-7B-Instruct-v0.2-AWQ

app.py

@@ -4,9 +4,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import re
 
 # Load model and tokenizer
-print("Loading gemma...")
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
+print("Loading Mistral-7B-Instruct AWQ...")
+tokenizer = AutoTokenizer.from_pretrained("TheBloke/Mistral-7B-Instruct-v0.2-AWQ", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    "TheBloke/Mistral-7B-Instruct-v0.2-AWQ",
+    device_map="auto",
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+)
 
 # Add pad token if it doesn't exist
 if tokenizer.pad_token is None:
@@ -88,33 +93,47 @@ def respond(message, history, max_length=150, temperature=0.9, top_p=0.9, top_k=
         if check_crisis_keywords(message):
             return get_crisis_response()
         
-        # Build conversation history - SIMPLIFIED for DialoGPT
-        conversation = ""
+        # Build conversation history using Mistral chat template
+        messages = []
+        
+        # Add system message for Aura personality
+        messages.append({"role": "system", "content": AURA_SYSTEM_PROMPT})
         
         # Only include last 2-3 exchanges to avoid overwhelming the model
         recent_history = history[-2:] if len(history) > 2 else history
         
         for user_msg, bot_msg in recent_history:
-            conversation += f"{user_msg}{tokenizer.eos_token}{bot_msg}{tokenizer.eos_token}"
+            messages.append({"role": "user", "content": user_msg})
+            if bot_msg:
+                messages.append({"role": "assistant", "content": bot_msg})
         
         # Add current message
-        conversation += f"{message}{tokenizer.eos_token}"
+        messages.append({"role": "user", "content": message})
+        
+        # Apply chat template
+        conversation = tokenizer.apply_chat_template(
+            messages, 
+            tokenize=False, 
+            add_generation_prompt=True
+        )
         
         # Tokenize
         input_ids = tokenizer.encode(conversation, return_tensors="pt")
         
-        # Generate response with configurable parameters
+        # Generate response with configurable parameters optimized for Mistral
         with torch.no_grad():
             chat_history_ids = model.generate(
-                input_ids,
-                max_length=max_length,
+                input_ids.to(model.device),
+                max_new_tokens=min(max_length - input_ids.shape[-1], 512),  # Use max_new_tokens instead
                 temperature=temperature,
                 top_p=top_p,
                 repetition_penalty=repetition_penalty,
                 do_sample=True,
                 top_k=top_k,
-                pad_token_id=tokenizer.eos_token_id,
-                no_repeat_ngram_size=3
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                no_repeat_ngram_size=2,
+                use_cache=True
             )
         
         # Decode only the new response

requirements.txt

@@ -1,3 +1,5 @@
 torch>=2.0.0,<2.2.0
-transformers>=4.30.0,<4.40.0
+transformers>=4.35.0,<4.40.0
+autoawq>=0.1.8
+accelerate>=0.20.0
 gradio>=3.50.0,<4.0.0

test_model.py

@@ -1,17 +1,22 @@
 #!/usr/bin/env python3
 """
-Test script to validate DialoGPT model response generation
+Test script to validate Mistral-7B-Instruct AWQ model response generation
 """
 
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 def test_model():
-    print("Loading DialoGPT-medium for testing...")
+    print("Loading Mistral-7B-Instruct AWQ for testing...")
     
     # Load model and tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+    tokenizer = AutoTokenizer.from_pretrained("TheBloke/Mistral-7B-Instruct-v0.2-AWQ", trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        "TheBloke/Mistral-7B-Instruct-v0.2-AWQ",
+        device_map="auto",
+        trust_remote_code=True,
+        torch_dtype=torch.float16
+    )
     
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
@@ -28,24 +33,33 @@ def test_model():
     for i, message in enumerate(test_messages):
         print(f"\n--- Test {i+1}: '{message}' ---")
         
-        # Simple conversation format
-        conversation = f"{message}{tokenizer.eos_token}"
+        # Use Mistral chat template format
+        messages = [
+            {"role": "user", "content": message}
+        ]
+        
+        # Apply chat template
+        conversation = tokenizer.apply_chat_template(
+            messages, 
+            tokenize=False, 
+            add_generation_prompt=True
+        )
+        
         input_ids = tokenizer.encode(conversation, return_tensors="pt")
         
-        # Generate response with conservative settings
+        # Generate response with settings optimized for Mistral AWQ
         with torch.no_grad():
             chat_history_ids = model.generate(
-                input_ids,
-                max_length=input_ids.shape[-1] + 50,
-                num_beams=5,
-                no_repeat_ngram_size=3,
+                input_ids.to(model.device),
+                max_new_tokens=100,
+                no_repeat_ngram_size=2,
                 do_sample=True,
-                early_stopping=True,
-                pad_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
                 eos_token_id=tokenizer.eos_token_id,
                 temperature=0.9,
                 top_k=50,
-                top_p=0.9
+                top_p=0.9,
+                use_cache=True
             )
         
         # Decode response