import gradio as gr
from transformers import pipeline

# Load your fine-tuned SpanBERT ADR model (replace with your actual model name)
ner_pipeline = pipeline(
    "token-classification",
    model="hazarri/fine_tuned_spanbert",  # make sure this matches your uploaded model repo name
    aggregation_strategy="simple"
)

def extract_adrs(text):
    results = ner_pipeline(text)
    if not results:
        return "No ADR entities found."
    
    formatted = []
    for r in results:
        formatted.append(f"{r['word']} ({r['entity_group']}, score={r['score']:.2f})")
    return "\n".join(formatted)

demo = gr.Interface(
    fn=extract_adrs,
    inputs=gr.Textbox(lines=5, placeholder="Enter a medical review..."),
    outputs="text",
    title="SpanBERT ADR Extraction API",
    description="Extracts adverse drug reactions (ADRs) from patient reviews using a fine-tuned SpanBERT model."
)

demo.launch()