Update app.py

app.py

@@ -9,7 +9,7 @@ import torch.nn.functional as F
 
 app = Flask(__name__)
 
-# ---------------- GLOBAL ERROR HANDLER ----------------
+# ---------------- ERROR HANDLER ----------------
 @app.errorhandler(Exception)
 def handle_error(e):
     return jsonify({"error": str(e)}), 500
@@ -44,23 +44,46 @@ def clean_text(text):
         outputs = bart_model.generate(**inputs, max_length=50)
         return bart_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+# ---------------- BETTER REFERENCE ----------------
 ref_texts = {
-    "emergency": "severe chest pain unconscious breathing failure stroke heart attack",
-    "urgent": "fever infection moderate pain vomiting weakness",
-    "non-urgent": "mild headache cold cough stable"
+    "emergency": "severe chest pain heart attack breathing difficulty unconscious stroke heavy bleeding",
+    "urgent": "high fever infection vomiting dehydration moderate pain weakness",
+    "non-urgent": "common cold runny nose mild headache sneezing cough minor symptoms"
 }
 
 ref_embeddings = {k: get_embedding(v) for k, v in ref_texts.items()}
 
+# ---------------- RULE ENGINE (KEY FIX) ----------------
+def apply_rules(text, label, score):
+    t = text.lower()
+
+    # 🚨 Emergency overrides
+    if any(x in t for x in ["chest pain", "heart", "breathing", "unconscious", "stroke", "bleeding"]):
+        return "emergency", max(score, 90)
+
+    # ⚠️ Urgent
+    if any(x in t for x in ["fever", "vomiting", "infection", "high temperature"]):
+        return "urgent", max(score, 60)
+
+    # ✅ Non-urgent
+    if any(x in t for x in ["cold", "runny nose", "sneezing", "mild headache"]):
+        return "non-urgent", min(score, 25)
+
+    return label, score
+
 # ---------------- CORE PIPELINE ----------------
 def run_pipeline(text):
 
+    # Translate
     try:
         text_en = GoogleTranslator(source="auto", target="en").translate(text)
     except:
         text_en = text
 
+    # Clean
     cleaned = clean_text(text_en)
+
+    # Embedding
     emb = get_embedding(cleaned)
 
     scores = {
@@ -71,8 +94,13 @@ def run_pipeline(text):
     label = max(scores, key=scores.get)
     confidence = scores[label]
 
-    score_map = {"emergency": 90, "urgent": 60, "non-urgent": 20}
-    score = int(score_map[label] * max(confidence, 0.5))
+    base_map = {"emergency": 90, "urgent": 60, "non-urgent": 20}
+
+    # 🔥 Improved scoring
+    score = int((base_map[label] * 0.7) + (confidence * 30))
+
+    # 🔥 Apply rules (CRITICAL)
+    label, score = apply_rules(cleaned, label, score)
 
     return {
         "processed_text": cleaned,
@@ -83,12 +111,11 @@ def run_pipeline(text):
 
 # ---------------- ROUTES ----------------
 
-# ✅ FIXED ROOT (NO HTML)
 @app.route("/", methods=["GET"])
 def root():
     return jsonify({
         "status": "running",
-        "message": "AI Triage API is live",
+        "message": "AI Triage API live",
         "endpoints": {
             "audio": "/process-audio",
             "text": "/process-text"
@@ -110,13 +137,10 @@ def process_audio():
 
     file.save(input_path)
 
-    try:
-        subprocess.run([
-            "ffmpeg", "-y", "-i", input_path,
-            "-ac", "1", "-ar", "16000", wav_path
-        ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-    except:
-        return jsonify({"error": "FFmpeg failed"}), 500
+    subprocess.run([
+        "ffmpeg", "-y", "-i", input_path,
+        "-ac", "1", "-ar", "16000", wav_path
+    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
 
     segments, _ = whisper_model.transcribe(wav_path)
     text = " ".join([seg.text for seg in segments]).strip()