""" 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": , "classifier": } 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)