import streamlit as st
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

# Page config
st.set_page_config(page_title="Biomedical NER App", layout="centered")

st.title("🧬 Biomedical Named Entity Recognition")
st.write("Extract diseases, drugs, genes, proteins, and chemicals from biomedical text.")

@st.cache_resource
def load_model():
    model_name = "d4data/biomedical-ner-all"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForTokenClassification.from_pretrained(model_name)

    ner_pipeline = pipeline(
        "ner",
        model=model,
        tokenizer=tokenizer,
        aggregation_strategy="simple"
    )
    return ner_pipeline

ner_pipeline = load_model()

text = st.text_area(
    "Enter Biomedical Text:",
    height=200,
    placeholder="Example: Aspirin is used to treat heart disease. BRCA1 mutation causes cancer."
)

if st.button("🔍 Extract Entities"):
    if text.strip() == "":
        st.warning("Please enter some text.")
    else:
        results = ner_pipeline(text)

        if not results:
            st.info("No biomedical entities found.")
        else:
            st.subheader("📌 Extracted Entities")
            for ent in results:
                st.markdown(
                    f"""
                    **Entity:** {ent['word']}  
                    **Type:** `{ent['entity_group']}`  
                    **Confidence:** `{round(ent['score'], 3)}`
                    ---
                    """
                )