Update app.py

app.py

@@ -226,8 +226,11 @@ from pydantic import BaseModel, Field
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Load model and tokenizer
-model_name = "FreedomIntelligence/Apollo-7B"
+model_name = "FreedomIntelligence/Apollo-2B"
 # model_name = "emilyalsentzer/Bio_ClinicalBERT"
 # model_name = "FreedomIntelligence/Apollo-2B"
 
@@ -294,16 +297,45 @@ def read_root():
     return {"message": "Apollo Medical Chatbot API is running"}
 
 
+# @app.post("/ask")
+# async def chat_fn(query: Query):
+    
+#     message = query.message
+#     logger.info(f"Received message: {message}")
+    
+#     prompt = generate_prompt(message)
+
+#     # 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'])
+
+#     # Parse answer
+#     answer = response.split("Answer:")[-1].strip() if "Answer:" in response else response.split("الإجابة:")[-1].strip()
+#     return {"Answer": answer}
+
 @app.post("/ask")
 async def chat_fn(query: Query):
     message = query.message
+    logger.info(f"Received message: {message}")
+    
     prompt = generate_prompt(message)
 
-    # 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'])
+    try:
+        start_time = time.time()
+        
+        loop = asyncio.get_event_loop()
+        response = await loop.run_in_executor(
+            executor,
+            lambda: pipe(prompt, max_new_tokens=150, temperature=0.7, do_sample=True, top_p=0.9)[0]['generated_text']
+        )
+        
+        duration = time.time() - start_time
+        logger.info(f"Model inference completed in {duration:.2f} seconds")
+
+        logger.info(f"Generated answer: {answer}")
+        return {"Answer": answer}
 
-    # Parse answer
-    answer = response.split("Answer:")[-1].strip() if "Answer:" in response else response.split("الإجابة:")[-1].strip()
-    return {"Answer": answer}
+    except Exception as e:
+        logger.error(f"Inference failed: {str(e)}")
+        raise HTTPException(status_code=500, detail="Model inference failed.")