Update app.py

app.py

@@ -14,6 +14,7 @@ DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
 # ── 2) Load & Filter Dataset ─────────────────────────────────────────────
 df = pd.read_csv("medquad.csv")
 df["text"] = df["question"].str.strip() + " " + df["answer"].str.strip()
+
 vc   = df["focus_area"].value_counts()
 keep = vc[vc >= MIN_FREQ].index
 df   = df[df["focus_area"].isin(keep)].reset_index(drop=True)
@@ -28,31 +29,42 @@ bert_model = transformers.AutoModel.from_pretrained(MODEL_NAME).to(DEVICE).eval(
 
 @torch.no_grad()
 def encode_text(s: str, max_length=MAX_LEN):
-    toks = tokenizer(
+    toks   = tokenizer(
         s,
         return_tensors="pt",
         truncation=True,
         max_length=max_length,
-        add_special_tokens=True
+        padding=False,
     ).to(DEVICE)
     hidden = bert_model(**toks).last_hidden_state
     return hidden[:,0].squeeze().cpu()
 
 # ── 4) Precompute Static Label Embeddings ─────────────────────────────────
-label_embs = torch.stack([
-    encode_text(lbl, max_length=16)
-    for lbl in labels
-])
+label_embs = torch.stack([encode_text(lbl, max_length=16) for lbl in labels])
 
-# ── 5) Classification fn ───────────────────────────────────────────────────
+# ── 5) Classification Function ────────────────────────────────────────────
 def predict_disease(symptoms: str):
-    q_emb = encode_text(symptoms).unsqueeze(0)
-    sims  = cosine_similarity(q_emb, label_embs)
-    idx   = sims.argmax(dim=1).item()
-    return id2label[idx]
+    if not symptoms.strip():
+        return "❗️ Please enter your symptoms."
+    try:
+        # 1) embed user text → [hidden_size]
+        q_emb = encode_text(symptoms)
+
+        # 2) compute cosine similarities → [num_labels]
+        sims = cosine_similarity(
+            label_embs,            # [num_labels, hidden_size]
+            q_emb.unsqueeze(0),    # [1, hidden_size]
+            dim=1
+        )
+
+        # 3) pick the best label index
+        best = sims.argmax().item()
+        return id2label[best]
+
+    except Exception as e:
+        return f"Error: {e}"
 
 # ── 6) Gradio Interface ───────────────────────────────────────────────────
-# Rename to "app" so Hugging Face picks it up automatically:
 app = gr.Interface(
     fn=predict_disease,
     inputs=gr.Textbox(
@@ -64,7 +76,7 @@ app = gr.Interface(
     description="PubMed-BERT + cosine similarity"
 )
 
-# Explicitly launch on HF’s host/port
+# ── 7) Launch ─────────────────────────────────────────────────────────────
 app.launch(
     server_name="0.0.0.0",
     server_port=int(os.environ.get("PORT", 7860)),