services/extractor.py

"""Skill extraction from job descriptions using GLiNER zero-shot NER.

This is domain-agnostic by design: GLiNER is a general-purpose zero-shot NER
model, not something fine-tuned on tech job postings, so the same extractor
works across software, healthcare, retail, hospitality, skilled trades,
logistics, etc. -- as long as the *label set* it's prompted with doesn't tilt
toward one domain. See `SkillExtractorConfig.labels` below.

Usage:
    extractor = SkillExtractor.get_instance()
    result = extractor.extract(job_description=text, job_title="Senior Backend Engineer")
    result.skills   # ["Python", "AWS", "Distributed Systems", ...]
    result.scores   # {"Python": 0.91, "AWS": 0.74, ...}  (model confidence per skill)

    # Works the same for a non-tech posting, no config changes needed:
    result = extractor.extract(job_description=nursing_posting_text, job_title="RN")
    result.skills   # ["BLS Certification", "Patient Charting", "Spanish", "EPIC", ...]

Design notes (why this differs from a "just call predict_entities" version):

* The label set sent to the model is deliberately industry-neutral: "tool,
  equipment or software" (covers Python/AWS *and* forklifts/POS systems/
  espresso machines), "certification or license" (AWS cert *and* RN license,
  CDL, ServSafe), plus "soft skill" and "language". A label like "technical
  skill" would quietly bias the model toward software-engineering postings and
  under-extract from everything else.

* Model loading is slow and memory-heavy, so it happens once per process behind a
  thread-safe singleton, and the heavy `gliner`/torch import is deferred until the
  first real instantiation -- importing this module (e.g. for its config/result
  types, or in unit tests) never requires torch to be installed.
* GLiNER's underlying transformer has a fixed token budget. Real job descriptions
  (responsibilities + requirements + benefits + boilerplate) routinely exceed it.
  Silently truncating means everything after the cut point is invisible to the
  model, so long input is split on paragraph/sentence boundaries into bounded
  chunks and the results are merged.
* A failure on one chunk (transient OOM, odd encoding, etc.) is logged and skipped
  rather than failing the whole request. Only total failure raises
  SkillExtractionError -- callers should be able to tell "extraction broke" apart
  from "this posting genuinely has zero recognizable skills".
* The original `.title()`-everything approach corrupts real skill casing:
  "AWS" -> "Aws", "Node.js" -> "Node.Js", ".NET" -> ".Net", "PostgreSQL" ->
  "Postgresql". Spans that already have internal casing are left untouched; only
  fully-lowercase spans get title-cased for display.
* Confidence scores from the model are preserved, used to pick the best-cased
  duplicate when the same skill surfaces with different capitalization across
  chunks, and used to rank/cap results so one giant posting can't return an
  unbounded list.
"""

from __future__ import annotations

import logging
import os
import re
import threading
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar, Optional

from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from gliner import GLiNER

logger = logging.getLogger(__name__)

__all__ = [
    "SkillExtractor",
    "SkillExtractorConfig",
    "SkillExtractionResult",
    "SkillExtractionError",
    "ModelLoadError",
]


# --------------------------------------------------------------------------- #
# Errors
# --------------------------------------------------------------------------- #


class SkillExtractionError(Exception):
    """Raised when skill extraction fails for every chunk of the input text."""


class ModelLoadError(SkillExtractionError):
    """Raised when the underlying GLiNER model fails to load."""


# --------------------------------------------------------------------------- #
# Config
# --------------------------------------------------------------------------- #


def _env_float(name: str, default: float) -> float:
    raw = os.environ.get(name)
    if raw is None:
        return default
    try:
        return float(raw)
    except ValueError:
        logger.warning("Invalid float for %s=%r; using default %s", name, raw, default)
        return default


def _env_int(name: str, default: int) -> int:
    raw = os.environ.get(name)
    if raw is None:
        return default
    try:
        return int(raw)
    except ValueError:
        logger.warning("Invalid int for %s=%r; using default %s", name, raw, default)
        return default


def _env_labels(name: str, default: tuple[str, ...]) -> tuple[str, ...]:
    raw = os.environ.get(name)
    if not raw:
        return default
    labels = tuple(label.strip() for label in raw.split(",") if label.strip())
    return labels or default


