Update app.py

app.py

@@ -223,6 +223,9 @@ import torch
 from fastapi import FastAPI, Request, HTTPException, status
 from pydantic import BaseModel, Field
 
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
 # Load model and tokenizer
 model_name = "FreedomIntelligence/Apollo-7B"
 # model_name = "emilyalsentzer/Bio_ClinicalBERT"
@@ -269,29 +272,33 @@ Question: {message}
 Answer:"""
 
 # Chat function
+# @app.post("/ask")
+# def chat_fn(message):
+#     prompt = generate_prompt(message)
+#     response = pipe(prompt,
+#                     max_new_tokens=512,
+#                     temperature=0.7,
+#                     do_sample = True,
+#                     top_p=0.9)[0]['generated_text']
+#     answer = response.split("Answer:")[-1].strip() if "Answer:" in response else response.split("الإجابة:")[-1].strip()
+#     return {"Answer": answer}
+
+executor = ThreadPoolExecutor()
+
+# Define request model
+class Query(BaseModel):
+    message: str
+
 @app.post("/ask")
-def chat_fn(message):
+async def chat_fn(query: Query):
+    message = query.message
     prompt = generate_prompt(message)
-    response = pipe(prompt,
-                    max_new_tokens=512,
-                    temperature=0.7,
-                    do_sample = True,
-                    top_p=0.9)[0]['generated_text']
-    answer = response.split("Answer:")[-1].strip() if "Answer:" in response else response.split("الإجابة:")[-1].strip()
-    return {"Answer": answer}
-
-# Gradio ChatInterface
-# demo = gr.ChatInterface(
-#     fn=chat_fn,
-#     title="🩺 Apollo Medical Chatbot",
-#     description="Multilingual (Arabic & English) medical Q&A chatbot powered by Apollo-7B model inference.",
-#     theme=gr.themes.Soft()
-# )
 
+    # Run blocking inference in thread
+    loop = asyncio.get_event_loop()
+    response = await loop.run_in_executor(executor,
+                                          lambda: pipe(prompt, max_new_tokens=512, temperature=0.7, do_sample=True, top_p=0.9)[0]['generated_text'])
 
-# if __name__ == "__main__":
-#     demo.launch(share=True)
-
-
-    
-
+    # Parse answer
+    answer = response.split("Answer:")[-1].strip() if "Answer:" in response else response.split("الإجابة:")[-1].strip()
+    return {"Answer": answer}