Update app.py

app.py

@@ -1,23 +1,67 @@
 from fastapi import FastAPI
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from pydantic import BaseModel
 import uvicorn
 
 app = FastAPI()
 
-# Load model
+# Load model + tokenizer
 model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
-chat_pipeline = pipeline("text-generation", model=model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
+
+chat_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+)
+
+SYSTEM_PROMPT = "You are Bella, a smart and witty female AI assistant. Keep responses short and to the point."
+
 
 class ChatRequest(BaseModel):
     message: str
 
+
 @app.post("/chat")
 async def chat(request: ChatRequest):
-    prompt = f"You are Bella, a smart and witty female AI assistant. Keep it very short.\nUser: {request.message}\nBella:"    
-    outputs = chat_pipeline(prompt, max_new_tokens=175, temperature=0.3)
-    response = outputs[0]['generated_text'].split("Bella:")[-1].split("User:")[0].strip()
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": request.message},
+    ]
+
+    # Build prompt using the model's chat template
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    outputs = chat_pipeline(
+        prompt,
+        max_new_tokens=175,
+        do_sample=True,
+        temperature=0.3,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
+    generated_text = outputs[0]["generated_text"]
+
+    # Strip off the prompt portion, leaving only the new reply
+    response = generated_text[len(prompt):].strip()
+
+    # Safety net: cut off if the model starts hallucinating a new turn
+    for stop_token in ["<|im_start|>", "User:", "system"]:
+        if stop_token in response:
+            response = response.split(stop_token)[0].strip()
+
+    if not response:
+        response = "Hmm, I didn't quite catch that — can you rephrase?"
+
     return {"response": response}
 
+
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)