app.py

from __future__ import annotations

from collections import Counter, defaultdict
from functools import lru_cache
from typing import Dict, List

import gradio as gr
from transformers import pipeline


# =============================================================================
# Configuration
# =============================================================================

APP_TITLE = "OTAR3088 CeLLaTe NER Model Demo"

# Replace these placeholders with your existing model metadata.
# Each model entry should include:
#   - repo_id: Hugging Face repo id
#   - description: short human-readable summary
#   - tag: optional label ("best", "fast", etc.)
#   - supported_entities: list of entity labels in canonical form
#
# Canonical labels used by this app:
#   3-class:   CELL_LINE, CELL_TYPE, TISSUE
#   2-class:   CELL_LINE, CELL_TISSUE
#
# Example shape:
# THREE_CLASS_MODELS = {
#     "bioformer-3class-baseline": {
#         "repo_id": "your-org/your-3class-model",
#         "description": "3-class Bioformer baseline.",
#         "tag": "best",
#         "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"],
#     },
# }
#
# TWO_CLASS_MODELS = {
#     "bioformer-2class-baseline": {
#         "repo_id": "your-org/your-2class-model",
#         "description": "2-class Bioformer baseline.",
#         "tag": "best",
#         "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
#     },
# }

THREE_CLASS_MODELS: Dict[str, dict] = {
    "Bioformer8l-3class-BaseModel": 
                {"repo_id": "OTAR3088/CeLLaTe-ner-3class-bioformer8l-baseline",
                "description": "3-class Bioformer baseline Model",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "Bioformer16l-3class-BaseModel": 
                {"repo_id": "Mardiyyah/CeLLaTe-ner-3class-bioformer16l-baseline",
                "description": "3-class Bioformer baseline Model",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "Bioformer16l-3class-BaseModel_old": 
                {"repo_id": "OTAR3088/CeLLaTe-ner-3class-bioformer16l-baseline",
                "description": "3-class Bioformer baseline Model old version",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "PubmedBert-3class-BaseModel":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-pubmedbert-baseline",
                "description": "3-class PubmedBert baseline Model",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "PubmedBert-3class-BaseModel_old":
                {
                "repo_id": "OTAR3088/CeLLaTe-ner-3class-pubmedbert-baseline",
                "description": "3-class PubmedBert baseline old Model",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "Bioformer16l-3class-tapt-tokenizer-original-standardmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-bioformer16l-tapt-tokenizer-original-baseline",
                "description": "3-class Bioformer tapt Model with original tokenizer and standard-masking",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },

    "Bioformer16l-3class-tapt-tokenizer-original-spanmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-bioformer16l-tapt-tokenizer-original-spanmask",
                "description": "3-class Bioformer tapt Model with original tokenizer and span-masking",
                "tag": "TAPT-adapted-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "Bioformer16l-3class-tapt-tokenizer-adapted-standardmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-bioformer16l-tapt-tokenizer-adapted-baseline",
                "description": "3-class PubmedBert baseline Model",
                "tag": "TAPT-adapted-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    
    "Bioformer16l-3class-tapt-tokenizer-adapted-spanmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-bioformer16l-tapt-tokenizer-adapted-spanmask",
                "description": "3-class Bioformer tapt Model with original tokenizer and span-masking",
                "tag": "TAPT-adapted-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "PubmedBert-3class-tapt-tokenizer-original-standardmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-pubmedbert-tapt-tokenizer-original-baseline",
                "description": "3-class Pubmedbert tapt Model with original tokenizer and standard-masking",
                "tag": "TAPT-adapted-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "PubmedBert-3class-tapt-tokenizer-original-spanmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-pubmedbert-tapt-tokenizer-original-spanmask",
                "description": "3-class Pubmedbert tapt Model with original tokenizer and span-masking",
                "tag": "TAPT-adapted-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "PubmedBert-3class-tapt-tokenizer-original-wwmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-pubmedbert-tapt-tokenizer-original-wwmask",
                "description": "3-class Pubmedbert Model with original tokenizer and wholeword-masking",
                "tag": "TAPT-adapted-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },

    "PubmedBert-3class-tapt-tokenizer-adapted-standardmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-pubmedbert-tapt-tokenizer-adapted-baseline",
                "description": "3-class Pubmedbert tapt Model with adapted tokenizer and standard-masking",
                "tag": "TAPT-adapted-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    "PubmedBert-3class-tapt-tokenizer-adapted-wwmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-3class-pubmedbert-tapt-tokenizer-adapted-wwmask",
                "description": "3-class Pubmedbert tapt Model with adapted tokenizer and wholeword-masking",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
     "PubmedBert-3class-tapt-tokenizer-adapted-oldversion_lr-4.27":
                {
                "repo_id": "OTAR3088/CeLLaTe_V3.3_lr-4.27",
                "description": "3-class Pubmedbert tapt Model with adapted tokenizer and wholeword-masking",
                "tag": "Current-best-3class-model",
                "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"]
                },
    
    
}

TWO_CLASS_MODELS: Dict[str, dict] = {
   "Bioformer8l-2class-BaseModel": 
                {"repo_id": "OTAR3088/CeLLaTe-ner-2class-bioformer8l-baseline",
                "description": "2-class Bioformer baseline Model",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },
    "Bioformer16l-2class-BaseModel": 
                {"repo_id": "Mardiyyah/CeLLaTe-ner-2class-bioformer16l-baseline",
                "description": "2-class Bioformer baseline Model",
                "tag": "Base-Model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },

    "Bioformer16l-2class-reinitllrd-gazetteers": 
                {"repo_id": "Mardiyyah/CeLLaTe-ner-2class-bioformer16l-reinitllrd-with-gazetteers-lr_2.754e-5",
                "description": "2-class Bioformer reinit-llrd Model",
                "tag": "Base-ReinitLLRD-Model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },

    "PubmedBert-2class-BaseModel":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-pubmedbert-baseline",
                "description": "2-class PubmedBert baseline Model",
                "tag": "Base-Model",
                "supported_entities":["CELL_LINE", "CELL_TISSUE"],
                },
    # "PubmedBert-2class-tapt-tokenizer-original-standardmasking": 
    #             {"repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-pubmedbert-tokenizer-original-baseline",
    #             "description": "2-class Pubmedbert tapt Model with original-tokenizer and standard-masking",
    #             "tag": "TAPT-Model",
    #             "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
    #             },
   
    "PubmedBert-2class-tapt-tokenizer-adapted-wwmask-combinedData-lr_3.89": 
                {"repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-pubmedbert-tokenizer-adapted-combinedData-lr_3.89",
                "description": "2-class Pubmedbert tapt Model with adapted tokenizer and wholeword-masking",
                "tag": "TAPT-Model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },

    
    "Bioformer16l-2class-tapt-tokenizer-original-standardmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-bioformer16l-tokenizer-original-baseline",
                "description": "2-class Bioformer tapt Model with original tokenizer and standard-masking",
                "tag": "TAPT-adapted model",
                "supported_entities":["CELL_LINE", "CELL_TISSUE"],
                },

    
    "Bioformer16l-2class-tapt-tokenizer-original-wwmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-bioformer16l-tokenizer-original-wwmask",
                "description": "2-class Bioformer tapt Model with original tokenizer and wholeword-masking",
                "tag": "TAPT-adapted model",
                "supported_entities":["CELL_LINE", "CELL_TISSUE"],
                },

    
    "Bioformer16l-2class-tapt-tokenizer-original-spanmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-bioformer16l-tokenizer-original-spanmask",
                "description": "2-class Bioformer tapt Model with original tokenizer and span-masking",
                "tag": "TAPT-adapted model",
                "supported_entities":["CELL_LINE", "CELL_TISSUE"],
                },
    
    "Bioformer16l-2class-tapt-tokenizer-adapted-standardmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-bioformer16l-tokenizer-adapted-baseline",
                "description": "2-class Bioformer tapt Model with adapted tokenizer and standard-masking",
                "tag": "TAPT-adapted model",
                "supported_entities":["CELL_LINE", "CELL_TISSUE"],
                },

    
    "Bioformer16l-2class-tapt-tokenizer-adapted-wwmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-bioformer16l-tokenizer-adapted-wwmask",
                "description": "2-class Bioformer tapt Model with adapted tokenizer and wholeword-masking",
                "tag": "TAPT-adapted model",
                "supported_entities":["CELL_LINE", "CELL_TISSUE"],
                },

    
    "Bioformer16l-2class-tapt-tokenizer-adapted-spanmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-bioformer16l-tokenizer-adapted-spanmask",
                "description": "2-class Bioformer tapt Model with adapted tokenizer and span-masking",
                "tag": "TAPT-adapted model",
                "supported_entities":["CELL_LINE", "CELL_TISSUE"],
                },
    

    "PubmedBert-2class-tapt-tokenizer-adapted-standardmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-pubmedbert-tokenizer-adapted-baseline",
                "description": "2-class Pubmedbert tapt Model with adapted tokenizer and standard-masking",
                "tag": "TAPT-adapted model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },
    
     "PubmedBert-2class-tapt-combData-tokenizer-adapted-standardmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-pubmedbert-tapt-combData-tokenizer-adapted-baseline",
                "description": "2-class Pubmedbert tapt Model with adapted tokenizer and standard-masking",
                "tag": "TAPT-adapted model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },
    
    "PubmedBert-2class-tapt-tokenizer-adapted-wwmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-pubmedbert-tokenizer-adapted-wwmask",
                "description": "2-class Pubmedbert tapt Model with adapted tokenizer and wholeword-masking",
                "tag": "TAPT-adapted model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },
    
    "PubmedBert-2class-tapt-tokenizer-adapted-wwmask-combinedData": 
                {"repo_id": "Mardiyyah/CeLLaTe-ner-2class-pubmedbert-tapt-combData-tokenizer-adapted-wwmask",
                "description": "2-class Pubmedbert tapt Model with adapted tokenizer and wholeword-masking",
                "tag": "TAPT-Model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },
    
    "PubmedBert-2class-tapt-tokenizer-adapted-spanmask":
                {
                "repo_id": "Mardiyyah/CeLLaTe-ner-2class-tapt-pubmedbert-tokenizer-adapted-spanmask",
                "description": "2-class Pubmedbert tapt Model with adapted tokenizer and span-masking",
                "tag": "TAPT-adapted model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },

    "PubmedBert-2class-tapt-tokenizer-adapted-spanmask-combinedData": 
                {"repo_id": "Mardiyyah/CeLLaTe-ner-2class-pubmedbert-tapt-combData-tokenizer-adapted-spanmask",
                "description": "2-class Pubmedbert tapt Model with adapted tokenizer and span-masking",
                "tag": "TAPT-Model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
                },


    "PubmedBert-2class-tapt-tokenizer-adapted-oldversion":
                {
                "repo_id": "OTAR3088/CeLLaTe_contracted_ent_Reinit_LLRD",
                "description": "2-class Pubmedbert tapt Model with adapted tokenizer and wholeword-masking",
                "tag": "Current-best-2class-model",
                "supported_entities": ["CELL_LINE", "CELL_TISSUE"]
                },
}

TASK_CONFIGS = {
    "3class": {
        "title": "3-class models",
        "kicker": "CellLine + CellType + Tissue",
        "default_model": next(iter(THREE_CLASS_MODELS), ""),
        "models": THREE_CLASS_MODELS,
        "supported_entities": ["CELL_LINE", "CELL_TYPE", "TISSUE"],
        "description": "Models trained to recognise CellLine, CellType, and Tissue separately.",
    },
    "2class": {
        "title": "2-class models",
        "kicker": "CellLine + Cell_Tissue",
        "default_model": next(iter(TWO_CLASS_MODELS), ""),
        "models": TWO_CLASS_MODELS,
        "supported_entities": ["CELL_LINE", "CELL_TISSUE"],
        "description": "Models trained to recognise CellLine and a combined Cell_Tissue class.",
    },
}

EXAMPLES = [
    "Aside from in vivo models, numerous studies investigating bacterial virulence and pathogenesis have also employed in vitro cell line models to gain an initial understanding of the intricate host-pathogen interactions. These studies, which are simpler and more cost-effective than those using in vivo models, serve as the foundation for many in vivo studies by providing additional data to support any conclusions [12]. Epithelial mucous membrane cells are the primary focus of most in vitro investigations due to them being usually the initial point of contact for infections [12,13].HeLa cells, which originate from human cervical epithelial cells, are thus frequently selected for bacterial adhesion and invasion and are particularly suitable for experiments [14]. A. baumannii frequently infects human epithelial tissues, such as the respiratory system, skin and mucosal linings [15]. HeLa cells are resilient and readily cultured in vitro, exhibiting a rapid growth rate.This ensures the availability of a uniform and consistent cell population for studies, rendering them economical and reliable.",
    "Based on the limma R package, a total of 2578 (DEGs 1398 downregulated and 1188 upregulated) were screened out from GEO: GSE225819 data, including 20 normal samples and 20 GIST samples with liver metastasis (|log2FC| > 1; P < 0.05), suggesting that these DEGs may be involved in liver metastasis in GIST patients (Figure (Figure1A).1A). The top 10 upregulated genes were PENK, IGF2, GPR20, CTSL, SCRG1, PNMAL1, NKX3-2, ANO1, PLAT, and BCHE. The top 10 downregulated genes were ATP4B, GKN1, MT1G, GKN2, ATP4A, SPINK1, TSPAN8, TFF1, KCNE2, and REG1A (Supplementary Table 1). Based on the Deseq2, 1386 DEGs (939 downregulated and 447 upregulated) were screened out in GSE155880, including seven Imatinib-sensitive samples and seven imatinib-resistant GIST patients (|log2FC| > 1; P < 0.05, Figure Figure1B).1B). The intersection of the two analyses indicated that only IGF2 was involved in the drug resistance regulation and GIST metastasis in these DEGs (Supplementary Table 2).Moreover, we evaluated IGF2 expression in the GIST cell line. By western blotting, expression levels of IGF2 in GIST882, GIST882-R, GIST-T1, and GIST-T1-R were higher than those in normal RGM-1. Furthermore, IGF2 was significantly over expressed in GIST882-R/GIST-T1-R compared with other cell lines GIST882/GIST-T1 (P < 0.01, P < 0.001; Figure Figure1C).1C). In addition, the expression levels of IGF2 in culture supernatants were measured using ELISA and compared (Figure (Figure1D).1D). We found that the ELISA and western blot results (P < 0.05, P < 0.001) were similar.IGF2 expression was high in drug-resistant GIST cell lines, suggesting that IGF2 overexpression may be closely related to drug resistance.",
    "Given the dynamic nature of T cell activation and the heterogeneity of CD4+ T cells, we mapped gene expression regulation using single-cell transcriptomes spanning four time points of CD4+ T cell activation. We reconstructed activation trajectories for naive and memory CD4+ T cells and identified eQTL effects manifesting at different time points and across different subpopulations of cells. We identified 127 genes with colocalizing eQTL and GWAS signals for immune-mediated diseases. Colocalizing genes were enriched in time-dependent eQTLs. Our data suggest that dysregulation of gene expression during T cell activation could underlie immune disease and emphasize the importance of context-specific gene expression regulation.",
    "To quantitatively characterise the macrophage foetal-like profile observed, we projected adult and foetal macrophages data from the human decidual-placental interface into our data set. This unique tissue setting includes both adult/maternal monocyte-derived macrophages and foetal/placental YS-derived macrophages (Hofbauer cells)20, thus avoiding the technical confounders of adult+foetal data set integration. The Hofbauer cells LR model had a higher mean prediction probability for iPSC-derived macrophages than any of the adult macrophage subtypes identified in the placenta (Supplementary Fig. 5D, E). As an exception, day31 + 1 macrophages presented a higher score with adult macrophages LR models (Supplementary Fig. 5D, E). This was likely because day31 + 1 macrophages showed an activated state, upregulating inflammatory cytokines such as CXCL8, CCL7 or IL1B, (Supplementary Fig. 5F) also found on monocyte-derived macrophages in the decidua (adult/maternal tissue). Next, we used a hepatocellular carcinoma data set27 and projected several tumour-associated macrophages (TAMs) LR models on our iPSC-derived data set (Supplementary Fig. 5G, H). We found the foetal-like FOLR2+ TAMs model showed a higher prediction score among end-stage macrophages (day31 + 7) while the SPP1+ TAMs LR model markedly captured macrophages on the activated state (day31 + 1, Supplementary Fig. 5H). Overall, this indicated that macrophages produced in the iPSC protocol have a strong foetal phenotype, and this could be relevant for their application as in vitro TAM models.",
    "Peripheral blood mononuclear cells (PBMCs) were isolated using Ficoll-Paque PLUS (GE Healthcare) density gradient centrifugation. Naive (CD25- CD45RA+ CD45RO-) and memory (CD25- CD45RA- CD45RO+) CD4+ T cells were isolated from the PBMC fraction using EasySep naive CD4+ T cell isolation kits and memory CD4+ T cell enrichment kits (StemCell Technologies) according to the manufacturer's instructions.Naive and memory T cells were then stimulated with anti-CD3/anti-CD28 human T-Activator Dynabeads (Invitrogen) at a 1:2 beads-to-cells ratio. Cells were harvested after 16 h, 40 h and 5 d of stimulation.In addition, unstimulated cells kept in culture without any beads for 16 h were used as a negative control (i.e., 0 h of activation).",
    "We performed dimensionality reduction and embedding using uniform manifold approximation (UMAP) (ref. 21) (Methods) and observed that cells separated by time point of stimulation, forming a gradual progression from resting to the most activated cell state (cells collected at 5 d) (Fig. 1b). This progression was accompanied by changes in activation markers. For example, an early activation marker, CD69, was upregulated at 16 h but downregulated at later time points, whereas expression of IL2RA, a marker of late activation, peaked at 40 h, remaining present at 5 d (Fig. 1c). A population of cells localized at the intermediate point between resting and 16 h-stimulated cells (Fig. 1b), and was composed of cells from the 16 h (74%) and 40 h (26%) time points. We hypothesized that this intermediate group represented an early activation state. By analyzing cells from these two time points independently, we observed that at each of these time points cells separated into two clear groups, one corresponding to the early activation state (Supplementary Fig. 3). Cells in the early activation group expressed fourfold fewer genes compared to other cells at their respective activation time points and showed lower expression of T cell activation markers19 (Supplementary Fig. 3). Furthermore, they showed a unique profile characterized by high expression of STAT1, IFIT3 and GBP1 (Fig. 1c).Therefore, these cells represent a distinct, early activation state, and we refer to them as lowly active.",
    "Finally, we leveraged this model to experimentally evaluate genes linked to immune-related phenotypes by GWAS. Interestingly, CRISPR/Cas9-mediated KO of GWAS hits (PRKCB, LSP1 and ICAM1) in iPSC-derived macrophages and DCs highlighted their potential role in physiological and pathological cell states of distinct cell types. Specifically, macrophage differentiation in KO lines altered inflammatory and extracellular matrix genes. Fibrosis constitutes a pathological feature of most chronic inflammatory diseases including the ones featured in our study87,88, and our results open an avenue for therapeutic intervention in these disorders. In line with this, we show that both the foetal-like FOLR2+ and the SPP1+ TAM states observed in liver cancer27 are recapitulated in this system.This suggests these cells could also be a faithful model to unravel the role of macrophage subtypes in the tumour microenvironment.",
]


ENTITY_META = {
    "CELL_LINE": {"display": "CellLine"},
    "CELL_TYPE": {"display": "CellType"},
    "TISSUE": {"display": "Tissue"},
    "CELL_TISSUE": {"display": "Cell_Tissue"},
}

LABEL_NORMALIZATION = {
    "CellLine": "CELL_LINE",
    "CELL_LINE": "CELL_LINE",
    "CELL-LINE": "CELL_LINE",
    "CellType": "CELL_TYPE",
    "CELL_TYPE": "CELL_TYPE",
    "CELL-TYPE": "CELL_TYPE",
    "Tissue": "TISSUE",
    "TISSUE": "TISSUE",
    "Cell_Tissue": "CELL_TISSUE",
    "CELL_TISSUE": "CELL_TISSUE",
    "CELL-TISSUE": "CELL_TISSUE",
}


# =============================================================================
# Helpers
# =============================================================================

def display_name(label: str) -> str:
    return ENTITY_META.get(label, {}).get("display", label)


def make_example_label(text: str, idx: int, max_chars: int = 92) -> str:
    one_line = " ".join(text.split())
    preview = one_line[:max_chars].rstrip()
    if len(one_line) > max_chars:
        preview += "..."
    return f"Example {idx + 1}: {preview}"


def build_task_registry(task_key: str) -> Dict[str, dict]:
    cfg = TASK_CONFIGS[task_key]
    models = cfg["models"]
    default_model = cfg["default_model"]

    if not models:
        return {}

    registry = {}
    for model_name, meta in models.items():
        registry[model_name] = {
            "repo_id": meta["repo_id"].strip(),
            "enabled": True,
            "tag": meta.get("tag", "").strip(),
            "description": meta.get("description", "").strip(),
            "supported_entities": meta.get("supported_entities", cfg["supported_entities"]),
        }

    if default_model in registry:
        registry[default_model]["tag"] = registry[default_model]["tag"] or "best"
    return registry


MODEL_REGISTRIES = {
    "3class": build_task_registry("3class"),
    "2class": build_task_registry("2class"),
}


def model_choices(task_key: str):
    registry = MODEL_REGISTRIES[task_key]
    choices = []
    for model_name, meta in registry.items():
        tag = meta.get("tag", "").strip()
        label = f"{model_name} · {tag}" if tag else model_name
        choices.append((label, model_name))
    return choices


def example_choices():
    return [(make_example_label(text, i), text.strip()) for i, text in enumerate(EXAMPLES)]



def model_catalog_rows():
    rows = []

    for task_key in ["2class", "3class"]:
        registry = MODEL_REGISTRIES.get(task_key, {})
        family_label = TASK_CONFIGS[task_key]["title"]

        for model_name, cfg in registry.items():
            rows.append([
                family_label,
                model_name,
                cfg.get("tag", ""),
                ", ".join(display_name(x) for x in cfg["supported_entities"]),
                cfg.get("description", ""),
                cfg.get("repo_id", ""),
            ])

    return rows

def render_hero() -> str:
    return f"""
    <div class="hero-banner">
      <div class="hero-inner">
        <div class="hero-kicker">🧬 {APP_TITLE}</div>
        <h1>Biomedical NER Explorer</h1>
        <p>
          Compare named entity recognition across two model families:
          3-class models (CellLine, CellType, Tissue) and 2-class models
          (CellLine, Cell_Tissue).
        </p>
        <div class="hero-chip-row">
          <span class="hero-chip">Biomedical NLP</span>
          <span class="hero-chip">Entity Extraction</span>
          <span class="hero-chip">Two evaluation setups</span>
        </div>
      </div>
    </div>
    """


def render_model_card(task_key: str, model_name: str) -> str:
    registry = MODEL_REGISTRIES[task_key]
    if not model_name or model_name not in registry:
        return "<div class='panel-card'><em>Select a model.</em></div>"

    cfg = registry[model_name]
    tag = cfg.get("tag", "").strip()
    tag_html = ""
    if tag:
        tag_class = "meta-pill meta-pill-best" if tag.lower() == "best" else "meta-pill"
        tag_html = f'<span class="{tag_class}">{tag}</span>'

    return f"""
    <div class="panel-card">
        <div class="panel-label">Model description</div>
        <div class="model-head">
            <div class="model-title">{model_name}</div>
            {tag_html}
        </div>
        <div class="model-desc">{cfg["description"]}</div>
        <div class="model-repo"><strong>Repo:</strong> {cfg["repo_id"]}</div>
    </div>
    """


def render_schema_card(task_key: str, model_name: str) -> str:
    registry = MODEL_REGISTRIES[task_key]
    if not model_name or model_name not in registry:
        return "<div class='panel-card'><em>Select a model.</em></div>"

    supported = registry[model_name]["supported_entities"]
    chips = "".join(f'<span class="entity-chip">{display_name(label)}</span>' for label in supported)

    return f"""
    <div class="panel-card">
        <div class="panel-label">Supported entities</div>
        <div class="schema-note">
            This model predicts <strong>{len(supported)}</strong> entity type(s).
        </div>
        <div class="chip-row">{chips}</div>
    </div>
    """


def render_summary_metrics_html(task_key: str, model_name: str, entities: list[dict] | None = None) -> str:
    registry = MODEL_REGISTRIES[task_key]
    if not model_name or model_name not in registry:
        return "<div class='panel-card'><em>Select a model.</em></div>"

    supported = registry[model_name]["supported_entities"]
    entities = entities or []

    counts = Counter(ent["entity"] for ent in entities)
    total = len(entities)
    active_types = sum(1 for label in supported if counts.get(label, 0) > 0)
    avg_conf = sum(ent.get("score", 0.0) for ent in entities) / total if total else 0.0

    return f"""
    <div class="panel-card summary-card">
        <div class="panel-label">Summary</div>
        <div class="metric-grid">
            <div class="metric-box">
                <div class="metric-value">{total}</div>
                <div class="metric-name">Spans</div>
            </div>
            <div class="metric-box">
                <div class="metric-value">{active_types}</div>
                <div class="metric-name">Types found</div>
            </div>
            <div class="metric-box">
                <div class="metric-value">{avg_conf:.2f}</div>
                <div class="metric-name">Avg. confidence</div>
            </div>
        </div>
    </div>
    """


def build_entity_summary_rows(task_key: str, model_name: str, entities: list[dict] | None = None):
    registry = MODEL_REGISTRIES[task_key]
    if not model_name or model_name not in registry:
        return []

    entities = entities or []
    supported = set(registry[model_name]["supported_entities"])
    grouped_scores = defaultdict(list)

    for ent in entities:
        label = ent["entity"]
        if label in supported:
            grouped_scores[label].append(float(ent.get("score", 0.0)))

    rows = []
    for label, scores in grouped_scores.items():
        avg_score = sum(scores) / len(scores) if scores else 0.0
        rows.append([display_name(label), len(scores), round(avg_score, 3)])

    rows.sort(key=lambda row: (-row[1], row[0]))
    return rows


def load_example(example_text: str):
    return example_text or ""


def normalize_entities(predictions):
    entities = []
    for pred in predictions:
        raw_label = pred.get("entity_group") or pred.get("entity") or "ENTITY"
        label = LABEL_NORMALIZATION.get(raw_label, raw_label.upper())
        entities.append(
            {
                "start": int(pred["start"]),
                "end": int(pred["end"]),
                "entity": label,
                "score": float(pred.get("score", 0.0)),
            }
        )
    return entities


def filter_entities_for_model(task_key: str, model_name: str, entities: list[dict]):
    registry = MODEL_REGISTRIES[task_key]
    supported = set(registry[model_name]["supported_entities"])
    return [ent for ent in entities if ent["entity"] in supported]


@lru_cache(maxsize=16)
def load_model(task_key: str, model_name: str):
    registry = MODEL_REGISTRIES[task_key]
    if model_name not in registry:
        raise KeyError(f"Unknown model: {model_name}")
    repo_id = registry[model_name]["repo_id"]
    return pipeline("ner", model=repo_id, aggregation_strategy="simple")


def run_inference(text: str, task_key: str, model_name: str):
    text = (text or "").strip()
    if not text:
        return (
            {"text": "", "entities": []},
            render_summary_metrics_html(task_key, model_name, []),
            [],
        )

    ner_pipe = load_model(task_key, model_name)
    predictions = ner_pipe(text)
    entities = normalize_entities(predictions)
    entities = filter_entities_for_model(task_key, model_name, entities)

    highlighted = {"text": text, "entities": entities}
    summary_metrics = render_summary_metrics_html(task_key, model_name, entities)
    summary_rows = build_entity_summary_rows(task_key, model_name, entities)

    return highlighted, summary_metrics, summary_rows


def clear_outputs(task_key: str, model_name: str):
    return "", {"text": "", "entities": []}, render_summary_metrics_html(task_key, model_name, []), []


def on_model_change(task_key: str, model_name: str):
    return (
        render_model_card(task_key, model_name),
        render_schema_card(task_key, model_name),
        {"text": "", "entities": []},
        render_summary_metrics_html(task_key, model_name, []),
        [],
    )


# =============================================================================
# UI
# =============================================================================

GR_THEME = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="violet",
    neutral_hue="slate",
)

