# BIOBERT-FTM — Gradio Space App
import os, json
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoConfig, AutoModelForTokenClassification, pipeline

# ---- Model sources (override via Space Secrets or env) ----
TOKENIZER_REPO = os.getenv("MLM_MODEL_REPO_ID", "Milad96/BIOBERT-FTM-mlm")
TASK_MODEL_REPO = os.getenv("TASK_MODEL_REPO_ID", "Milad96/BIOBERT-FTM-tasks")
NER_SUBFOLDER   = os.getenv("NER_SUBFOLDER", "ner-spyysalo_bc2gm_corpus")

DEVICE = 0 if torch.cuda.is_available() else -1

tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_REPO, use_fast=True)
config = AutoConfig.from_pretrained(TASK_MODEL_REPO, subfolder=NER_SUBFOLDER)
model = AutoModelForTokenClassification.from_pretrained(
    TASK_MODEL_REPO, subfolder=NER_SUBFOLDER, torch_dtype=torch.float32
)

ner = pipeline(
    "token-classification",
    model=model,
    tokenizer=tokenizer,
    device=DEVICE,
    aggregation_strategy="simple"
)

fill_mask = pipeline(
    "fill-mask",
    model=TOKENIZER_REPO,
    tokenizer=TOKENIZER_REPO,
    device=DEVICE
)

EXAMPLE_TEXTS = [
    "Mutations in BRCA1 and TP53 are frequently reported in breast cancer.",
    "The expression of EGFR and KRAS was quantified by qPCR.",
    "BRCA2 interacts with RAD51 during homologous recombination.",
]

MASK_EXAMPLES = [
    "SARS-CoV-2 infects [MASK] epithelial cells in the respiratory tract.",
    "Dexamethasone is commonly used to treat [MASK] inflammation.",
]

def run_fill_mask(text: str, topk: int = 5):
    if not text or "[MASK]" not in text:
        return "Input must contain [MASK] token.", f"Device: {'cuda' if DEVICE==0 else 'cpu'}"
    try:
        res = fill_mask(text, top_k=int(topk))
        if isinstance(res, dict):  # older transformers may return dict for top_k=1
            res = [res]
        lines = []
        for p in res:
            seq = p.get("sequence", "")
            score = float(p.get("score", 0.0))
            lines.append(f"→ {seq} (score={score:.4f})")
        return "\n".join(lines), f"Device: {'cuda' if DEVICE==0 else 'cpu'}"
    except Exception as e:
        return f"⚠️ Error: {e}", f"Device: {'cuda' if DEVICE==0 else 'cpu'}"


def run_ner(text: str, score_threshold: float = 0.50):
    if not text:
        return "", "[]"
    try:
        preds = ner(text)
        if not isinstance(preds, list):
            preds = []
        spans = []
        for s in preds:
            try:
                sc = float(s.get("score", 0.0))
                if sc < float(score_threshold):
                    continue
                start = int(s.get("start", 0))
                end = int(s.get("end", 0))
                spans.append({
                    "entity": s.get("entity_group") or s.get("entity") or "",
                    "start": start,
                    "end": end,
                    "text": text[start:end],
                    "score": sc
                })
            except Exception:
                # Skip any malformed span safely
                continue
        # Markdown summary
        if spans:
            lines = [f"Found {len(spans)} spans:"]
            for sp in spans:
                lines.append(
                    f"- [{sp['entity']}] {sp['text']} (start={sp['start']}, end={sp['end']}, score={sp['score']:.4f})"
                )
            md = "\n".join(lines)
        else:
            md = "No spans above threshold."
        return md, json.dumps(spans, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"⚠️ Error: {e}", "[]"

with gr.Blocks(title="BIOBERT-FTM Demo", theme=gr.themes.Default()) as demo:
    gr.Markdown("""
    # 🧬 BIOBERT-FTM — Demo
    **Fill‑Mask** (MLM) and **Bio‑NER** (Gene) demo powered by your fine‑tuned checkpoints.
    """)

    with gr.Tabs():
        with gr.TabItem("Fill-Mask"):
            inp_mask = gr.Textbox(label="Input (must contain [MASK])", value=MASK_EXAMPLES[0])
            topk = gr.Slider(1, 10, value=5, step=1, label="Top‑K")
            btn = gr.Button("Run Fill-Mask")
            out_txt = gr.Textbox(label="Predictions", lines=6)
            out_meta = gr.Textbox(label="Meta")
            btn.click(run_fill_mask, inputs=[inp_mask, topk], outputs=[out_txt, out_meta])
            gr.Examples(MASK_EXAMPLES, inputs=[inp_mask])

        with gr.TabItem("NER (Genes)"):
            inp = gr.Textbox(label="Input text", value=EXAMPLE_TEXTS[0])
            thr = gr.Slider(0, 1, value=0.50, step=0.01, label="Score threshold")
            go = gr.Button("Run NER")
            out_md = gr.Markdown()
            out_json = gr.JSON()
            go.click(run_ner, inputs=[inp, thr], outputs=[out_md, out_json])
            gr.Examples(EXAMPLE_TEXTS, inputs=[inp])

if __name__ == "__main__":
    # When running inside Spaces, just `python app.py` is enough.
    demo.queue().launch()