export_to_hf.py

"""Convert TaoTrain checkpoint + tokenizer assets into a standard HF package."""

import json
import shutil
from pathlib import Path

import torch


IGNORED_CHECKPOINT_SUFFIXES = (
    ".rotary.inv_freq",
)


def normalize_checkpoint(checkpoint):
    if isinstance(checkpoint, dict):
        if "model_state" in checkpoint:
            return checkpoint["model_state"], checkpoint.get("config", {})
        if "model_state_dict" in checkpoint:
            return checkpoint["model_state_dict"], checkpoint.get("config", {})
    return checkpoint, {}


def infer_special_token_paths(repo_dir):
    subdir_metadata = repo_dir / "tokenizer" / "tokenizer.special_tokens.json"
    root_metadata = repo_dir / "tokenizer.special_tokens.json"
    subdir_model = repo_dir / "tokenizer" / "tokenizer.model"
    subdir_vocab = repo_dir / "tokenizer" / "tokenizer.vocab"
    if subdir_metadata.exists() and subdir_model.exists() and subdir_vocab.exists():
        return subdir_model, subdir_metadata, subdir_vocab
    if root_metadata.exists():
        return repo_dir / "tokenizer.model", root_metadata, repo_dir / "tokenizer.vocab"
    return repo_dir / "tokenizer" / "tokenizer.model", subdir_metadata, repo_dir / "tokenizer" / "tokenizer.vocab"


def sanitize_model_state(model_state):
    """Drop deterministic non-persistent buffers that should not participate in HF weight export."""
    sanitized = {}
    ignored = []
    for key, value in model_state.items():
        if key.endswith(IGNORED_CHECKPOINT_SUFFIXES):
            ignored.append(key)
            continue
        sanitized[key] = value
    return sanitized, ignored


def write_clean_tokenizer_metadata(repo_dir, special_tokens):
    tokenizer_config = {
        "backend": "custom",
        "bos_token": "<BOS>",
        "eos_token": "<EOS>",
        "unk_token": "<UNK>",
        "pad_token": "<PAD>",
        "tokenizer_class": "TaoNetTokenizer",
        "model_max_length": 1000000000000000019884624838656,
        "extra_special_tokens": [
            token
            for token in special_tokens
            if token not in {"<UNK>", "<BOS>", "<EOS>", "<PAD>"}
        ],
    }
    special_tokens_map = {
        "unk_token": "<UNK>",
        "bos_token": "<BOS>",
        "eos_token": "<EOS>",
        "pad_token": "<PAD>",
        "additional_special_tokens": [
            token
            for token in special_tokens
            if token not in {"<UNK>", "<BOS>", "<EOS>", "<PAD>"}
        ],
    }

    with open(repo_dir / "tokenizer_config.json", "w", encoding="utf-8") as handle:
        json.dump(tokenizer_config, handle, indent=2)
        handle.write("\n")

    with open(repo_dir / "special_tokens_map.json", "w", encoding="utf-8") as handle:
        json.dump(special_tokens_map, handle, indent=2)
        handle.write("\n")

    added_tokens_path = repo_dir / "added_tokens.json"
    if added_tokens_path.exists():
        added_tokens_path.unlink()


def main():
    from configuration_taonet import TaoNetConfig
    from modeling_taonet import TaoNetForCausalLM
    from tokenization_taonet import TaoNetTokenizer

    repo_dir = Path(__file__).resolve().parent
    checkpoint_path = repo_dir / "checkpoints" / "sft" / "final_model.pt"

    checkpoint = torch.load(checkpoint_path, map_location="cpu")
    model_state, train_config = normalize_checkpoint(checkpoint)
    model_state, ignored_keys = sanitize_model_state(model_state)
    model_config = dict(train_config.get("model", {}))

    metadata_model_path, metadata_path, vocab_path = infer_special_token_paths(repo_dir)
    with open(metadata_path, "r", encoding="utf-8") as handle:
        metadata = json.load(handle)

    special_tokens = metadata.get("special_tokens", {})
    hf_config = TaoNetConfig.from_taotrain_model_config(
        model_config,
        vocab_size=model_config.get("vocab_size", sum(1 for _ in open(vocab_path, "r", encoding="utf-8"))),
        pad_token_id=special_tokens.get("<PAD>", 3),
        bos_token_id=special_tokens.get("<BOS>", 1),
        eos_token_id=special_tokens.get("<EOS>", 2),
        unk_token_id=special_tokens.get("<UNK>", 0),
    )
    hf_config.architectures = ["TaoNetForCausalLM"]
    hf_config.auto_map = {
        "AutoConfig": "configuration_taonet.TaoNetConfig",
        "AutoModelForCausalLM": "modeling_taonet.TaoNetForCausalLM",
    }

    model = TaoNetForCausalLM(hf_config)
    current_state = model.model.state_dict()
    missing = sorted(set(current_state) - set(model_state))
    unexpected = sorted(set(model_state) - set(current_state))
    if missing or unexpected:
        if missing:
            print("Missing keys while loading model state:")
            for key in missing:
                print(f"  - {key}")
        if unexpected:
            print("Unexpected keys while loading model state:")
            for key in unexpected:
                print(f"  - {key}")
        raise ValueError("Checkpoint/model key mismatch detected. Refusing to export a partial model.")

    model.model.load_state_dict(model_state, strict=True)
    model.tie_weights()

    exported_state = model.model.state_dict()
    mismatched_tensors = []
    for key, value in model_state.items():
        exported_value = exported_state[key]
        if value.shape != exported_value.shape:
            mismatched_tensors.append((key, "shape"))
            continue
        if value.dtype.is_floating_point:
            if not torch.equal(value, exported_value.to(dtype=value.dtype)):
                mismatched_tensors.append((key, "value"))
        else:
            if not torch.equal(value, exported_value):
                mismatched_tensors.append((key, "value"))

    if mismatched_tensors:
        print("Tensor mismatches detected after strict load:")
        for key, mismatch_type in mismatched_tensors[:20]:
            print(f"  - {key} ({mismatch_type})")
        raise ValueError("Checkpoint tensors do not match the HF wrapper after load.")

    model.save_pretrained(repo_dir, safe_serialization=False)

    tokenizer = TaoNetTokenizer(
        vocab_file=str(metadata_model_path),
        special_tokens_file=str(metadata_path),
    )
    tokenizer.save_pretrained(repo_dir)
    write_clean_tokenizer_metadata(repo_dir, metadata.get("configured_special_tokens", []))

    shutil.copyfile(metadata_model_path, repo_dir / "tokenizer.model")
    shutil.copyfile(metadata_path, repo_dir / "tokenizer.special_tokens.json")
    shutil.copyfile(vocab_path, repo_dir / "tokenizer.vocab")

    if ignored_keys:
        print("Ignored non-persistent checkpoint buffers:")
        for key in ignored_keys:
            print(f"  - {key}")
    print(f"Saved Hugging Face package to: {repo_dir}")


if __name__ == "__main__":
    main()