Update app.py

app.py

@@ -8,47 +8,70 @@ model_id = "facebook/MobileLLM-R1-950M"
 pipe = pipeline(
     "text-generation",
     model=model_id,
-    torch_dtype=torch.float16,
+    torch_dtype="auto",
     device_map="auto",
 )
 
 @spaces.GPU(duration=120)
 def respond(message, history):
-    # Build prompt from history
-    prompt = ""
+    # Build messages list from history
+    messages = []
+    
+    # Add system message based on content type detection
+    if any(kw in message.lower() for kw in ["python", "def ", "function"]):
+        messages.append({
+            "role": "system",
+            "content": (
+                "\nYou are a helpful and harmless assistant. You should think step-by-step before responding to the instruction below.\n\n"
+                "Please use python programming language only.\n"
+                "You must use ```python for just the final solution code block with the following format:\n"
+                "```python\n# Your code here\n```\n"
+            )
+        })
+    elif any(kw in message.lower() for kw in ["c++", "cpp", "#include", "cout"]):
+        messages.append({
+            "role": "system",
+            "content": (
+                "\nYou are a helpful and harmless assistant. You should think step-by-step before responding to the instruction below.\n\n"
+                "Please use c++ programming language only.\n"
+                "You must use ```cpp for just the final solution code block with the following format:\n"
+                "```cpp\n// Your code here\n```\n"
+            )
+        })
+    elif any(kw in message.lower() for kw in ["compute", "calculate", "math", "+", "-", "*", "/"]):
+        messages.append({
+            "role": "system",
+            "content": "Please reason step by step, and put your final answer within \\boxed{}."
+        })
+    else:
+        messages.append({
+            "role": "system",
+            "content": "You are a helpful AI assistant."
+        })
+    
+    # Add conversation history
     for user_msg, assistant_msg in history:
         if user_msg:
-            prompt += f"User: {user_msg}\n"
+            messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
-            prompt += f"Assistant: {assistant_msg}\n"
+            messages.append({"role": "assistant", "content": assistant_msg})
     
     # Add current message
-    prompt += f"User: {message}\nAssistant: "
-    
-    # Generate response with streaming
-    streamer = pipe.tokenizer.decode
-    
-    # Generate tokens
-    inputs = pipe.tokenizer(prompt, return_tensors="pt").to(pipe.model.device)
-    
-    with torch.no_grad():
-        outputs = pipe.model.generate(
-            **inputs,
-            max_new_tokens=10000,
-            temperature=0.7,
-            do_sample=True,
-            pad_token_id=pipe.tokenizer.eos_token_id,
-        )
+    messages.append({"role": "user", "content": message})
     
-    # Decode the generated tokens, skipping the input tokens
-    generated_tokens = outputs[0][inputs['input_ids'].shape[-1]:]
+    # Generate response
+    outputs = pipe(
+        messages,
+        max_new_tokens=8192,
+        temperature=0.7,
+        do_sample=True,
+    )
     
-    # Stream the output token by token
+    # Extract and stream the generated text
+    full_response = outputs[0]["generated_text"][-1]["content"]
     response_text = ""
-    for i in range(len(generated_tokens)):
-        token = generated_tokens[i:i+1]
-        token_text = pipe.tokenizer.decode(token, skip_special_tokens=True)
-        response_text += token_text
+    for char in full_response:
+        response_text += char
         yield response_text
 
 # Create the chat interface