Update app.py

app.py

@@ -2,55 +2,52 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Modern, tiny model trained for instruction following
-MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_NAME = "Agent_1.5"
 
+# Fast tokenizer and CPU-optimized loading
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)
 
 def respond(message, history):
-    # Format the prompt so the model knows it's a chat
-    # We keep it simple to ensure it follows the "Thoughts/Response" rule
-    prompt = f"<|im_start|>system\nYou are agent_1. Rules: Show reasoning under [Thoughts] (1 line), then final answer under [Response]. Stay concise.<|im_end|>\n"
+    # System prompt to enforce the "Agent_1" persona and formatting
+    system_instruction = "You are agent_1. Rules: 1. Start with [Thoughts] (one sentence). 2. End with [Response]. 3. Be concise."
     
-    # Add conversation history for actual "smart" context
-    for user_msg, bot_msg in history[-3:]: # Last 3 turns to save memory
-        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n{bot_msg}<|im_end|>\n"
-    
-    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+    # Building the chat template
+    messages = [{"role": "system", "content": system_instruction}]
+    for user_msg, bot_msg in history[-2:]:  # Last 2 turns to keep CPU speed high
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
+    messages.append({"role": "user", "content": message})
+
+    # Use the official chat template for better logic
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(input_text, return_tensors="pt")
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=100,
-            temperature=0.4, # Lower temp = more logical/less "crazy"
+            max_new_tokens=80,
+            temperature=0.3, # Low temp = smarter, less "rambling"
             do_sample=True,
-            repetition_penalty=1.2,
-            eos_token_id=tokenizer.eos_token_id
+            repetition_penalty=1.1,
+            pad_token_id=tokenizer.eos_token_id
         )
 
-    # Extract only the new generated text
-    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response_part = full_output.split("assistant")[-1].strip()
-    
-    # Ensure it follows the [Thoughts] format even if it forgets
-    if "[Thoughts]" not in response_part:
-        response_part = f"[Thoughts] Analyzing... \n[Response] {response_part}"
-    
-    # Clean up any "User:" hallucinations
-    final_output = response_part.split("<|im_start|>")[0].strip()
-    
-    return final_output
+    # Decode only the new part (the assistant's response)
+    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Extract response after the last "assistant" block
+    response = full_text.split("assistant")[-1].strip()
+
+    # Fallback if the model forgets the [Thoughts] tag
+    if "[Thoughts]" not in response:
+        response = f"[Thoughts] Processing request...\n[Response] {response.replace('[Response]', '').strip()}"
+
+    return response
 
 demo = gr.ChatInterface(
     fn=respond,
-    title="Agent_1: Enhanced SmolLM",
-    description="Swapped DistilGPT2 for SmolLM2-135M ",
-    examples=["Who are you?", "What is 2+2?", "Tell me a joke."],
-    theme="soft"
+    title="Agent_1 (Smol-Instruct)",
+    description=" using SmolLM2-135M. ",
 )
 
 if __name__ == "__main__":