padoc/preprocess.py

"""Create a PaDoc-ready checkpoint from a standard image-text model."""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path

import torch
from transformers import AutoModelForImageTextToText, AutoProcessor

from .constants import (
    DEFAULT_FORK_TOKEN_MAP,
    DEFAULT_SPECIAL_TOKENS,
    PADOC_CONFIG_KEY,
    PADOC_FORK_MAP_KEY,
    PADOC_SPECIAL_TOKENS_KEY,
)

logger = logging.getLogger(__name__)


def get_padoc_metadata(model_or_config) -> dict:
    config = getattr(model_or_config, "config", model_or_config)
    metadata = getattr(config, PADOC_CONFIG_KEY, None)
    if metadata is None and hasattr(config, "text_config"):
        metadata = getattr(config.text_config, PADOC_CONFIG_KEY, None)
    if not isinstance(metadata, dict) or not metadata.get(PADOC_FORK_MAP_KEY):
        raise ValueError("Checkpoint has no padoc.fork_token_map metadata.")
    return metadata


def get_fork_token_map(model_or_config) -> dict[str, str]:
    return dict(get_padoc_metadata(model_or_config)[PADOC_FORK_MAP_KEY])


def _initialize_new_rows(model, token_ids: list[int], old_vocab_size: int, seed: int) -> None:
    if not token_ids:
        return
    with torch.no_grad(), torch.random.fork_rng():
        torch.manual_seed(seed)
        input_weights = model.get_input_embeddings().weight
        old_input = input_weights[:old_vocab_size].float()
        input_mean = old_input.mean(0)
        input_std = old_input.std(0)
        for token_id in token_ids:
            row = input_mean + torch.randn_like(input_mean) * input_std
            input_weights[token_id].copy_(row.to(input_weights.dtype))

        output = model.get_output_embeddings()
        if output is not None and output.weight is not input_weights:
            output_weights = output.weight
            old_output = output_weights[:old_vocab_size].float()
            output_mean = old_output.mean(0)
            output_std = old_output.std(0)
            for token_id in token_ids:
                row = output_mean + torch.randn_like(output_mean) * output_std
                output_weights[token_id].copy_(row.to(output_weights.dtype))


def preprocess_model(
    base_model: str | Path,
    output_dir: str | Path,
    *,
    special_tokens: list[str] | None = None,
    fork_token_map: dict[str, str] | None = None,
    dtype: torch.dtype = torch.bfloat16,
    seed: int = 42,
) -> Path:
    """Register atomic fork tokens and persist their mapping in config.json."""

    special_tokens = list(special_tokens or DEFAULT_SPECIAL_TOKENS)
    fork_token_map = dict(fork_token_map or DEFAULT_FORK_TOKEN_MAP)
    referenced = set(fork_token_map) | set(fork_token_map.values())
    if not referenced <= set(special_tokens):
        missing = sorted(referenced - set(special_tokens))
        raise ValueError(f"Fork map references tokens absent from special_tokens: {missing}")

    model = AutoModelForImageTextToText.from_pretrained(str(base_model), dtype=dtype)
    processor = AutoProcessor.from_pretrained(str(base_model))
    tokenizer = processor.tokenizer
    old_vocab_size = len(tokenizer)

    new_tokens = [
        token
        for token in special_tokens
        if len(tokenizer.encode(token, add_special_tokens=False)) != 1
    ]
    if new_tokens:
        tokenizer.add_special_tokens({"additional_special_tokens": new_tokens})
        model.resize_token_embeddings(len(tokenizer))
        new_ids = [tokenizer.encode(token, add_special_tokens=False)[0] for token in new_tokens]
        _initialize_new_rows(model, new_ids, old_vocab_size, seed)

    for token in special_tokens:
        ids = tokenizer.encode(token, add_special_tokens=False)
        if len(ids) != 1:
            raise ValueError(f"Special token {token!r} is not atomic: {ids}")

    metadata = {
        PADOC_SPECIAL_TOKENS_KEY: special_tokens,
        PADOC_FORK_MAP_KEY: fork_token_map,
    }
    setattr(model.config, PADOC_CONFIG_KEY, metadata)
    if hasattr(model.config, "text_config"):
        setattr(model.config.text_config, PADOC_CONFIG_KEY, metadata)

    output_path = Path(output_dir).expanduser().resolve()
    output_path.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(output_path)
    processor.save_pretrained(output_path)

    config_path = output_path / "config.json"
    with config_path.open(encoding="utf-8") as handle:
        config = json.load(handle)
    config[PADOC_CONFIG_KEY] = metadata
    with config_path.open("w", encoding="utf-8") as handle:
        json.dump(config, handle, indent=2, ensure_ascii=False)
        handle.write("\n")
    logger.info("Saved PaDoc-ready checkpoint to %s", output_path)
    return output_path


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="Create a PaDoc-ready checkpoint.")
    parser.add_argument("--base-model", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--dtype", choices=("bfloat16", "float32"), default="bfloat16")
    args = parser.parse_args(argv)
    logging.basicConfig(level=logging.INFO)
    preprocess_model(
        args.base_model,
        args.output,
        seed=args.seed,
        dtype=torch.bfloat16 if args.dtype == "bfloat16" else torch.float32,
    )


if __name__ == "__main__":
    main()