Blame for app.py

import os
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import gc
import html
import json
from functools import lru_cache
from pathlib import Path

import gradio as gr
import h5py
import numpy as np
import torch
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoTokenizer,
    pipeline,
)

BASE_DIR = Path(__file__).resolve().parent
CFG = json.loads((BASE_DIR / "config.json").read_text(encoding="utf-8"))

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

TASK_A_LABELS = {int(k): v for k, v in CFG["task_a_labels"].items()}
TASK_B_LABELS = {int(k): v for k, v in CFG["task_b_labels"].items()}

# sagorsarker/mbert-bengali-ner semantic labels
NER_ID2LABEL = {
    0: "O",
    1: "B-PER",
    2: "I-PER",
    3: "B-ORG",
    4: "I-ORG",
    5: "B-LOC",
    6: "I-LOC",
}
NER_LABEL2ID = {v: k for k, v in NER_ID2LABEL.items()}


# ============================================================
# LOCAL H5 WEIGHT LOADING
# ============================================================
def get_local_weight_path(filename: str) -> Path:
    path = BASE_DIR / filename
    if not path.exists():
        raise FileNotFoundError(
            f"{filename} was not found in the Space repository root. "
            f"Upload {filename} beside app.py."
        )
    return path


def list_h5_datasets(h5_file):
    names = []

    def visitor(name, obj):
        if isinstance(obj, h5py.Dataset):
            names.append(name)

    h5_file.visititems(visitor)
    return names


def load_pytorch_h5_weights(model: torch.nn.Module, h5_path: Path):
    """
    Load the custom HDF5 format used in training:

        with h5py.File(h5_path, "w") as f:
            for name, param in model.state_dict().items():
                f.create_dataset(name, data=param.cpu().numpy())

    This is NOT a Keras H5 model.

    Tensors are copied directly into the rebuilt PyTorch model state one-by-one.
    """
    state = model.state_dict()
    model_keys = set(state.keys())

    with h5py.File(h5_path, "r") as f:
        file_keys = set(list_h5_datasets(f))

        # Compatibility-only buffers can differ across transformers versions.
        def benign(key):
            return key.endswith("position_ids") or key.endswith("token_type_ids")

        missing = sorted(
            key for key in (model_keys - file_keys)
            if not benign(key)
        )
        unexpected = sorted(
            key for key in (file_keys - model_keys)
            if not benign(key)
        )

        if missing or unexpected:
            raise RuntimeError(
                "\n".join([
                    f"Architecture mismatch for {h5_path.name}.",
                    f"Missing learned keys ({len(missing)}): {missing[:15]}",
                    f"Unexpected learned keys ({len(unexpected)}): {unexpected[:15]}",
                    "The model must be recreated with the same Transformers "
                    "architecture/version family used during training."
                ])
            )

        with torch.no_grad():
            for key in file_keys:
                if key not in state:
                    continue

                arr = np.asarray(f[key])
                src = torch.from_numpy(arr)
                dst = state[key]

                if tuple(src.shape) != tuple(dst.shape):
                    raise RuntimeError(
                        f"Shape mismatch at {key}: "
                        f"H5={tuple(src.shape)} vs model={tuple(dst.shape)}"
                    )

                dst.copy_(src.to(dtype=dst.dtype, device=dst.device))

    del state
    gc.collect()
    return model


# ============================================================
# XLM-R TOKENIZER + CLASSIFIERS
# ============================================================
@lru_cache(maxsize=1)
def get_xlmr_tokenizer():
    return AutoTokenizer.from_pretrained(CFG["base_model"])


def build_classifier(num_labels: int, labels: dict):
    id2label = {int(i): label for i, label in labels.items()}
    label2id = {label: i for i, label in id2label.items()}

    # Same architecture as:
    # AutoModelForSequenceClassification.from_pretrained(
    #     "xlm-roberta-base", num_labels=...
    # )
    #
    # We instantiate from config so we don't first download a second full copy
    # of XLM-R base weights. The local H5 contains the COMPLETE fine-tuned state_dict.
    config = AutoConfig.from_pretrained(
        CFG["base_model"],
        num_labels=num_labels,
        id2label=id2label,
        label2id=label2id,
    )

    return AutoModelForSequenceClassification.from_config(config)