@dataclass(frozen=True)
class SkillExtractorConfig:
    """Runtime configuration, overridable via env vars so behavior can be tuned
    per-deployment without a code change."""

    model_name: str = field(
        default_factory=lambda: os.environ.get(
            "SKILL_EXTRACTOR_MODEL", "urchade/gliner_small-v2.1"
        )
    )
    device: str = field(
        default_factory=lambda: os.environ.get("SKILL_EXTRACTOR_DEVICE", "cpu")
    )
    threshold: float = field(
        default_factory=lambda: _env_float("SKILL_EXTRACTOR_THRESHOLD", 0.3)
    )
    # Conservative character budget per chunk. The small GLiNER model is trained on
    # sequences up to a few hundred subword tokens; ~4 chars/token in English gives
    # headroom without needing the tokenizer loaded just to size chunks.
    max_chunk_chars: int = field(
        default_factory=lambda: _env_int("SKILL_EXTRACTOR_MAX_CHUNK_CHARS", 1200)
    )
    # Hard cap on total input size, to bound worst-case latency on malformed/huge input.
    max_total_chars: int = field(
        default_factory=lambda: _env_int("SKILL_EXTRACTOR_MAX_TOTAL_CHARS", 20_000)
    )
    max_skills_returned: int = field(
        default_factory=lambda: _env_int("SKILL_EXTRACTOR_MAX_SKILLS", 100)
    )
    cache_dir: Optional[str] = field(
        default_factory=lambda: os.environ.get("SKILL_EXTRACTOR_CACHE_DIR")
    )
    # Deliberately industry-neutral. "tool, equipment or software" covers
    # both "Python"/"AWS" and "forklift"/"point-of-sale system"/"espresso
    # machine"; "certification or license" covers an AWS cert as readily as
    # an RN license, CDL, or food-handler's permit. Avoid narrow labels like
    # "technical skill" -- they nudge the model toward tech postings and
    # under-extract from healthcare, retail, hospitality, trades, etc.
    # Override with a comma-separated SKILL_EXTRACTOR_LABELS env var if a
    # deployment wants to tune this further.
    labels: tuple[str, ...] = field(
        default_factory=lambda: _env_labels(
            "SKILL_EXTRACTOR_LABELS",
            (
                "skill",
                "qualification",
                "certification or license",
                "tool, equipment or software",
                "soft skill",
                "language",
            ),
        )
    )


DEFAULT_CONFIG = SkillExtractorConfig()


# --------------------------------------------------------------------------- #
# Result model
# --------------------------------------------------------------------------- #


class SkillExtractionResult(BaseModel):
    skills: list[str] = Field(default_factory=list)
    # Per-skill model confidence (0-1), keyed by the exact string in `skills`.
    # Useful for downstream filtering/ranking and for debugging low-quality
    # extractions without re-running the model.
    scores: dict[str, float] = Field(default_factory=dict)


# --------------------------------------------------------------------------- #
# Text cleanup
# --------------------------------------------------------------------------- #

# Generic noise/ATS-template filtering -- these are extraction artifacts and
# Workday/ATS boilerplate that show up regardless of industry, not a
# tech-specific allow/deny list. Safe to extend for other ATS platforms
# (Greenhouse, Lever, iCIMS, etc.) as they're observed in production data.
_STOPLIST = frozenset(
    {
        "r", "e", "m", "com", "dos", "nas", "modo", "history", "job description",
        "requisition", "advertised", "layer", "scale", "systems learn",
        "human intelligence", "subject matter", "team building",
    }
)

