Updated docker file

app.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
@@ -22,15 +23,29 @@ app.add_middleware(
 class ChatRequest(BaseModel):
     message: str
 
-@app.post("/chat")
-async def chat(req: ChatRequest):
-    try:
-        messages = [{"role": "user", "content": req.message}]
-        response = client.chat_completion(
+# @app.post("/chat")
+# async def chat(req: ChatRequest):
+#     try:
+#         messages = [{"role": "user", "content": req.message}]
+#         response = client.chat_completion(
+#             model=MODEL_ID,
+#             messages=messages,
+#             temperature=0.7,
+#         )
+#         return {"response": response.choices[0].message.content}
+#     except Exception as e:
+#         return {"error": str(e)}
+
+
+async def chat_endpoint(data: ChatRequest):
+    def stream():
+        for chunk in client.text_generation(
             model=MODEL_ID,
-            messages=messages,
+            prompt=data.message,
+            stream=True,
+            max_new_tokens=512,
             temperature=0.7,
-        )
-        return {"response": response.choices[0].message.content}
-    except Exception as e:
-        return {"error": str(e)}
+        ):
+            yield chunk
+
+    return StreamingResponse(stream(), media_type="text/plain")