@lru_cache(maxsize=1)
def get_task_a_model():
    model = build_classifier(2, TASK_A_LABELS)
    model = load_pytorch_h5_weights(
        model,
        get_local_weight_path(CFG["task_a_file"])
    )
    model.to(DEVICE)
    model.eval()
    return model


@lru_cache(maxsize=1)
def get_task_b_model():
    # Lazy load: this 10-class model is loaded only when Task A says Accidental.
    model = build_classifier(10, TASK_B_LABELS)
    model = load_pytorch_h5_weights(
        model,
        get_local_weight_path(CFG["task_b_file"])
    )
    model.to(DEVICE)
    model.eval()
    return model


def predict_classifier(text: str, model, labels: dict):
    tokenizer = get_xlmr_tokenizer()

    encoded = tokenizer(
        text,
        max_length=int(CFG["max_length"]),
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )
    encoded = {k: v.to(DEVICE) for k, v in encoded.items()}

    with torch.inference_mode():
        outputs = model(
            input_ids=encoded["input_ids"],
            attention_mask=encoded["attention_mask"],
        )
        probs = torch.softmax(outputs.logits, dim=-1)[0].cpu().numpy()

    pred_id = int(np.argmax(probs))
    scores = {
        labels[i]: float(probs[i])
        for i in range(len(labels))
    }

    return pred_id, labels[pred_id], scores


# ============================================================
# BENGALI NER / EXACT LOCATION EXTRACTION
# ============================================================
@lru_cache(maxsize=1)
def get_ner_pipeline():
    tokenizer = AutoTokenizer.from_pretrained(CFG["ner_model"])
    model = AutoModelForTokenClassification.from_pretrained(CFG["ner_model"])

    # Force semantic NER label mapping so LOC filtering works reliably.
    model.config.id2label = NER_ID2LABEL
    model.config.label2id = NER_LABEL2ID

    return pipeline(
        "token-classification",
        model=model,
        tokenizer=tokenizer,
        aggregation_strategy="simple",
        device=0 if DEVICE.type == "cuda" else -1,
    )


def extract_locations(text: str):
    entities = get_ner_pipeline()(text)

    spans = []

    for ent in entities:
        group = str(ent.get("entity_group", "")).upper()

        if group != "LOC":
            continue

        start = int(ent["start"])
        end = int(ent["end"])
        score = float(ent.get("score", 0.0))

        if 0 <= start < end <= len(text):
            spans.append((start, end, score))

    # Merge overlapping/adjacent location pieces.
    spans.sort(key=lambda x: (x[0], x[1]))
    merged = []

    for start, end, score in spans:
        if merged and start <= merged[-1][1]:
            ps, pe, pscore = merged[-1]
            merged[-1] = (
                ps,
                max(pe, end),
                max(pscore, score)
            )
        else:
            merged.append((start, end, score))

    locations = []
    for start, end, _ in merged:
        location = text[start:end].strip()
        if location and location not in locations:
            locations.append(location)

    return locations, merged


def highlight_locations(text: str, spans):
    if not spans:
        return (
            '<div class="location-box">'
            + html.escape(text)
            + '<div class="empty-location">No location detected.</div>'
            + '</div>'
        )

    pieces = []
    cursor = 0

    for start, end, score in spans:
        pieces.append(html.escape(text[cursor:start]))

        loc = html.escape(text[start:end])

        pieces.append(
            f'<mark title="NER confidence: {score:.3f}">{loc}</mark>'
        )
        cursor = end

    pieces.append(html.escape(text[cursor:]))

    return (
        '<div class="location-box">'
        + "".join(pieces)
        + "</div>"
    )


