sentiment_deploy.py

"""
sentiment_deploy.py
===================

Self-contained, picklable deployment wrapper for the Route C (BERTweet-large)
sentiment classifier, compatible with the case-manual API template.

The API loads a single ``*.model`` pickle and expects a dict::

    {"vectorizer": <obj with .transform(list[str])>,
     "classifier": <obj with .predict(X)>}

A HuggingFace transformer does not fit that interface, so this module
provides two small adapters:

* ``BertweetVectorizer`` -- a pass-through "vectorizer". It applies the SAME
  light cleaning used at training time and returns the (cleaned) strings.
  ``fit`` / ``fit_transform`` are no-ops, so the wrapper is safe even if the
  API template erroneously calls ``fit_transform`` at inference time.

* ``BertweetClassifier`` -- holds the fine-tuned weights, config and tokenizer
  files *inside the pickle* (no external paths, no dependence on a checkpoint
  directory). It rebuilds the model lazily on first use and maps the internal
  class indices {0,1,2} back to the API label space {-1, 0, 1}.

This module is model-agnostic: the config + state_dict + tokenizer are captured
dynamically from whatever model you pass in, so it serves ``vinai/bertweet-base``
and ``vinai/bertweet-large`` identically. The only large-specific default is
``max_length=512`` (base maxes out at 128 tokens; large supports up to 512).

fp16 storage note: shrink_model.py can cast the stored weights to float16 to
halve the file size. PyTorch on CPU cannot run float16 matmuls, so when the
weights are stored as fp16 (marked by ``_weights_dtype == "float16"``),
``_build_model`` up-casts them back to float32 at load time. The on-disk file
stays small; serve-time inference is unaffected.

IMPORTANT (pickle/__main__ caveat): because the API loads the pickle in a
*separate process*, the classes referenced by the pickle must be importable
there. Defining them in THIS module (not in a notebook's __main__) is what
makes the round-trip work. Ship ``sentiment_deploy.py`` alongside the API
``app.py``.
"""

from __future__ import annotations

import io
import os
import tempfile
from typing import List, Sequence

# Heavy deps (torch / transformers / bs4) are imported lazily inside methods
# so this module can be imported in lightweight contexts and unit-tested.

# Internal index -> API label.  Training uses 0=Negative, 1=Neutral, 2=Positive.
# The API/case-manual label space is -1=Negative, 0=Neutral, 1=Positive.
INDEX_TO_API_LABEL = {0: -1, 1: 0, 2: 1}

# Default max sequence length. BERTweet-large (RoBERTa-large backbone,
# max_position_embeddings=514) supports up to 512 tokens; base maxes at 128.
DEFAULT_MAX_LENGTH = 512


def normalize_text(x) -> str:
    """Light, rule-based cleaning applied identically at train and serve time.

    Only does what BERTweet's own tokenizer normalisation does NOT do: strip
    HTML (reviews contain markup) and collapse whitespace. Mention/URL/emoji
    handling is delegated to the tokenizer (``normalization=True``) so that
    train and serve stay perfectly consistent and there is no MNTN/URL skew.
    """
    if x is None:
        return ""
    x = str(x)
    if "<" in x and ">" in x:  # only pay BeautifulSoup cost when markup is likely
        try:
            from bs4 import BeautifulSoup
            x = BeautifulSoup(x, "html.parser").get_text(separator=" ")
        except Exception:
            pass
    x = " ".join(x.split())  # collapse all whitespace runs
    return x


class BertweetVectorizer:
    """Pass-through 'vectorizer' for API compatibility.

    Tokenisation happens inside the classifier, so ``transform`` just returns
    the cleaned strings. ``fit``/``fit_transform`` are no-ops -- importantly,
    ``fit_transform`` does NOT re-fit anything, so the buggy template call
    ``vectorizer.fit_transform(text)`` at inference behaves like ``transform``.
    """

    def fit(self, X=None, y=None):
        return self

    def transform(self, X: Sequence[str]) -> List[str]:
        if isinstance(X, str):
            X = [X]
        return [normalize_text(t) for t in X]

    def fit_transform(self, X: Sequence[str], y=None) -> List[str]:
        return self.transform(X)


