Update app.py

app.py

@@ -4,24 +4,22 @@ from transformers import pipeline
 # Load your fine-tuned SpanBERT ADR model
 ner_pipeline = pipeline(
     "token-classification",
-    model="hazarri/fine_tuned_spanbert",  # replace with your actual model repo name
+    model="hazarri/fine_tuned_spanbert",  # replace with your actual model name
     aggregation_strategy="simple"
 )
 
 def extract_adrs(text):
-    # Run the model on the input text
+    # Run the model
     results = ner_pipeline(text)
 
-    # If no ADRs found, return an empty list
+    # Return empty list if no entities detected
     if not results:
         return []
 
-    # Format the results as a list
-    formatted = [
-        f"{r['word']} ({r['entity_group']}, score={r['score']:.2f})"
-        for r in results
-    ]
-    return formatted
+    # Extract only the ADR words
+    adrs = [r["word"] for r in results]
+
+    return adrs
 
 # Gradio interface
 demo = gr.Interface(