OBLITERATUS/tests/fixtures/tiny_offline_model.py

"""Build a deterministic, synthetic Hugging Face causal language model."""

from __future__ import annotations

import json
from pathlib import Path

import torch
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.pre_tokenizers import Whitespace
from transformers import (
    GPT2Config,
    GPT2LMHeadModel,
    MixtralConfig,
    MixtralForCausalLM,
    PreTrainedTokenizerFast,
)


FIXTURE_SEED = 20260814
FIXTURE_VOCAB = {
    "<pad>": 0,
    "<eos>": 1,
    "<unk>": 2,
    "harmful": 3,
    "harmless": 4,
    "request": 5,
    "answer": 6,
    "hello": 7,
    "world": 8,
    "safe": 9,
    "test": 10,
}


def _fixture_tokenizer() -> PreTrainedTokenizerFast:
    tokenizer_backend = Tokenizer(WordLevel(FIXTURE_VOCAB, unk_token="<unk>"))
    tokenizer_backend.pre_tokenizer = Whitespace()
    return PreTrainedTokenizerFast(
        tokenizer_object=tokenizer_backend,
        pad_token="<pad>",
        eos_token="<eos>",
        unk_token="<unk>",
    )


def build_tiny_offline_model(destination: Path) -> Path:
    """Create a tiny random-init GPT-2 model without downloads or caches."""
    destination = Path(destination)
    destination.mkdir(parents=True, exist_ok=False)

    torch.manual_seed(FIXTURE_SEED)
    tokenizer = _fixture_tokenizer()
    config = GPT2Config(
        vocab_size=len(FIXTURE_VOCAB),
        n_positions=128,
        n_ctx=128,
        n_embd=16,
        n_layer=1,
        n_head=2,
        n_inner=32,
        bos_token_id=1,
        eos_token_id=1,
        pad_token_id=0,
    )
    model = GPT2LMHeadModel(config)
    model.save_pretrained(destination, safe_serialization=True)
    tokenizer.save_pretrained(destination)

    manifest = {
        "fixture": "tiny-offline-gpt2",
        "provenance": "generated locally from configuration with random initialization",
        "training_data": None,
        "third_party_weights": None,
        "license": "AGPL-3.0-only (part of the OBLITERATUS test suite)",
        "seed": FIXTURE_SEED,
        "architecture": {
            "model_type": "gpt2",
            "layers": 1,
            "hidden_size": 16,
            "attention_heads": 2,
            "vocabulary_size": len(FIXTURE_VOCAB),
        },
    }
    (destination / "fixture-provenance.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return destination


def build_tiny_offline_moe_model(destination: Path) -> Path:
    """Create a tiny random-init Mixtral model with fused 3D expert tensors.

    transformers 5 stores routed experts as ``experts.gate_up_proj`` and
    ``experts.down_proj`` parameters of shape ``(num_experts, ...)``, which is
    the layout frontier MoE checkpoints expose to weight surgery.
    """
    destination = Path(destination)
    destination.mkdir(parents=True, exist_ok=False)

    torch.manual_seed(FIXTURE_SEED)
    tokenizer = _fixture_tokenizer()
    config = MixtralConfig(
        vocab_size=len(FIXTURE_VOCAB),
        hidden_size=16,
        intermediate_size=32,
        num_hidden_layers=2,
        num_attention_heads=2,
        num_key_value_heads=2,
        num_local_experts=4,
        num_experts_per_tok=2,
        max_position_embeddings=2048,
        bos_token_id=1,
        eos_token_id=1,
        pad_token_id=0,
    )
    model = MixtralForCausalLM(config)
    model.save_pretrained(destination, safe_serialization=True)
    tokenizer.save_pretrained(destination)

    manifest = {
        "fixture": "tiny-offline-mixtral-moe",
        "provenance": "generated locally from configuration with random initialization",
        "training_data": None,
        "third_party_weights": None,
        "license": "AGPL-3.0-only (part of the OBLITERATUS test suite)",
        "seed": FIXTURE_SEED,
        "architecture": {
            "model_type": "mixtral",
            "layers": 2,
            "hidden_size": 16,
            "attention_heads": 2,
            "experts": 4,
            "experts_per_token": 2,
            "vocabulary_size": len(FIXTURE_VOCAB),
        },
    }
    (destination / "fixture-provenance.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return destination