Update app.py

app.py

@@ -1,11 +1,30 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
+from transformers import pipeline
 
 app = FastAPI()
 
+pipe = pipeline(
+    "text-generation",
+    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+)
+
 class Req(BaseModel):
     message: str
 
 @app.post("/chat")
 def chat(req: Req):
-    return {"reply": f"AI says: {req.message}"}
+    prompt = f"<s>[INST] {req.message} [/INST]"
+
+    out = pipe(
+        prompt,
+        max_new_tokens=100,
+        temperature=0.7
+    )
+
+    text = out[0]["generated_text"]
+
+    # remove prompt from output
+    reply = text.split("[/INST]")[-1].strip()
+
+    return {"reply": reply}