# app_offline_ner_min.py
import os
os.environ["TRANSFORMERS_OFFLINE"] = "1"   # force offline per HF docs
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
import gradio as gr

# point to your local snapshot downloaded by prepare_model.py
# path: ./models/biomedical-ner-all
HERE = os.path.dirname(os.path.abspath(__file__))
LOCAL_MODEL_DIR = os.path.join(HERE, "models", "biomedical-ner-all")

# load strictly from disk
tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_DIR, local_files_only=True)
model = AutoModelForTokenClassification.from_pretrained(LOCAL_MODEL_DIR, local_files_only=True)

device = 0 if torch.cuda.is_available() else -1
ner_pipe = pipeline(
    task="token-classification",                # NER
    model=model,
    tokenizer=tokenizer,
    aggregation_strategy="simple",              # merge subword tokens into entities
    device=device
)

def run_ner(text: str):
    if not text.strip():
        return {"text": "", "entities": []}, []

    out = ner_pipe(text)

    highlighted = {
        "text": text,
        "entities": [
            {
                "entity": r["entity_group"],
                "start": int(r["start"]),
                "end": int(r["end"]),
                "score": float(r["score"]),
            }
            for r in out
        ],
    }

    # list-of-lists in a fixed column order
    rows = [
        [r["entity_group"], r["word"], float(r["score"]), int(r["start"]), int(r["end"])]
        for r in out
    ]
    return highlighted, rows

with gr.Blocks() as demo:
    gr.Markdown("# 🩺 Biomedical NER (offline, local model)")
    inp = gr.Textbox(label="Enter text", value="Patient has a history of asthma treated with albuterol.")
    ner_view = gr.HighlightedText(label="Entities", combine_adjacent=True)
    table = gr.Dataframe(
        label="Raw predictions",
        headers=["entity", "word", "score", "start", "end"],  # <-- headers for list-of-lists
        interactive=False,
    )
    inp.change(run_ner, inp, [ner_view, table])

demo.launch(debug=True)