Update app.py

app.py

@@ -9,7 +9,7 @@ llm = Llama(
     n_ctx=2048,
     n_threads=2,
     n_batch=128,
-    verbose=False,
+    verbose=False
 )
 
 class Message(BaseModel):
@@ -21,21 +21,21 @@ class ChatRequest(BaseModel):
 
 @app.post("/v1/chat")
 def chat(req: ChatRequest):
-    prompt = ""
+    prompt = "<|system|>\nYou are a helpful, concise chatbot.\n"
     for m in req.messages:
-        prompt += f"{m.role.upper()}: {m.content}\n"
-    prompt += "ASSISTANT:"
+        prompt += f"<|{m.role}|>\n{m.content}\n"
+    prompt += "<|assistant|>\n"
 
     output = llm(
         prompt,
         max_tokens=256,
         temperature=0.7,
         top_p=0.9,
-        stop=["USER:"]
+        stop=["<|user|>", "<|system|>"]
     )
 
     return {
-        "model": "tinyllama-1.1b-chat-q4_k_m",
+        "model": "edyx-convo",
         "text": output["choices"][0]["text"].strip(),
         "tokens": output["usage"]["total_tokens"],
     }