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 (
'
'
+ html.escape(text)
+ '
No location detected.
'
+ '
'
)
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'{loc}'
)
cursor = end
pieces.append(html.escape(text[cursor:]))
return (
'