Update app.py

app.py

@@ -2,15 +2,14 @@ import gradio as gr
 from transformers import pipeline
 import torch
 import os
-from pathlib import Path
+from typing import Union, List
 
 # ==============================
 # ⚙️ Configuration
 # ==============================
-# Use a smaller zero-shot model to avoid memory/storage limits
 MODEL_NAME = os.getenv("MODEL_NAME", "MoritzLaurer/deberta-v3-base-zeroshot-v2.0")
 
-# Optional: specify cache dir outside your repo
+# Optional: specify cache dir
 HF_HOME = os.getenv("HF_HOME", "/tmp/huggingface")
 os.environ["TRANSFORMERS_CACHE"] = HF_HOME
 
@@ -33,28 +32,46 @@ except Exception as e:
 # ==============================
 # 🧠 Classification Function
 # ==============================
-def classify_severity(side_effects_str: str):
-    if not side_effects_str or not side_effects_str.strip():
-        return {}
-
-    side_effects_list = [se.strip() for se in side_effects_str.split(",") if se.strip()]
-    candidate_labels = ["Mild", "Moderate", "Severe", "Life-threatening", "Death"]
-
+def classify_severity(side_effects_input: Union[str, List[str]]):
+    """
+    Accepts a string (comma-separated) or a list of side effects.
+    Returns a dictionary mapping side effects -> severity.
+    """
     try:
+        # Handle both list and string input types
+        if isinstance(side_effects_input, str):
+            side_effects_list = [
+                se.strip() for se in side_effects_input.split(",") if se.strip()
+            ]
+        elif isinstance(side_effects_input, list):
+            side_effects_list = [se.strip() for se in side_effects_input if se.strip()]
+        else:
+            return {"error": "Invalid input format. Expected string or list."}
+
+        if not side_effects_list:
+            return {}
+
+        candidate_labels = ["Mild", "Moderate", "Severe", "Life-threatening", "Death"]
+
+        # Perform classification
         results = severity_classifier(
             sequences=side_effects_list,
             candidate_labels=candidate_labels,
             multi_label=False
         )
 
+        # Normalize single result into list
         if isinstance(results, dict):
             results = [results]
 
-        return {
+        # Convert to dict { side_effect: severity }
+        output_dict = {
             result["sequence"]: result["labels"][0].lower()
             for result in results
         }
 
+        return output_dict
+
     except Exception as e:
         return {"error": str(e)}
 
@@ -63,10 +80,12 @@ def classify_severity(side_effects_str: str):
 # ==============================
 demo = gr.Interface(
     fn=classify_severity,
-    inputs=gr.Textbox(
-        label="Side Effects",
-        placeholder="Enter side effects, separated by commas (e.g., headache, nausea, dizziness)"
-    ),
+    inputs=[
+        gr.Textbox(
+            label="Side Effects",
+            placeholder="Enter side effects, separated by commas (e.g., headache, nausea, dizziness)"
+        )
+    ],
     outputs=gr.JSON(label="Severity Classification"),
     title="⚕️ Side Effect Severity Classifier",
     description="Classifies side effect severity using a DeBERTa/BART model for zero-shot classification.",
@@ -86,6 +105,7 @@ if __name__ == "__main__":
 
 
 
+
 # import gradio as gr
 # from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
 # from pathlib import Path