src/models/turn_encoder.py

"""Per-turn text encoder.

Architecture: a **frozen** fine-tuned ``roberta-base``; the ``[CLS]`` hidden state
of each turn becomes ``e_t in R^768``. We expose a small :class:`TurnEncoder`
``Protocol`` so the rest of the pipeline (and tests) can swap in a deterministic
:class:`FakeEncoder` that needs no ``transformers`` / ``torch``.

``transformers``/``torch`` are imported lazily inside :class:`RobertaEncoder` so
this module imports cheaply.
"""
from __future__ import annotations

import hashlib
from typing import List, Optional, Protocol, TYPE_CHECKING, runtime_checkable

if TYPE_CHECKING:  # pragma: no cover
    import numpy as np

from ..config import DIMS


@runtime_checkable
class TurnEncoder(Protocol):
    """Contract: map a turn's text to a fixed-size embedding vector."""

    dim: int

    def encode(self, text: str) -> "np.ndarray":  # [dim]
        ...

    def encode_batch(self, texts: List[str]) -> "np.ndarray":  # [B, dim]
        ...


class FakeEncoder:
    """Deterministic hash -> vector encoder for unit tests (no heavy deps).

    Maps text to a reproducible pseudo-random unit-norm vector via SHA-256 seeding,
    so tests get stable, distinct embeddings for distinct turns without a model.
    """

    def __init__(self, dim: int = DIMS.turn_embed) -> None:
        self.dim = dim

    def _seed(self, text: str) -> int:
        h = hashlib.sha256((text or "").encode("utf-8")).hexdigest()
        return int(h[:16], 16)

    def encode(self, text: str) -> "np.ndarray":
        import numpy as np

        rng = np.random.default_rng(self._seed(text))
        v = rng.standard_normal(self.dim).astype("float32")
        n = np.linalg.norm(v)
        return v / n if n > 0 else v

    def encode_batch(self, texts: List[str]) -> "np.ndarray":
        import numpy as np

        if not texts:
            return np.zeros((0, self.dim), dtype="float32")
        return np.vstack([self.encode(t) for t in texts])


class RobertaEncoder:
    """Frozen fine-tuned RoBERTa turn encoder (lazy ``transformers``/``torch``).

    Loads from ``models/roberta_turn_encoder`` if present, else falls back to the
    base ``roberta-base`` (so the API can boot untrained). Parameters are frozen
    and the model is put in ``eval`` mode; ``encode`` returns the ``[CLS]`` vector.
    """

    def __init__(
        self,
        model_name_or_path: Optional[str] = None,
        device: Optional[str] = None,
        max_length: int = 256,
    ) -> None:
        self.dim = DIMS.turn_embed
        self.max_length = max_length
        self._model_ref = model_name_or_path
        self._device = device
        self._tokenizer = None
        self._model = None

    def _ensure_loaded(self) -> None:  # pragma: no cover - needs transformers/torch
        if self._model is not None:
            return
        import torch
        from transformers import AutoModel, AutoTokenizer

        from .. import config as cfg

        ref = self._model_ref
        if ref is None:
            ref = (
                str(cfg.ROBERTA_ENCODER_DIR)
                if cfg.ROBERTA_ENCODER_DIR.exists()
                else cfg.CONFIG.roberta_model_name
            )
        self._tokenizer = AutoTokenizer.from_pretrained(ref)
        self._model = AutoModel.from_pretrained(ref)
        for p in self._model.parameters():
            p.requires_grad_(False)
        self._model.eval()
        self._device = self._device or ("cuda" if torch.cuda.is_available() else "cpu")
        self._model.to(self._device)

    def encode(self, text: str) -> "np.ndarray":  # pragma: no cover
        return self.encode_batch([text])[0]

    def encode_batch(self, texts: List[str]) -> "np.ndarray":  # pragma: no cover
        import numpy as np
        import torch

        self._ensure_loaded()
        if not texts:
            return np.zeros((0, self.dim), dtype="float32")
        enc = self._tokenizer(
            list(texts),
            padding=True,
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt",
        ).to(self._device)
        with torch.no_grad():
            out = self._model(**enc)
        cls = out.last_hidden_state[:, 0, :]  # [B, 768] CLS token
        return cls.detach().cpu().numpy().astype("float32")