_BOILERPLATE_LINE_PATTERNS = [
    re.compile(r"^\s*apply\s*$", re.IGNORECASE),
    re.compile(r"^\s*locations?\s*$", re.IGNORECASE),
    re.compile(r"^\s*time\s*type\s*$", re.IGNORECASE),
    re.compile(r"^\s*posted\s+on\s*$", re.IGNORECASE),
    re.compile(r"^\s*time\s+left\s+to\s+apply\s*$", re.IGNORECASE),
    re.compile(r"^\s*job\s+requisition\s+id\s*$", re.IGNORECASE),
    re.compile(r"^\s*end\s+date\s*:.*$", re.IGNORECASE),
    re.compile(r"^\s*posted\s+today\s*$", re.IGNORECASE),
    re.compile(r"^\s*r-\d+\s*$", re.IGNORECASE),
    re.compile(r"^\s*my\s+career\s+development\s+portal\s*:?\s*$", re.IGNORECASE),
    re.compile(r"^\s*career\s+development\s+portal\b.*$", re.IGNORECASE),
]

_SENTENCE_SPLIT_RE = re.compile(r"(?<=[.!?])\s+")


def _strip_boilerplate(text: str) -> str:
    lines = text.splitlines()
    kept = [
        line
        for line in lines
        if not any(pattern.match(line) for pattern in _BOILERPLATE_LINE_PATTERNS)
    ]
    cleaned = re.sub(r"\n{3,}", "\n\n", "\n".join(kept))
    return cleaned.strip()


def _normalize_casing(text: str) -> str:
    """Title-case only spans that came back fully lowercase (e.g. "machine
    learning" -> "Machine Learning"). Anything with existing internal casing is
    left untouched, since blind .title() mangles real skill names: "AWS" ->
    "Aws", "Node.js" -> "Node.Js", ".NET" -> ".Net", "PostgreSQL" ->
    "Postgresql"."""
    if text == text.lower():
        return text.title()
    return text


def _clean_span(text: str) -> str:
    text = re.sub(r"\s+", " ", text).strip()
    return _normalize_casing(text)


def _chunk_text(text: str, max_chars: int) -> list[str]:
    """Split text into pieces no longer than max_chars, breaking on paragraph
    then sentence boundaries so a skill phrase doesn't get cut in half. Falls
    back to a hard split only if a single sentence itself exceeds max_chars."""
    if not text.strip():
        return []
    if len(text) <= max_chars:
        return [text]

    paragraphs = [p for p in text.split("\n\n") if p.strip()]
    chunks: list[str] = []
    current = ""

    def flush() -> None:
        nonlocal current
        if current.strip():
            chunks.append(current.strip())
        current = ""

    def add_piece(piece: str, joiner: str) -> None:
        nonlocal current
        candidate = f"{current}{joiner}{piece}" if current else piece
        if len(candidate) <= max_chars:
            current = candidate
            return
        flush()
        if len(piece) <= max_chars:
            current = piece
            return
        # Single sentence/paragraph still too long: hard-split as last resort.
        for i in range(0, len(piece), max_chars):
            chunks.append(piece[i : i + max_chars])

    for para in paragraphs:
        candidate = f"{current}\n\n{para}" if current else para
        if len(candidate) <= max_chars:
            current = candidate
            continue
        flush()
        if len(para) <= max_chars:
            current = para
            continue
        for sentence in _SENTENCE_SPLIT_RE.split(para):
            if sentence.strip():
                add_piece(sentence.strip(), " ")

    flush()
    return chunks


def _clean_skills(
    scored_spans: list[tuple[str, float]],
    *,
    max_results: int,
) -> tuple[list[str], dict[str, float]]:
    """Dedupe case-insensitively, drop stoplisted/junk spans, keep the
    highest-confidence (text, score) per skill, then rank by score and cap."""
    best: dict[str, tuple[str, float]] = {}

    for raw_text, score in scored_spans:
        text = _clean_span(raw_text)
        if len(text) < 3 or not re.search(r"[a-zA-Z]", text):
            continue
        key = text.lower()
        if key in _STOPLIST:
            continue
        existing = best.get(key)
        if existing is None or score > existing[1]:
            best[key] = (text, score)

    ranked = sorted(best.values(), key=lambda pair: pair[1], reverse=True)
    if max_results > 0:
        ranked = ranked[:max_results]

    skills = [text for text, _ in ranked]
    scores = {text: score for text, score in ranked}
    return skills, scores


# --------------------------------------------------------------------------- #
# Extractor
# --------------------------------------------------------------------------- #


