Update app.py

app.py

@@ -24,24 +24,28 @@ print("Model loaded successfully!")
 
 
 class ChatRequest(BaseModel):
-    model: str
+    model: str = "auric-ai"
     messages: list
     stream: bool = False
+    max_tokens: int = 512
+    temperature: float = 0.1
 
 
 @app.post("/v1/chat/completions")
 async def chat(req: ChatRequest):
 
-    user_message = req.messages[-1]["content"]
-
-    prompt = f"""
-You are a helpful assistant.
-
-User: {user_message}
-Assistant:
-"""
+    # Use Qwen2's chat template for proper system/user/assistant formatting
+    prompt = tokenizer.apply_chat_template(
+        req.messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
 
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+    prompt_length = inputs["input_ids"].shape[-1]
+
+    temperature = max(req.temperature, 0.01)
+    max_tokens = min(req.max_tokens, 2048)
 
     # ---------------- STREAM MODE ----------------
 
@@ -56,8 +60,8 @@ Assistant:
         generation_kwargs = dict(
             **inputs,
             streamer=streamer,
-            max_new_tokens=512,
-            temperature=0.7,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
             do_sample=True
         )
 
@@ -96,22 +100,20 @@ Assistant:
 
         output = model.generate(
             **inputs,
-            max_new_tokens=80,
-            temperature=0.7,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
             do_sample=True
         )
 
-        response = tokenizer.decode(output[0], skip_special_tokens=True)
-
-        if "Assistant:" in response:
-            response = response.split("Assistant:")[-1].strip()
+        # Only decode the newly generated tokens, not the prompt
+        response = tokenizer.decode(output[0][prompt_length:], skip_special_tokens=True)
 
         return {
             "choices": [
                 {
                     "message": {
                         "role": "assistant",
-                        "content": response
+                        "content": response.strip()
                     }
                 }
             ]