class BertweetClassifier:
    """Self-contained, picklable BERTweet sequence classifier.

    Parameters
    ----------
    model : transformers PreTrainedModel (fine-tuned)
    tokenizer : transformers PreTrainedTokenizer
    max_length : int
        Token cap at inference. 512 for bertweet-large (default), 128 for base.
    batch_size : int
        Inference batch size. Keep modest for large on CPU (the HF free tier).
    """

    def __init__(self, model=None, tokenizer=None,
                 max_length: int = DEFAULT_MAX_LENGTH, batch_size: int = 16):
        self.max_length = int(max_length)
        self.batch_size = int(batch_size)
        self.index_to_api = dict(INDEX_TO_API_LABEL)

        # Serialised payload (populated from the live objects). Kept so the
        # pickle is fully self-contained.
        self._config = None
        self._state_dict = None       # dict[str, torch.Tensor] on CPU
        self._tokenizer_files = None  # dict[str, bytes]
        self._weights_dtype = None    # set to "float16" by shrink_model.py

        if model is not None and tokenizer is not None:
            self._capture(model, tokenizer)

        # Live objects (rebuilt lazily; never pickled).
        self._model = None
        self._tok = None

    # ---- serialisation helpers -------------------------------------------
    def _capture(self, model, tokenizer):
        """Snapshot weights/config/tokenizer into picklable payload."""
        self._config = model.config
        self._state_dict = {k: v.detach().cpu() for k, v in model.state_dict().items()}
        with tempfile.TemporaryDirectory() as d:
            tokenizer.save_pretrained(d)
            files = {}
            for name in os.listdir(d):
                path = os.path.join(d, name)
                if os.path.isfile(path):
                    with open(path, "rb") as fh:
                        files[name] = fh.read()
            self._tokenizer_files = files

    def __getstate__(self):
        # Exclude live (non-portable) objects from the pickle.
        return {
            "max_length": self.max_length,
            "batch_size": self.batch_size,
            "index_to_api": self.index_to_api,
            "_config": self._config,
            "_state_dict": self._state_dict,
            "_tokenizer_files": self._tokenizer_files,
            "_weights_dtype": self._weights_dtype,
        }

    def __setstate__(self, state):
        self.__dict__.update(state)
        # Back-compat: older pickles won't carry _weights_dtype.
        self._weights_dtype = state.get("_weights_dtype", None)
        self._model = None
        self._tok = None

    # ---- lazy rebuild -----------------------------------------------------
    def _build_model(self, config, state_dict):
        """Rebuild the HF model from config + state_dict (no hub download).

        If the stored weights are float16 (produced by shrink_model.py for a
        smaller file), up-cast float tensors back to float32 here, because
        PyTorch on CPU cannot run float16 matmuls. This keeps the on-disk file
        small while keeping serve-time inference correct on the HF free tier.
        Float32 pickles are loaded unchanged (the cast block is skipped).
        """
        import torch
        from transformers import AutoModelForSequenceClassification

        if getattr(self, "_weights_dtype", None) == "float16":
            state_dict = {
                k: (v.to(torch.float32)
                    if torch.is_tensor(v) and v.is_floating_point() else v)
                for k, v in state_dict.items()
            }

        # The number of output labels must match the fine-tuned head, otherwise
        # from_config() defaults to 2 labels and load_state_dict() raises a size
        # mismatch on classifier.out_proj.* (which only surfaces on the first
        # prediction, i.e. as a 500 on every POST). Derive it from the stored
        # head weight so it is always correct for base or large.
        head_key = next(
            (k for k in ("classifier.out_proj.weight", "classifier.weight")
             if k in state_dict),
            None,
        )
        if head_key is not None:
            n_labels = int(state_dict[head_key].shape[0])
            config.num_labels = n_labels

        model = AutoModelForSequenceClassification.from_config(config)
        # strict=False tolerates harmless extra/missing keys (e.g. pooler /
        # position_ids buffers that differ across transformers versions); the
        # head and encoder weights still load by name.
        missing, unexpected = model.load_state_dict(state_dict, strict=False)
        # Guard: if any parameter (not just a buffer) failed to load, fail
        # loudly rather than serving a half-random model.
        real_missing = [m for m in missing if "position_ids" not in m]
        if real_missing:
            raise RuntimeError(
                f"State dict missing parameters after load: {real_missing[:8]}"
            )
        return model

    def _ensure(self):
        if self._model is not None:
            return
        import torch
        from transformers import AutoTokenizer

        # tokenizer
        self._tokdir = tempfile.mkdtemp(prefix="bertweet_tok_")
        for name, data in self._tokenizer_files.items():
            with open(os.path.join(self._tokdir, name), "wb") as fh:
                fh.write(data)
        self._tok = AutoTokenizer.from_pretrained(
            self._tokdir, normalization=True, use_fast=False)

        # model
        model = self._build_model(self._config, self._state_dict)
        self._device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(self._device)
        model.eval()
        self._model = model

    # ---- inference --------------------------------------------------------
    def predict(self, X: Sequence[str]):
        """Return a list of API labels in {-1, 0, 1} for the input texts.

        Accepts either raw strings or strings already passed through the
        vectorizer (cleaning is idempotent, so both work).
        """
        import numpy as np
        import torch

        if isinstance(X, str):
            X = [X]
        texts = [normalize_text(t) for t in X]
        self._ensure()

        preds = []
        for i in range(0, len(texts), self.batch_size):
            batch = texts[i:i + self.batch_size]
            enc = self._tok(batch, max_length=self.max_length, truncation=True,
                            padding=True, return_tensors="pt")
            enc = {k: v.to(self._device) for k, v in enc.items()}
            with torch.no_grad():
                logits = self._model(**enc).logits
            idx = logits.argmax(dim=1).cpu().numpy()
            preds.extend(int(self.index_to_api[int(j)]) for j in idx)
        return preds

    # convenience
    def predict_proba(self, X: Sequence[str]):
        import torch
        import torch.nn.functional as F
        if isinstance(X, str):
            X = [X]
        texts = [normalize_text(t) for t in X]
        self._ensure()
        out = []
        for i in range(0, len(texts), self.batch_size):
            batch = texts[i:i + self.batch_size]
            enc = self._tok(batch, max_length=self.max_length, truncation=True,
                            padding=True, return_tensors="pt")
            enc = {k: v.to(self._device) for k, v in enc.items()}
            with torch.no_grad():
                logits = self._model(**enc).logits
            out.append(F.softmax(logits, dim=1).cpu().numpy())
        return out and __import__("numpy").vstack(out)