char_encoder.py

"""
Shared char-level tokenizer + encoder for the prefix -> official-name model.

Why char-level: user inputs are joined-without-spaces ("choachukang"),
initialisms ("cck") and short forms ("ns4"). WordPiece tokenizers shatter
these into meaningless pieces. A character model sees the raw letters, which
is exactly what these patterns are about.

This module is imported by train_prefix_encoder.py, test_prefix_encoder.py and
generate_embeddings_bin.py so that all three use IDENTICAL tokenization and
architecture. The vocab is derived deterministically from a constant alphabet,
so no external files are needed; it is also exported to char_vocab.json for the
mobile app.
"""

import json
import os

import torch
import torch.nn as nn

# =========================
# Shared config
# =========================
MAX_LEN = 40          # characters (official names + variants fit comfortably)
EMBEDDING_DIM = 128   # output embedding dim (kept == old model for mobile compat)
CHAR_DIM = 48         # character embedding size
HIDDEN = 96           # GRU hidden size (per direction)

PAD_TOKEN = "<pad>"
UNK_TOKEN = "<unk>"

# Lowercased character set the model understands. Anything else -> <unk>.
# Covers letters, digits, space and the punctuation that survives in official
# names (apostrophe, hyphen, period, ampersand, slash, comma).
ALPHABET = "abcdefghijklmnopqrstuvwxyz0123456789 '-.&/,"

VOCAB_FILE = os.path.join("artifacts_mobile", "char_vocab.json")


class CharTokenizer:
    """Deterministic character tokenizer (lowercases input, maps OOV to <unk>)."""

    def __init__(self, max_len: int = MAX_LEN):
        self.max_len = max_len
        # index 0 = pad, 1 = unk, then the alphabet
        self.itos = [PAD_TOKEN, UNK_TOKEN] + list(ALPHABET)
        self.stoi = {ch: i for i, ch in enumerate(self.itos)}
        self.pad_id = self.stoi[PAD_TOKEN]
        self.unk_id = self.stoi[UNK_TOKEN]

    @property
    def vocab_size(self) -> int:
        return len(self.itos)

    def _ids(self, text: str):
        text = (text or "").lower()
        ids = [self.stoi.get(ch, self.unk_id) for ch in text][: self.max_len]
        mask = [1] * len(ids)
        pad = self.max_len - len(ids)
        if pad > 0:
            ids += [self.pad_id] * pad
            mask += [0] * pad
        return ids, mask

    def encode_one(self, text: str):
        ids, mask = self._ids(text)
        return {
            "input_ids": torch.tensor(ids, dtype=torch.long),
            "attention_mask": torch.tensor(mask, dtype=torch.long),
        }

    def encode_batch(self, texts):
        ids, masks = [], []
        for t in texts:
            i, m = self._ids(t)
            ids.append(i)
            masks.append(m)
        return {
            "input_ids": torch.tensor(ids, dtype=torch.long),
            "attention_mask": torch.tensor(masks, dtype=torch.long),
        }

    def save(self, path: str = VOCAB_FILE):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(
                {"itos": self.itos, "max_len": self.max_len},
                f,
                ensure_ascii=False,
            )

    @classmethod
    def from_pretrained(cls, pretrained_dir_or_repo: str):
        """Load tokenizer from a local directory or Hugging Face Hub repo.

        Expects ``config.json`` for max_len, or falls back to MAX_LEN.
        Also reads ``char_vocab.json`` if present to restore the exact vocab
        (needed if the alphabet ever changes).
        """
        try:
            from huggingface_hub import snapshot_download
            local_dir = snapshot_download(
                pretrained_dir_or_repo,
                allow_patterns=["config.json", "char_vocab.json"],
            )
        except (ImportError, ValueError):
            local_dir = pretrained_dir_or_repo

        max_len = MAX_LEN
        config_path = os.path.join(local_dir, "config.json")
        if os.path.isfile(config_path):
            with open(config_path, encoding="utf-8") as f:
                cfg = json.load(f)
            max_len = cfg.get("max_len", MAX_LEN)

        tok = cls(max_len=max_len)

        vocab_path = os.path.join(local_dir, "char_vocab.json")
        if os.path.isfile(vocab_path):
            with open(vocab_path, encoding="utf-8") as f:
                data = json.load(f)
            if isinstance(data, dict) and "itos" in data:
                tok.itos = data["itos"]
                tok.stoi = {ch: i for i, ch in enumerate(tok.itos)}
                tok.pad_id = tok.stoi[PAD_TOKEN]
                tok.unk_id = tok.stoi.get(UNK_TOKEN, 1)

        return tok


