model/tokenizer_utils.py

"""
Standardized tokenizer loading for ArcisVLM.

ALL scripts should use load_tokenizer() from this module instead of
manually loading/training tokenizers. This prevents the critical bug
where eval scripts train a dummy tokenizer that doesn't match training.
"""

import os
import sys

from model.tokenizer import BPETokenizer


# Priority order for tokenizer discovery — 32K tokenizer ONLY
# The old 8K tokenizer.json is INCOMPATIBLE with v4 model and must not be used
TOKENIZER_SEARCH_PATHS = [
    "checkpoints/v4_tokenizer_32k.json",
    "checkpoints/tokenizer_32k.json",
]


def load_tokenizer(config: dict = None, checkpoint_dir: str = None) -> BPETokenizer:
    """Load the correct tokenizer, matching the model config.

    Args:
        config: Model config dict (used to get vocab_size)
        checkpoint_dir: Directory containing tokenizer files

    Returns:
        BPETokenizer with correct vocab

    Raises:
        FileNotFoundError: If no tokenizer file found (NEVER falls back to dummy)
    """
    vocab_size = 32768
    if config:
        vocab_size = config.get("tokenizer", {}).get("vocab_size",
                     config.get("decoder", {}).get("vocab_size", 32768))

    tokenizer = BPETokenizer(vocab_size=vocab_size)

    # Build search paths — v4_tokenizer_32k.json has highest priority
    # NO reference to old tokenizer.json (8K) — it's incompatible
    search_paths = []
    if checkpoint_dir:
        search_paths.append(os.path.join(checkpoint_dir, "v4_tokenizer_32k.json"))
        search_paths.append(os.path.join(checkpoint_dir, "tokenizer_32k.json"))
    search_paths.extend(TOKENIZER_SEARCH_PATHS)

    # Try each path
    for path in search_paths:
        if os.path.exists(path):
            tokenizer.load(path)
            actual_vocab = tokenizer.vocab_size if hasattr(tokenizer, 'vocab_size') else len(getattr(tokenizer, 'vocab', {}))
            print(f"  Tokenizer loaded: {path} ({actual_vocab} tokens)")

            # Warn on mismatch
            if actual_vocab > 0 and actual_vocab != vocab_size:
                print(f"  [WARN] Tokenizer vocab ({actual_vocab}) != config vocab ({vocab_size})")

            return tokenizer

    # NEVER train a dummy tokenizer — that causes 0% benchmarks
    available = [p for p in search_paths if os.path.exists(os.path.dirname(p) or ".")]
    raise FileNotFoundError(
        f"No tokenizer found! Searched: {search_paths}\n"
        f"Download from HuggingFace:\n"
        f"  python3 -c \"from huggingface_hub import hf_hub_download; "
        f"hf_hub_download('hardiksa/arcisvlm', 'checkpoints/v4_tokenizer_32k.json', local_dir='.')\""
    )


def validate_tokenizer_model_match(tokenizer: BPETokenizer, model) -> bool:
    """Check that tokenizer vocab matches model decoder vocab.

    Returns True if match, prints warning and returns False if mismatch.
    """
    if model is None:
        return True

    decoder = getattr(model, 'decoder', None)
    if decoder is None:
        return True

    # Get decoder vocab size from embedding layer
    tok_embed = getattr(decoder, 'token_embedding', None)
    if tok_embed is not None:
        model_vocab = tok_embed.num_embeddings
        tok_vocab = tokenizer.vocab_size if hasattr(tokenizer, 'vocab_size') else len(getattr(tokenizer, 'vocab', {}))

        if tok_vocab != model_vocab:
            print(f"  [ERROR] Tokenizer vocab ({tok_vocab}) != model decoder vocab ({model_vocab})")
            print(f"  This WILL cause garbage output. Fix tokenizer before proceeding.")
            return False

    return True