Updated docker file

app.py

@@ -1,10 +1,10 @@
 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
 import os
 
+
 HF_TOKEN = os.getenv("HF_TOKEN")
 MODEL_ID = "google/gemma-2b-it"
 client = InferenceClient(token=HF_TOKEN)
@@ -23,29 +23,15 @@ 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(
-#             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(
+@app.post("/chat")
+async def chat(req: ChatRequest):
+    try:
+        messages = [{"role": "user", "content": req.message}]
+        response = client.chat_completion(
             model=MODEL_ID,
-            prompt=data.message,
-            stream=True,
-            max_new_tokens=512,
+            messages=messages,
             temperature=0.7,
-        ):
-            yield chunk
-
-    return StreamingResponse(stream(), media_type="text/plain")
+        )
+        return {"response": response.choices[0].message.content}
+    except Exception as e:
+        return {"error": str(e)}