apps/accounts/ner/sbert.py

"""Layer 4: SBERT + pgvector cosine — degraded-fallback canonical mapping.

Uses sentence-transformers/all-MiniLM-L6-v2 (Apache 2.0, 22M params, 384 dims)
to embed candidate noun phrases from the resume text, then queries the
SkillEmbedding table via pgvector cosine distance to find the nearest
catalog skill. This is the always-available fallback for when higher BERT
layers can't load (memory-limited Render tier, offline, etc.) AND for
paraphrased skill language that lexical + BERT miss.

Threshold: 0.65 similarity (equivalent to cosine_distance <= 0.35). The
dispatcher's alias resolver uses the canonical ``skill_name`` keyed output
directly — no further fuzzy-match hop needed.
"""
from __future__ import annotations

import logging

from . import NerLayer


logger = logging.getLogger(__name__)

_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
_SIM_THRESHOLD = 0.65

_model = None
_nlp = None


def _get_model():
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer
        logger.info("sbert: loading %s (first call may be slow)", _MODEL_ID)
        _model = SentenceTransformer(_MODEL_ID)
    return _model


def _get_nlp():
    """Reuse spaCy's en_core_web_lg for noun-phrase extraction.

    If skillNer already loaded it, great — it's a module-level singleton in
    that layer too. Otherwise we load our own.
    """
    global _nlp
    if _nlp is None:
        # Reuse skillner's loaded model if possible to save ~560 MB of RAM.
        from . import skillner as _skillner_mod
        cached = getattr(_skillner_mod, "_nlp", None)
        if cached is not None:
            _nlp = cached
        else:
            import spacy
            # Prefer `_lg` for slightly better noun-chunk quality; fall back
            # to `_sm` which is the minimum viable install documented in
            # INSTALL.md.
            try:
                _nlp = spacy.load("en_core_web_lg")
            except OSError:
                _nlp = spacy.load("en_core_web_sm")
    return _nlp


def _candidate_phrases(text: str) -> list[str]:
    nlp = _get_nlp()
    doc = nlp(text[:4000])
    seen: set[str] = set()
    out: list[str] = []
    for chunk in doc.noun_chunks:
        phrase = chunk.text.strip(" ,.;:-").lower()
        # Skip very short fragments + stopwords-only spans — they lead to
        # spurious cosine matches against broad skill names.
        if len(phrase) < 3 or phrase in seen:
            continue
        seen.add(phrase)
        out.append(phrase)
    return out


class SBertLayer(NerLayer):
    name = "sbert"

    def predict(self, text: str) -> dict[str, float]:
        if not text.strip():
            return {}
        # DB lookup: pgvector cosine against SkillEmbedding. Import Django
        # models lazily so module import doesn't touch the DB.
        from pgvector.django import CosineDistance

        from apps.skills.models import Skill, SkillEmbedding

        phrases = _candidate_phrases(text)
        if not phrases:
            return {}

        model = _get_model()
        embeddings = model.encode(phrases, normalize_embeddings=True).tolist()

        out: dict[str, float] = {}
        # One query per phrase — each hits the HNSW index. 20 phrases × ~2 ms
        # per query = ~40 ms total, acceptable for on-demand parse.
        for phrase, vec in zip(phrases, embeddings):
            top = (
                SkillEmbedding.objects
                .annotate(distance=CosineDistance("embedding", vec))
                .order_by("distance")
                .select_related("skill")
                .first()
            )
            if top is None:
                continue
            # Cosine distance = 1 - cosine similarity for normalized vectors.
            similarity = 1.0 - float(top.distance)
            if similarity < _SIM_THRESHOLD:
                continue
            name = top.skill.skill_name
            if similarity > out.get(name, 0.0):
                out[name] = similarity
        return out

    def available(self) -> bool:
        try:
            import sentence_transformers  # noqa: F401
            import pgvector  # noqa: F401
            from apps.skills.models import SkillEmbedding  # noqa: F401
            return True
        except Exception:
            return False


layer = SBertLayer()