Upload char_encoder.py with huggingface_hub

0beee7d verified 4 months ago

8.92 kB

"""
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

	"""
	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