CSS = """
:root {
    --page-bg: #f8fafc;
    --card-bg: #ffffff;
    --card-border: #e5e7eb;
    --text-main: #0f172a;
    --text-muted: #475569;
    --text-subtle: #64748b;
    --hero-title: #ffffff;
    --hero-body: rgba(255, 255, 255, 0.95);
    --hero-chip-bg: rgba(255, 255, 255, 0.14);
    --hero-chip-border: rgba(255, 255, 255, 0.26);
    --hero-chip-text: #ffffff;
}

.dark,
[data-theme="dark"] {
    --page-bg: #0b1220;
    --card-bg: #111827;
    --card-border: rgba(255, 255, 255, 0.08);
    --text-main: #e5e7eb;
    --text-muted: #cbd5e1;
    --text-subtle: #94a3b8;
    --hero-chip-bg: rgba(255, 255, 255, 0.12);
    --hero-chip-border: rgba(255, 255, 255, 0.18);
    --hero-chip-text: #ffffff;
}

.gradio-container {
    font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif !important;
    font-size: 16px;
    line-height: 1.7;
    background: var(--page-bg);
    color: var(--text-main);
}

.gradio-container p,
.gradio-container li,
.gradio-container label,
.gradio-container input,
.gradio-container textarea,
.gradio-container select,
.gradio-container .markdown,
.gradio-container .prose {
    font-size: 16px !important;
    line-height: 1.75 !important;
    color: var(--text-main);
}

/* Tabs */
.gradio-container .tab-nav,
.gradio-container .tabs button,
.gradio-container .tabitem {
    font-size: 15px !important;
    font-weight: 600;
}

/* Hero */
.hero-banner {
    background: linear-gradient(135deg, #5b7cfa 0%, #7c4dff 100%);
    border-radius: 22px;
    padding: 38px 28px;
    color: var(--hero-title);
    margin: 8px 0 18px 0;
    box-shadow: 0 14px 32px rgba(76, 81, 191, 0.18);
}

.hero-inner {
    text-align: center;
    max-width: 920px;
    margin: 0 auto;
}

.hero-kicker {
    font-size: 1rem;
    font-weight: 700;
    opacity: 0.95;
    margin-bottom: 10px;
    letter-spacing: 0.02em;
    color: var(--hero-title);
}

.hero-banner h1 {
    font-size: clamp(2.5rem, 4vw, 3.4rem);
    line-height: 1.08;
    margin: 0 0 12px 0;
    font-weight: 800;
    color: var(--hero-title);
}

.hero-banner p {
    margin: 0 auto;
    max-width: 760px;
    font-size: 1.08rem;
    line-height: 1.7;
    opacity: 0.96;
    color: var(--hero-body);
}

.hero-chip-row {
    display: flex;
    justify-content: center;
    gap: 10px;
    flex-wrap: wrap;
    margin-top: 16px;
}

.hero-chip {
    background: var(--hero-chip-bg);
    color: var(--hero-chip-text);
    border: 1px solid var(--hero-chip-border);
    padding: 7px 12px;
    border-radius: 999px;
    font-size: 0.92rem;
    font-weight: 600;
}

/* Cards */
.panel-card {
    background: var(--card-bg);
    border: 1px solid var(--card-border);
    border-radius: 18px;
    padding: 20px 20px;   /* increased padding */
    box-shadow: 0 1px 2px rgba(15, 23, 42, 0.03);
    color: var(--text-main);
}

.panel-label {
    display: inline-block;
    font-size: 0.88rem;
    font-weight: 700;
    color: #4f46e5;
    background: #eef2ff;
    padding: 4px 9px;
    border-radius: 8px;
    margin-bottom: 12px;
}

.dark .panel-label,
[data-theme="dark"] .panel-label {
    background: rgba(99, 102, 241, 0.16);
    color: #c7d2fe;
}

.model-head {
    display: flex;
    align-items: center;
    gap: 8px;
    flex-wrap: wrap;
}

.model-title {
    font-size: 1.05rem;
    font-weight: 800;
    color: var(--text-main);
}

.model-desc,
.model-repo,
.schema-note {
    font-size: 1rem;
    line-height: 1.65;
    color: var(--text-muted);
}

.model-desc {
    margin-top: 8px;
}

.model-repo {
    margin-top: 10px;
    color: var(--text-subtle);
    word-break: break-word;
}

.meta-pill {
    display: inline-block;
    padding: 4px 10px;
    border-radius: 999px;
    font-size: 0.82rem;
    font-weight: 700;
    border: 1px solid #dbeafe;
    background: #eff6ff;
    color: #1d4ed8;
}

.dark .meta-pill,
[data-theme="dark"] .meta-pill {
    border-color: rgba(147, 197, 253, 0.25);
    background: rgba(30, 41, 59, 0.9);
    color: #93c5fd;
}

.meta-pill-best {
    border-color: #bbf7d0;
    background: #f0fdf4;
    color: #166534;
}

.dark .meta-pill-best,
[data-theme="dark"] .meta-pill-best {
    border-color: rgba(74, 222, 128, 0.25);
    background: rgba(20, 83, 45, 0.35);
    color: #86efac;
}

.chip-row {
    display: flex;
    gap: 8px;
    flex-wrap: wrap;
    margin-top: 10px;
}

.entity-chip {
    display: inline-block;
    padding: 6px 10px;
    border-radius: 999px;
    font-size: 0.9rem;
    font-weight: 700;
    color: var(--text-main);
    border: 1px solid #dbe4f0;
    background: #f8fafc;
}

.dark .entity-chip,
[data-theme="dark"] .entity-chip {
    border-color: rgba(255, 255, 255, 0.08);
    background: rgba(15, 23, 42, 0.75);
}

/* Summary metrics */
.metric-grid {
    display: grid;
    grid-template-columns: repeat(3, 1fr);
    gap: 10px;
    margin-bottom: 12px;
}

.metric-box {
    background: #f8fafc;
    border: 1px solid #e5e7eb;
    border-radius: 14px;
    padding: 16px 12px;   /* slightly roomier */
    text-align: center;
}

.dark .metric-box,
[data-theme="dark"] .metric-box {
    background: rgba(15, 23, 42, 0.75);
    border-color: rgba(255, 255, 255, 0.08);
}

.metric-value {
    font-size: 1.5rem;    /* larger summary numbers */
    font-weight: 800;
    color: var(--text-main);
    line-height: 1.1;
}

.metric-name {
    margin-top: 4px;
    font-size: 0.9rem;
    color: var(--text-subtle);
}

/* Input/output */
#input-box textarea {
    font-size: 17px !important;
    line-height: 1.8 !important;
    color: var(--text-main) !important;
}

#tagged-output,
#tagged-output * {
    font-size: 16px !important;
    line-height: 1.8 !important;
}

#tagged-output {
    min-height: 320px;
}

#summary-table {
    margin-top: 12px;
}

/* About tab content */
#about-tab,
#about-tab p,
#about-tab li,
#about-tab h1,
#about-tab h2,
#about-tab h3,
#about-tab h4 {
    font-size: 17px !important;
    line-height: 1.9 !important;
    color: var(--text-main);
}

#about-tab {
    padding-top: 6px;
    padding-bottom: 10px;
}

@media (max-width: 900px) {
    .hero-banner h1 {
        font-size: 1.9rem;
    }

    .hero-banner p {
        font-size: 1rem;
    }

    .metric-grid {
        grid-template-columns: 1fr;
    }
}
"""


