Update app.py

app.py

@@ -223,9 +223,9 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Load model and tokenizer
-model_name = "FreedomIntelligence/Apollo-7B"
+# model_name = "FreedomIntelligence/Apollo-7B"
 # model_name = "emilyalsentzer/Bio_ClinicalBERT"
-# model_name = "FreedomIntelligence/Apollo-2B"
+model_name = "FreedomIntelligence/Apollo-2B"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
@@ -300,45 +300,48 @@ 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)
 
-    try:
-        start_time = time.time()
+    # Run blocking inference in thread
+    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'])
+
+    # Parse answer
+    answer = response.split("Answer:")[-1].strip() if "Answer:" in response else response.split("الإجابة:")[-1].strip()
+    return {
+        "response": response,
+        "Answer": answer
+    }
+
+# @app.post("/ask")
+# async def chat_fn(query: Query):
+#     message = query.message
+#     logger.info(f"Received message: {message}")
+    
+#     prompt = generate_prompt(message)
+
+#     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.6, do_sample=True, top_p=0.8)[0]['generated_text']
-        )
+#         loop = asyncio.get_event_loop()
+#         response = await loop.run_in_executor(
+#             executor,
+#             lambda: pipe(prompt, max_new_tokens=150, temperature=0.6, do_sample=True, top_p=0.8)[0]['generated_text']
+#         )
         
-        duration = time.time() - start_time
-        logger.info(f"Model inference completed in {duration:.2f} seconds")
+#         duration = time.time() - start_time
+#         logger.info(f"Model inference completed in {duration:.2f} seconds")
 
-        logger.info(f"Generated answer: {answer}")
-        return {"Answer": answer}
+#         logger.info(f"Generated answer: {answer}")
+#         return {"Answer": answer}
 
-    except Exception as e:
-        logger.error(f"Inference failed: {str(e)}")
-        raise HTTPException(status_code=500, detail="Model inference TimeOut failed.")
+#     except Exception as e:
+#         logger.error(f"Inference failed: {str(e)}")
+#         raise HTTPException(status_code=500, detail="Model inference TimeOut failed.")