"""``CaptionTokenizer`` — typed wrapper around ``tf.keras.layers.TextVectorization``. Why a wrapper instead of using the Keras layer directly? 1. **Stable interface for the model.** The model code calls ``tokenizer.encode(captions)`` and ``tokenizer.decode_id(idx)``. The fact that those happen to delegate to a Keras layer is an implementation detail. In Phase 5 we may swap the implementation for HuggingFace ``tokenizers`` without rewriting the encoder, decoder, or inference loop. 2. **Persistence.** The notebook saves the *vocabulary list* with pickle, but loading requires re-instantiating a layer and calling ``set_vocabulary``. That ceremony belongs inside the wrapper, not at every call site. 3. **A JSON sidecar.** Pickle is fast but opaque, and loading one can run code. We additionally write a ``vocab.json`` file (a JSON array, one token per line, UTF-8) so humans and other tools can inspect the vocabulary. ``load`` reads and validates that file first. It reads a pickle only when there is no JSON file, and refuses any pickle that names a global other than the two NumPy uses to rebuild string scalars. The wrapper preserves the notebook's behaviour exactly: ``standardize=None``, ``output_sequence_length`` defaults to ``max_length``, and ``encode`` accepts either a single string or a list of strings (matching the layer's call form used in cells 7 and 25). """ from __future__ import annotations import json import pickle from collections.abc import Iterable from pathlib import Path from typing import Any VOCAB_PICKLE_FILENAME = "vocab.pkl" VOCAB_JSON_FILENAME = "vocab.json" # TextVectorization's padding and OOV tokens. ``get_vocabulary()`` returns them # at ids 0 and 1, so every saved vocabulary starts with them. _RESERVED_TOKENS = ("", "[UNK]") # The only globals a vocabulary pickle may name. The notebook pickles # ``get_vocabulary()``, whose entries can be ``numpy.str_``, and unpickling # those needs NumPy's dtype and scalar reconstructor. Neither calls arbitrary # code: they only build a dtype and a scalar from the pickled values. _ALLOWED_PICKLE_GLOBALS = frozenset({("numpy", "dtype"), ("numpy.core.multiarray", "scalar")}) class _VocabularyUnpickler(pickle.Unpickler): """Unpickler that refuses every global outside ``_ALLOWED_PICKLE_GLOBALS``.""" def find_class(self, module: str, name: str) -> Any: if (module, name) not in _ALLOWED_PICKLE_GLOBALS: raise pickle.UnpicklingError(f"refused global {module}.{name}") return super().find_class(module, name) def read_vocabulary(directory: str | Path, vocab_size: int) -> list[str]: """Read and validate a saved vocabulary list, without importing TensorFlow. ``vocab.json`` wins when both files exist. ``vocab.pkl`` is read only when there is no JSON file, through ``_VocabularyUnpickler``. Args: directory: Directory containing ``vocab.json`` or ``vocab.pkl``. vocab_size: Maximum vocabulary size the model was built with. Returns: The tokens in id order, as plain ``str``. Raises: FileNotFoundError: Neither file exists. ValueError: The file can't be parsed, or its list isn't a vocabulary this tokenizer could have saved. """ directory = Path(directory) js = directory / VOCAB_JSON_FILENAME pkl = directory / VOCAB_PICKLE_FILENAME if js.is_file(): source = js try: with js.open(encoding="utf-8") as f: vocab = json.load(f) except ValueError as exc: # JSONDecodeError, UnicodeDecodeError raise ValueError(f"Could not parse tokenizer vocabulary {js}: {exc}") from exc elif pkl.is_file(): source = pkl try: with pkl.open("rb") as f: vocab = _VocabularyUnpickler(f).load() except Exception as exc: # malformed pickles raise many exception types raise ValueError(f"Could not read tokenizer vocabulary {pkl}: {exc}") from exc else: raise FileNotFoundError( f"No tokenizer vocabulary found in {directory!s}. " f"Expected '{VOCAB_JSON_FILENAME}' (preferred) or " f"'{VOCAB_PICKLE_FILENAME}'. Train the model with " "`python -m scripts.train --config configs/base.yaml` to " "produce the artefacts, or point BACKEND_TOKENIZER_DIR at a " "directory that contains them." ) return _validate_vocabulary(vocab, source, vocab_size) def _validate_vocabulary(vocab: object, source: Path, vocab_size: int) -> list[str]: """Return ``vocab`` as plain strings, or raise ``ValueError`` naming ``source``.""" if not isinstance(vocab, list): raise ValueError( f"Tokenizer vocabulary {source} must be a list of strings, " f"got {type(vocab).__name__}." ) for index, token in enumerate(vocab): if not isinstance(token, str): raise ValueError( f"Tokenizer vocabulary {source} has a non-string entry at id {index}: " f"{type(token).__name__}." ) tokens = [str(token) for token in vocab] # numpy.str_ -> str if tuple(tokens[:2]) != _RESERVED_TOKENS: raise ValueError( f"Tokenizer vocabulary {source} must start with the padding and OOV tokens " f"{list(_RESERVED_TOKENS)!r}, got {tokens[:2]!r}." ) seen: set[str] = set() for token in tokens: if token in seen: raise ValueError(f"Tokenizer vocabulary {source} repeats the token {token!r}.") seen.add(token) if len(tokens) > vocab_size: raise ValueError( f"Tokenizer vocabulary {source} has {len(tokens)} entries, " f"more than vocab_size={vocab_size}." ) return tokens class CaptionTokenizer: """Wrapper that owns a fitted ``TextVectorization`` layer + lookup tables.""" def __init__(self, vocab_size: int, max_length: int) -> None: """Construct an unfit tokenizer. Args: vocab_size: Maximum vocabulary size (notebook: ``VOCABULARY_SIZE``). max_length: Pad/truncate every caption to this many tokens (notebook: ``MAX_LENGTH``). """ self.vocab_size = vocab_size self.max_length = max_length self._layer = None self._idx2word = None self._word2idx = None # ----------------------------------------------------------------- fit ---- def fit(self, captions: Iterable[str]) -> None: """Adapt the underlying TextVectorization layer to the given captions. Args: captions: An iterable of *already preprocessed* captions (i.e. lower-cased, punctuation-stripped, wrapped in ``[start] ... [end]``). Mirrors notebook cell 7 which calls ``tokenizer.adapt(captions['caption'])`` *after* cell 4 has applied ``preprocess`` to every row. """ import tensorflow as tf layer = tf.keras.layers.TextVectorization( max_tokens=self.vocab_size, standardize=None, output_sequence_length=self.max_length, ) layer.adapt(list(captions)) self._layer = layer self._build_lookups() # ----------------------------------------------------------- properties --- @property def vocabulary(self) -> list[str]: """Return the fitted vocabulary list (same order as TextVectorization).""" layer = self._require_fit() return list(layer.get_vocabulary()) @property def vocabulary_size(self) -> int: """Number of tokens in the fitted vocabulary.""" return int(self._require_fit().vocabulary_size()) @property def layer(self): """Direct access to the inner Keras layer. Exposed because the model's ``Embeddings`` layer (notebook cell 19) needs ``tokenizer.vocabulary_size()`` at construction time. Phase 1b replaces this with a constructor argument and removes the property. """ return self._require_fit() # -------------------------------------------------------- encode/decode --- def encode(self, text): """Encode ``text`` (str or list[str]) to integer-id tensor. Mirrors ``tokenizer(text)`` in notebook cells 7 and 25. Single string returns a 1-D tensor of shape ``[max_length]``; list returns 2-D. """ return self._require_fit()(text) def decode_id(self, idx) -> str: """Inverse-lookup a single integer id to its string token. Mirrors notebook cell 25's ``idx2word(pred_idx).numpy().decode('utf-8')``. """ self._require_fit() # By invariant, _idx2word is set together with _layer in fit/load. assert self._idx2word is not None word = self._idx2word(idx) return word.numpy().decode("utf-8") def word_to_id(self, word: str) -> int: """Look up a single word's integer id, returning 1 (the OOV id) if absent. Used by beam search to seed beams with the ``[start]`` token without going through ``TextVectorization``'s padded-string path. """ self._require_fit() assert self._word2idx is not None return int(self._word2idx(word).numpy()) # ---------------------------------------------------------- persistence --- def save(self, directory: str | Path) -> None: """Save the vocabulary to ``directory/vocab.pkl`` and ``vocab.json``. The pickle matches notebook cell 9 exactly so old artefacts remain loadable. The JSON sidecar is human-inspectable. """ self._require_fit() directory = Path(directory) directory.mkdir(parents=True, exist_ok=True) vocab = self.vocabulary with (directory / VOCAB_PICKLE_FILENAME).open("wb") as f: pickle.dump(vocab, f) with (directory / VOCAB_JSON_FILENAME).open("w", encoding="utf-8") as f: json.dump(vocab, f, ensure_ascii=False, indent=2) @classmethod def load( cls, directory: str | Path, vocab_size: int, max_length: int, ) -> CaptionTokenizer: """Load a previously saved vocabulary into a new tokenizer. Args: directory: Directory containing ``vocab.json`` (preferred) or ``vocab.pkl``. vocab_size: Maximum vocabulary size — must match the saved vocab. max_length: Pad/truncate length — must match training-time value. Returns: A fitted ``CaptionTokenizer`` ready to ``encode`` and ``decode_id``. Raises: FileNotFoundError: Neither vocabulary file exists. ValueError: The vocabulary is malformed (see ``read_vocabulary``). """ vocab = read_vocabulary(directory, vocab_size) import tensorflow as tf tok = cls(vocab_size=vocab_size, max_length=max_length) layer = tf.keras.layers.TextVectorization( max_tokens=vocab_size, standardize=None, output_sequence_length=max_length, ) layer.set_vocabulary(vocab) tok._layer = layer tok._build_lookups() return tok # -------------------------------------------------------------- internal -- def _build_lookups(self) -> None: """Construct ``StringLookup`` (idx → word) for inference decoding. Called only from ``fit()`` and ``load()``, *after* ``self._layer`` has been assigned, so the assertion below is a defensive no-op for mypy. """ import tensorflow as tf assert self._layer is not None vocab = self._layer.get_vocabulary() self._word2idx = tf.keras.layers.StringLookup(mask_token="", vocabulary=vocab) self._idx2word = tf.keras.layers.StringLookup(mask_token="", vocabulary=vocab, invert=True) def _require_fit(self): """Validate that the tokenizer has been fitted; return the inner layer. Returning the layer (rather than only raising on the unfit state) gives callers a non-``None``-typed local for the rest of their body — which is what mypy needs to prove ``layer.get_vocabulary()`` etc. are valid calls. Costs one attribute lookup at runtime. """ if self._layer is None: raise RuntimeError( "CaptionTokenizer not fitted. Call `.fit(captions)` or " "`.load(directory, ...)` first." ) return self._layer