""" 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 = "" UNK_TOKEN = "" # Lowercased character set the model understands. Anything else -> . # 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 ).""" 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