# app.py
import gradio as gr
from gliner import GLiNER
import re

# Load GLiNER-BioMed Small model
model = GLiNER.from_pretrained("Ihor/gliner-biomed-bi-small-v1.0", trust_remote_code=True)
labels = ["Population", "Intervention", "Comparison", "Outcome"]

def parse_nbib(file):
    content = file.read().decode("utf-8")
    entries = re.findall(r'(?=PMID- .+?)(.*?)(?=(?:PMID- |\Z))', content, re.DOTALL)
    refs = []
    for entry in entries:
        title = ""
        abstract = ""
        for line in entry.splitlines():
            if line.startswith("TI  - "):
                title = line.replace("TI  - ", "").strip()
            elif line.startswith("AB  - "):
                abstract += " " + line.replace("AB  - ", "").strip()
        if title or abstract:
            refs.append(f"{title}. {abstract.strip()}")
    return refs

def extract_pico_from_nbib(file):
    refs = parse_nbib(file)
    results = []
    for ref in refs:
        entities = model.predict_entities(ref, labels, threshold=0.5)
        results.append({"text": ref, "pico": entities})
    return results

def extract_from_text(text):
    entities = model.predict_entities(text, labels, threshold=0.5)
    return {ent['label']: ent['text'] for ent in entities}

demo = gr.Interface(
    fn=extract_from_text,
    inputs=gr.Textbox(lines=10, placeholder="Paste biomedical abstract here..."),
    outputs="json",
    title="PICO Extractor (GLiNER-BioMed Small)",
    description="Extract PICO (Population, Intervention, Comparison, Outcome) elements from biomedical text using the GLiNER-BioMed Small model."
)

nbib_demo = gr.Interface(
    fn=extract_pico_from_nbib,
    inputs=gr.File(file_types=[".nbib"]),
    outputs="json",
    title="PICO Extractor from NBIB",
    description="Upload .nbib files to extract PICO elements using GLiNER-BioMed Small."
)

tabs = gr.TabbedInterface([demo, nbib_demo], ["From Text", "From NBIB File"])

if __name__ == "__main__":
    tabs.launch()