class CharEncoder(nn.Module):
    """Char embedding -> BiGRU -> masked mean-pool -> projection.

    Shared (Siamese) weights encode both the user input and the official name.
    Masked mean-pool over time keeps it TorchScript-trace friendly at fixed
    MAX_LEN (unlike packed variable-length sequences).
    """

    def __init__(self, vocab_size: int, embedding_dim: int = EMBEDDING_DIM):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, CHAR_DIM, padding_idx=0)
        self.gru = nn.GRU(
            CHAR_DIM,
            HIDDEN,
            num_layers=1,
            batch_first=True,
            bidirectional=True,
        )
        self.proj = nn.Linear(HIDDEN * 2, embedding_dim)

    def forward(self, input_ids, attention_mask):
        emb = self.embedding(input_ids)                 # [B, T, CHAR_DIM]
        out, _ = self.gru(emb)                           # [B, T, 2*HIDDEN]

        # Masked mean-pool over valid (non-pad) positions.
        mask = attention_mask.unsqueeze(-1).type_as(out)  # [B, T, 1]
        summed = (out * mask).sum(dim=1)                  # [B, 2*HIDDEN]
        counts = mask.sum(dim=1).clamp(min=1.0)           # [B, 1]
        pooled = summed / counts
        return self.proj(pooled)                          # [B, EMBEDDING_DIM]

    def save_pretrained(self, save_directory: str):
        """Save model weights + config to a directory."""
        os.makedirs(save_directory, exist_ok=True)
        torch.save(self.state_dict(), os.path.join(save_directory, "prefix_encoder.pt"))
        config = {
            "model_type": "char_level_siamese_encoder",
            "architectures": ["CharEncoder"],
            "vocab_size": self.embedding.num_embeddings,
            "max_len": MAX_LEN,
            "char_dim": CHAR_DIM,
            "hidden_dim": HIDDEN,
            "embedding_dim": self.proj.out_features,
            "num_gru_layers": 1,
            "bidirectional": True,
        }
        with open(os.path.join(save_directory, "config.json"), "w", encoding="utf-8") as f:
            json.dump(config, f, ensure_ascii=False, indent=2)

    @classmethod
    def from_pretrained(cls, pretrained_dir_or_repo: str, map_location="cpu"):
        """Load model from a local directory or Hugging Face Hub repo.

        Example::

            model = CharEncoder.from_pretrained("username/sg-transit-prefix-encoder")
            tokenizer = CharTokenizer.from_pretrained("username/sg-transit-prefix-encoder")
        """
        try:
            from huggingface_hub import snapshot_download
            local_dir = snapshot_download(
                pretrained_dir_or_repo,
                allow_patterns=["config.json", "prefix_encoder.pt", "char_vocab.json"],
            )
        except (ImportError, ValueError):
            local_dir = pretrained_dir_or_repo

        # Vocab size comes from char_vocab.json (the ground truth), not config.json.
        # config.json can drift if it was hand-edited.
        vocab_size = None
        embedding_dim = EMBEDDING_DIM
        vocab_path = os.path.join(local_dir, "char_vocab.json")
        if os.path.isfile(vocab_path):
            with open(vocab_path, encoding="utf-8") as f:
                data = json.load(f)
            if isinstance(data, dict) and "itos" in data:
                vocab_size = len(data["itos"])

        config_path = os.path.join(local_dir, "config.json")
        if os.path.isfile(config_path):
            with open(config_path, encoding="utf-8") as f:
                cfg = json.load(f)
            if vocab_size is None:
                vocab_size = cfg.get("vocab_size", 45)
            embedding_dim = cfg.get("embedding_dim", EMBEDDING_DIM)

        if vocab_size is None:
            raise FileNotFoundError(f"Cannot determine vocab_size from {local_dir}")

        model = cls(vocab_size=vocab_size, embedding_dim=embedding_dim)
        weights_path = os.path.join(local_dir, "prefix_encoder.pt")
        if not os.path.isfile(weights_path):
            raise FileNotFoundError(f"prefix_encoder.pt not found in {local_dir}")
        state = torch.load(weights_path, map_location=map_location, weights_only=True)
        model.load_state_dict(state)
        model.eval()
        return model