apps/accounts/ner/nucha.py

"""Layer 1: Nucha/Nucha_ITSkillNER_BERT — BERT NER for IT skills.

Model: https://huggingface.co/Nucha/Nucha_ITSkillNER_BERT
  - MIT license, 108.9M params, BERT architecture
  - Labels: HSKILL (hard skill) + SSKILL (soft skill), F1 ~0.90
  - Downloads ~440 MB into %USERPROFILE%\\.cache\\huggingface on first use.

Contract: returns {skill_span_text: confidence}. The dispatcher's alias
resolver decides which canonical Skill catalog row that maps to.
"""
from __future__ import annotations

import logging

from . import NerLayer


logger = logging.getLogger(__name__)

_MODEL_ID = "Nucha/Nucha_ITSkillNER_BERT"
_pipeline = None  # set on first call, persists for the process lifetime


def _get_pipeline():
    """Lazy singleton — the first call does the ~440 MB download; later
    calls reuse the cached model. Raises on network failure; the dispatcher
    catches and moves to the next layer.
    """
    global _pipeline
    if _pipeline is None:
        from transformers import pipeline  # local import keeps lexical-only CI fast
        logger.info("nucha: loading %s pipeline (first call may be slow)", _MODEL_ID)
        _pipeline = pipeline(
            task="ner",
            model=_MODEL_ID,
            aggregation_strategy="simple",
        )
    return _pipeline


class NuchaLayer(NerLayer):
    name = "nucha"

    def predict(self, text: str) -> dict[str, float]:
        if not text.strip():
            return {}
        pipe = _get_pipeline()
        # The pipeline returns a list of span dicts with keys
        # {entity_group, score, word, start, end}. aggregation_strategy="simple"
        # merges adjacent B-/I- tokens into full spans.
        spans = pipe(text[:4000])  # cap at ~4K chars to bound latency
        out: dict[str, float] = {}
        for span in spans:
            label = span.get("entity_group") or span.get("entity") or ""
            if not label.upper().endswith("SKILL"):
                # Tokens labelled 'O' or padding — drop.
                continue
            term = (span.get("word") or "").strip(" ,.;:-")
            if len(term) < 2:
                continue
            conf = float(span.get("score") or 0.0)
            # Merge duplicates — keep the max confidence seen.
            if conf > out.get(term, 0.0):
                out[term] = conf
        return out

    def available(self) -> bool:
        try:
            import transformers  # noqa: F401
            return True
        except Exception:
            return False


# Module-level singleton for the dispatcher to import cheaply.
layer = NuchaLayer()