# ============================================================
# MAIN PIPELINE
# ============================================================
def analyze_news(text):
    text = (text or "").strip()

    if not text:
        raise gr.Error("Please enter a Bangla news text.")

    # ---------------- Task A ----------------
    a_id, a_label, a_scores = predict_classifier(
        text,
        get_task_a_model(),
        TASK_A_LABELS
    )

    # ---------------- Task B ----------------
    if a_id == 1:
        _, b_label, b_scores = predict_classifier(
            text,
            get_task_b_model(),
            TASK_B_LABELS
        )
        b_result = b_label
    else:
        b_result = "Not Applicable"
        b_scores = {
            "Not Applicable — Task A predicted Non-Accidental": 1.0
        }

    # ---------------- NER ----------------
    locations, spans = extract_locations(text)

    location_text = (
        " | ".join(locations)
        if locations
        else "No location detected"
    )

    summary = f"""
### Analysis Result

**Subtask A — Accident Identification:** {a_label}

**Subtask B — Accident Category:** {b_result}

**Exact Location(s):** {location_text}
"""

    return (
        summary,
        a_scores,
        b_scores,
        location_text,
        highlight_locations(text, spans),
    )


# ============================================================
# GRADIO UI
# ============================================================
CSS = """
.gradio-container {
    max-width: 1180px !important;
    margin: 0 auto !important;
}

.header-card {
    border: 1px solid var(--border-color-primary);
    border-radius: 18px;
    padding: 24px 26px;
    margin-bottom: 18px;
    background: var(--background-fill-secondary);
}

.header-card h1 {
    margin: 0 0 10px 0;
    font-size: 1.7rem;
    line-height: 1.25;
}

.header-card p {
    margin: 0;
    opacity: .8;
}

.location-box {
    border: 1px solid var(--border-color-primary);
    border-radius: 14px;
    padding: 18px;
    min-height: 110px;
    line-height: 2.1;
    font-size: 1.05rem;
    white-space: pre-wrap;
}

.location-box mark {
    background: #facc15;
    color: #111827;
    padding: 2px 5px;
    border-radius: 5px;
    font-weight: 700;
}

.empty-location {
    opacity: .65;
    margin-top: 12px;
    font-size: .9rem;
}
"""

with gr.Blocks(css=CSS, title="Accident Causal Identification") as demo:

    gr.HTML(
        f"""
        <div class="header-card">
            <h1>{html.escape(CFG["paper_title"])}</h1>
            <p>
                Bangla news analysis using fine-tuned XLM-RoBERTa
                for accident classification and Bengali mBERT NER
                for exact location extraction.
            </p>
        </div>
        """
    )

    with gr.Row():

        with gr.Column(scale=7):

            news_text = gr.Textbox(
                label="Bangla News Text",
                placeholder="এখানে বাংলা সংবাদ লিখুন...",
                lines=10,
            )

            analyze_button = gr.Button(
                "Analyze News",
                variant="primary"
            )

        with gr.Column(scale=4):

            gr.Markdown(
                """
### Pipeline

**Subtask A**  
Accidental / Non-Accidental

**Subtask B**  
10 accident categories

**Exact Location**  
Bengali mBERT NER

**XLM-RoBERTa max length:** `256`
                """
            )

    result_summary = gr.Markdown()

    with gr.Row():

        task_a_probabilities = gr.Label(
            label="Subtask A Probabilities",
            num_top_classes=2
        )

        task_b_probabilities = gr.Label(
            label="Subtask B Probabilities",
            num_top_classes=10
        )

    exact_locations = gr.Textbox(
        label="Extracted Exact Location(s)",
        interactive=False
    )

    gr.Markdown("### Location Highlighting")

    highlighted_news = gr.HTML()

    gr.Examples(
        examples=[
            ["গাজীপুরের জয়দেবপুর স্টেশন এলাকায় ট্রেনে কাটা পড়ে এক যুবকের মৃত্যু হয়েছে।"],
            ["রাজধানী ঢাকার পল্লবী এলাকায় নির্মাণাধীন ভবন থেকে ইট পড়ে এক পথচারীর মৃত্যু হয়েছে।"],
            ["রোয়ালখালীতে হাতির হামলায় তিনজনের মৃত্যু হয়েছে।"],
            ["বাংলাদেশের অর্থনীতি নিয়ে নতুন প্রতিবেদন প্রকাশ করেছে একটি আন্তর্জাতিক সংস্থা।"],
        ],
        inputs=news_text
    )

    analyze_button.click(
        fn=analyze_news,
        inputs=news_text,
        outputs=[
            result_summary,
            task_a_probabilities,
            task_b_probabilities,
            exact_locations,
            highlighted_news,
        ]
    )


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