def build_task_panel(task_key: str):
    cfg = TASK_CONFIGS[task_key]
    registry = MODEL_REGISTRIES[task_key]
    default_model = cfg["default_model"]

    if not registry:
        gr.Markdown(
            f"### {cfg['title']}\n\n"
            "Model registry is empty. Paste your existing Hugging Face model metadata into the "
            f"`{task_key.upper()}_MODELS` dictionary in the script."
        )
        return

    gr.Markdown(f"### {cfg['title']}")
    gr.Markdown(cfg["description"])

    with gr.Row():
        with gr.Column(scale=4, min_width=330):
            model_choice = gr.Dropdown(
                choices=model_choices(task_key),
                value=default_model,
                label="Model",
                info="Choose the NER model to run",
            )

            model_info = gr.HTML(value=render_model_card(task_key, default_model))
            schema_info = gr.HTML(value=render_schema_card(task_key, default_model))

            example_choice = gr.Dropdown(
                choices=example_choices(),
                value=None,
                label="Example text",
                info="Select a sample to load into the input box",
            )

            with gr.Row():
                run_btn = gr.Button("Run extraction", variant="primary")
                clear_btn = gr.Button("Clear")

        with gr.Column(scale=8):
            input_text = gr.Textbox(
                label="Input text",
                placeholder="Paste a biomedical abstract or paragraph here...",
                lines=12,
                max_lines=16,
                elem_id="input-box",
            )

            with gr.Row(equal_height=True):
                with gr.Column(scale=8):
                    output_highlight = gr.HighlightedText(
                        label="Tagged entities",
                        color_map=None,
                        show_legend=False,
                        show_inline_category=True,
                        combine_adjacent=True,
                        adjacent_separator=" ",
                        elem_id="tagged-output",
                    )

                with gr.Column(scale=4):
                    summary_metrics = gr.HTML(
                        value=render_summary_metrics_html(task_key, default_model, []),
                    )
                    summary_table = gr.Dataframe(
                        value=[],
                        headers=["Entity", "Count", "Avg. confidence"],
                        datatype=["str", "number", "number"],
                        interactive=False,
                        label="Entity summary",
                        elem_id="summary-table",
                    )

    example_choice.change(
        fn=load_example,
        inputs=example_choice,
        outputs=input_text,
    )

    model_choice.change(
        fn=lambda model: on_model_change(task_key, model),
        inputs=model_choice,
        outputs=[model_info, schema_info, output_highlight, summary_metrics, summary_table],
    )

    run_btn.click(
        fn=lambda text, model: run_inference(text, task_key, model),
        inputs=[input_text, model_choice],
        outputs=[output_highlight, summary_metrics, summary_table],
    )

    clear_btn.click(
        fn=lambda model: clear_outputs(task_key, model),
        inputs=model_choice,
        outputs=[input_text, output_highlight, summary_metrics, summary_table],
    )


def build_ui():
    with gr.Blocks(
        theme=GR_THEME,
        css=CSS,
        title=APP_TITLE,
    ) as demo:
        gr.HTML(render_hero())

        with gr.Tabs():
            with gr.Tab("3-class models"):
                build_task_panel("3class")

            with gr.Tab("2-class models"):
                build_task_panel("2class")

            with gr.Tab("Model catalogue"):
                gr.Dataframe(
                    value=model_catalog_rows(),
                    headers=["Family", "Model", "Tag", "Supported entities", "Description", "Repo"],
                    interactive=False,
                    label="Model catalogue",
                )
            with gr.Tab("About", elem_id="about-tab"):
                gr.Markdown(
                    """
                    ### About this demo
                    
                    This Space showcases CeLLaTe biomedical NER models for extracting entity spans from biomedical text.
                    
                    **Current behavior**
                    - The tagged-entity box uses Gradio's default highlight coloring rather than a hard-coded color map.
                    - The selected model's supported schema is shown explicitly in a separate card.
                    - The summary panel sits beside the tagged output and includes average confidence per entity type found.
                    
                    """
                                    )


    return demo


if __name__ == "__main__":
    app = build_ui()
    app.launch()