class SkillExtractor:
    """Thread-safe singleton wrapper around a GLiNER zero-shot NER model,
    specialized for pulling skills/qualifications out of job descriptions.

    Use `SkillExtractor.get_instance()` rather than constructing directly so the
    (expensive) model load happens exactly once per process.
    """

    _instance: ClassVar[Optional["SkillExtractor"]] = None
    _instance_lock: ClassVar[threading.Lock] = threading.Lock()

    def __init__(self, config: SkillExtractorConfig = DEFAULT_CONFIG):
        self._config = config
        self._inference_lock = threading.Lock()
        self._model: "GLiNER" = self._load_model(config)

    @staticmethod
    def _load_model(config: SkillExtractorConfig) -> "GLiNER":
        # Imported lazily: importing this module should never require torch/gliner
        # to be installed (keeps config/result types usable in lightweight contexts
        # and unit tests, and keeps process startup fast for callers who don't
        # immediately need inference).
        from gliner import GLiNER

        logger.info("Loading GLiNER model %r on device %r", config.model_name, config.device)
        kwargs: dict[str, Any] = {}
        if config.cache_dir:
            kwargs["cache_dir"] = config.cache_dir

        try:
            model = GLiNER.from_pretrained(config.model_name, **kwargs)
        except Exception as exc:  # network/hub failures, corrupt cache, OOM, etc.
            raise ModelLoadError(
                f"Failed to load GLiNER model '{config.model_name}': {exc}"
            ) from exc

        try:
            model = model.to(config.device)
        except Exception as exc:
            logger.warning(
                "Could not move model to device %r (%s); falling back to CPU",
                config.device,
                exc,
            )
            model = model.to("cpu")

        logger.info("GLiNER model loaded successfully")
        return model

    @classmethod
    def get_instance(cls, config: SkillExtractorConfig = DEFAULT_CONFIG) -> "SkillExtractor":
        """Returns the process-wide singleton, creating it on first call. Note:
        if an instance already exists, a different `config` passed here is
        ignored -- call `reset_instance()` first if you need to reconfigure
        (mainly useful in tests)."""
        if cls._instance is None:
            with cls._instance_lock:
                if cls._instance is None:
                    cls._instance = cls(config)
        return cls._instance

    @classmethod
    def reset_instance(cls) -> None:
        """Drops the cached singleton so the next get_instance() call reloads
        the model. Mainly for tests."""
        with cls._instance_lock:
            cls._instance = None

    def health_check(self) -> bool:
        """Cheap readiness check, e.g. for a k8s readiness/liveness probe."""
        return self._model is not None

    def extract(self, job_description: str, job_title: str = "") -> SkillExtractionResult:
        if not job_description or not job_description.strip():
            return SkillExtractionResult()

        config = self._config

        if len(job_description) > config.max_total_chars:
            logger.warning(
                "job_description length %d exceeds max_total_chars=%d; truncating",
                len(job_description),
                config.max_total_chars,
            )
            job_description = job_description[: config.max_total_chars]

        cleaned_description = _strip_boilerplate(job_description)
        full_text = (
            f"{job_title}\n\n{cleaned_description}" if job_title else cleaned_description
        )

        chunks = _chunk_text(full_text, config.max_chunk_chars)
        if not chunks:
            return SkillExtractionResult()

        scored_spans: list[tuple[str, float]] = []
        failures = 0

        for chunk in chunks:
            try:
                with self._inference_lock:
                    entities = self._model.predict_entities(
                        chunk, list(config.labels), threshold=config.threshold
                    )
            except Exception as exc:
                failures += 1
                logger.warning("Inference failed on a %d-char chunk: %s", len(chunk), exc)
                continue

            for entity in entities:
                text = (entity.get("text") or "").strip()
                if not text:
                    continue
                score = float(entity.get("score", 0.0))
                scored_spans.append((text, score))

        if failures and failures == len(chunks):
            raise SkillExtractionError(
                f"Skill extraction failed on all {len(chunks)} chunk(s) of input"
            )

        skills, scores = _clean_skills(scored_spans, max_results=config.max_skills_returned)
        return SkillExtractionResult(skills=skills, scores=scores)