update app.py

app.py

@@ -15,50 +15,56 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 )
 
-# Ensure tokenizer has a chat template
-if tokenizer.chat_template is None:
-    # Fallback: define a basic one if needed (check model card for correct format)
-    tokenizer.chat_template = "{% for message in messages %}{{message['role']}}: {{message['content']}}\n{% endfor %}"
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 model.eval()
 print(f"✅ Model loaded on {device}!")
 
+# Define a strong system prompt
+SYSTEM_PROMPT = (
+    "You are an expert AI assistant. Provide clear, accurate, and concise answers to the user's questions. "
+    "Do not add extra commentary, disclaimers, or summaries unless asked. Answer directly."
+)
+
 @app.get("/")
 def root():
     return {"message": "MobileLLM-Pro API is running!"}
 
 @app.get("/generate")
-def generate(prompt: str, max_tokens: int = 50):
+def generate(prompt: str, max_tokens: int = 256):
     try:
-        # Format as a chat with user message
-        messages = [{"role": "user", "content": prompt}]
+        # Construct full chat with system + user
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": prompt}
+        ]
         
         # Apply chat template
         input_text = tokenizer.apply_chat_template(
             messages,
-            tokenize=False,  # We'll tokenize next
-            add_generation_prompt=True  # Adds assistant start token
+            tokenize=False,
+            add_generation_prompt=True  # Ensures <|assistant|> or equivalent is added
         )
         
-        # Tokenize
+        # Tokenize and move to device
         inputs = tokenizer(input_text, return_tensors="pt").to(device)
         
-        # Generate
+        # Generate response
         outputs = model.generate(
             **inputs,
             max_new_tokens=max_tokens,
             do_sample=True,
             temperature=0.7,
+            top_p=0.9,
             pad_token_id=tokenizer.eos_token_id
         )
         
-        # Decode only the generated part (after input)
-        generated_tokens = outputs[0][inputs['input_ids'].shape[1]:]
+        # Extract only the generated part (after input)
+        input_len = inputs.input_ids.shape[1]
+        generated_tokens = outputs[0][input_len:]
         result = tokenizer.decode(generated_tokens, skip_special_tokens=True)
         
-        return {"input": prompt, "output": result}
+        return {"input": prompt, "output": result.strip()}
     
     except Exception as e:
         return {"error": str(e)}