data_studio/ner.py

from __future__ import annotations

from functools import lru_cache
import re
from typing import Any

from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline

from data_studio.config import StudioSettings
from data_studio.utils import utc_now_iso


def _normalize_entity_label(label: str) -> str:
    normalized = label.upper().replace("B-", "").replace("I-", "")
    return normalized


def _normalize_entity_text(text: str) -> str:
    cleaned = text.replace("##", "")
    cleaned = re.sub(r"\s+", " ", cleaned).strip()
    cleaned = re.sub(r"\s+([,.;:!?)])", r"\1", cleaned)
    cleaned = re.sub(r"([(])\s+", r"\1", cleaned)
    return cleaned


def _extract_entity_text(text: str, item: dict[str, Any]) -> str:
    start = item.get("start")
    end = item.get("end")
    if isinstance(start, int) and isinstance(end, int) and 0 <= start < end <= len(text):
        span = _normalize_entity_text(text[start:end])
        if span:
            return span
    return _normalize_entity_text(str(item.get("word", "")))


def _is_useful_entity_text(text: str) -> bool:
    if not text or "#" in text:
        return False
    if len(text.replace(" ", "")) < 2:
        return False
    if not re.search(r"[\u0C00-\u0C7FA-Za-z]", text):
        return False
    if re.fullmatch(r"[\W_]+", text):
        return False
    return True


@lru_cache(maxsize=2)
def _ner_pipeline(model_id: str, token: str):
    tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
    model = AutoModelForTokenClassification.from_pretrained(model_id, token=token)
    return pipeline(
        "token-classification",
        model=model,
        tokenizer=tokenizer,
        aggregation_strategy="simple",
        device=-1,
    )


def should_run_ner(sample: dict[str, Any]) -> bool:
    return sample.get("queue_name") == "pure_telugu" and sample.get("sample_type") == "document"


def build_ner_suggestions(settings: StudioSettings, text: str) -> dict[str, Any]:
    if not settings.active_ner_model_id:
        return {
            "model_version": None,
            "generated_at": None,
            "entities": [],
            "suggested_tags": [],
            "status": "disabled",
            "error": "NER model is not configured.",
        }

    ner_pipe = _ner_pipeline(settings.active_ner_model_id, settings.hf_token)
    raw_entities = ner_pipe(text)
    entities: list[dict[str, Any]] = []
    suggested_tags: list[str] = []
    seen = set()
    for item in raw_entities:
        label = _normalize_entity_label(str(item.get("entity_group", "")))
        score = round(float(item.get("score", 0.0)), 4)
        entity_text = _extract_entity_text(text, item)
        if not entity_text or not label:
            continue
        if score < 0.60:
            continue
        if not _is_useful_entity_text(entity_text):
            continue
        key = (label, entity_text.casefold())
        if key in seen:
            continue
        seen.add(key)
        entities.append(
            {
                "text": entity_text,
                "label": label,
                "score": score,
                "start": item.get("start"),
                "end": item.get("end"),
            }
        )
        suggested_tags.append(entity_text)
    return {
        "model_version": settings.active_ner_model_id,
        "postprocess_version": "v2",
        "generated_at": utc_now_iso(),
        "entities": entities,
        "suggested_tags": suggested_tags[:24],
        "status": "ready",
        "error": None,
    }