Update app.py

app.py

@@ -91,6 +91,8 @@ def load_models():
         model_name = "M4-ai/Orca-2.0-Tau-1.8B"
         models['llm_tokenizer'] = AutoTokenizer.from_pretrained(model_name)
         models['llm_model'] = AutoModelForCausalLM.from_pretrained(model_name)
+        models['gen_tokenizer'] = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-1.7B-Instruct")
+        models['gen_model'] = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-1.7B-Instruct")
         print("Models loaded successfully")
         return True
     except Exception as e:
@@ -501,6 +503,49 @@ def translate_en_to_ar(text):
         print(f"Error during English to Arabic translation: {e}")
         return None
 
+
+
+# Medical context prompt
+MEDICAL_PROMPT = """You are a medical doctor who provides accurate and reliable health information based on current medical knowledge. 
+Only answer medical questions and provide information from reliable healthcare sources. 
+If a question is not medical in nature, politely explain that you can only address health-related queries.
+Question: {question}
+Answer: """
+
+def generate_response(question, max_length=350):
+    tok = models['gen_tokenizer']
+    mod = models['gen_model']    
+    # Prepare prompt
+    full_prompt = MEDICAL_PROMPT.format(question=question)
+    
+    # Generate response
+    inputs = tok(full_prompt, return_tensors="pt", truncation=True)
+    outputs = mod.generate(
+        inputs.input_ids,
+        max_length=max_length,
+        num_beams=4,
+        temperature=0.7,
+        early_stopping=True
+    )    
+    response = tok.decode(outputs[0], skip_special_tokens=True)
+    # Extract only the answer part
+    answer = response.split("Answer: ")[-1].strip()
+    return answer
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 @app.get("/")
 async def root():
     return {"message": "Welcome to the FastAPI application! Use the /health endpoint to check health, and /api/query for processing queries."}
@@ -516,6 +561,38 @@ async def health_check():
     }
     return status
 
+
+
+
+
+
+@app.post("/api/ask")
+async def chat_endpoint(chat_query: ChatQuery):
+    try:
+        query_text = chat_query.query
+        language_code = chat_query.language_code        
+        if language_code == 0:
+            query_text = translate_ar_to_en(query_text)
+        # Generate response
+        answer = generate_response(query_text)        
+        # Translate back to Arabic if needed
+        if query.language_code == 0:
+            answer = translate_en_to_ar(answer)
+            if not answer:
+                return Response(answer="", error="Translation failed")
+        
+        return Response(answer=answer)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+
+
+
+
+
+    
+
 @app.post("/api/chat")
 async def chat_endpoint(chat_query: ChatQuery):
     try: