optimum-intel-internal-testing/tiny-random-qwen3-forced-aligner

"""Create a tiny random `qwen3_asr` (Qwen3-ForcedAligner) fixture.

Produces a ``Qwen3ASRForConditionalGeneration`` whose top-level ``model_type`` is
``qwen3_asr`` and whose ``thinker_config.model_type`` is ``qwen3_forced_aligner``
(so the model builds the forced-aligner ``classify_num`` head instead of the
vocabulary head), wrapping a ``qwen3_asr_audio_encoder`` audio tower and a
``qwen3_asr_text`` decoder.

The tiny artifact preserves the real architecture. Weights are random; the
original 0.6B checkpoint is never downloaded.

Requirements: the third-party ``qwen_asr`` package (registers the trust-remote-code
architecture with ``AutoConfig``/``AutoModel``), transformers, torch, soundfile.

Usage:
    python qwen3_forced_aligner_tiny_model_generation.py
    # writes ./tiny-random-qwen3-forced-aligner next to this script.
"""

import json
import os
import sys
from pathlib import Path

import numpy as np
import torch

# qwen_asr must be imported to register the model with AutoConfig/AutoModel.
import qwen_asr  # noqa: F401
from transformers import AutoConfig, AutoProcessor
from qwen_asr.core.transformers_backend.modeling_qwen3_asr import Qwen3ASRForConditionalGeneration


SEED = 42
MODEL_ID = "Qwen/Qwen3-ForcedAligner-0.6B"
OUTPUT_DIR = Path(__file__).resolve().parent / "tiny-random-qwen3-forced-aligner"

# The forced-aligner path is selected by the instance model_type of the thinker
# sub-config. ``PretrainedConfig.to_dict`` serializes the class attribute
# ("qwen3_asr_thinker"), so the saved config.json must be patched back to this
# value for the reloaded model to keep the classify head.
FORCED_ALIGNER_MODEL_TYPE = "qwen3_forced_aligner"

MAX_PARAMS = 25_000_000
MAX_BYTES = 100 * 1024 * 1024


def register_default_rope_shim():
    """Compatibility shim for transformers >= 5.

    The ``qwen_asr`` remote code targets transformers 4.57, where the
    ``"default"`` RoPE type is an entry in ``ROPE_INIT_FUNCTIONS``. Transformers 5
    dropped that dict entry (turning it into a rotary-embedding method), so the
    remote code's ``ROPE_INIT_FUNCTIONS[self.rope_type]`` lookup and the tf5
    weight-init path both fail for ``rope_type == "default"``. Re-register an
    equivalent classic-RoPE initializer and attach it to the remote rotary
    modules under the name tf5 expects.
    """
    from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
    from qwen_asr.core.transformers_backend import modeling_qwen3_asr as mod

    def _compute_default_rope_parameters(config, device=None, seq_len=None, **kwargs):
        rope_params = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {}
        base = rope_params.get("rope_theta") if isinstance(rope_params, dict) else None
        if base is None:
            base = getattr(config, "rope_theta", 10000.0)
        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
        inv_freq = 1.0 / (
            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
        )
        return inv_freq, 1.0

    ROPE_INIT_FUNCTIONS.setdefault("default", _compute_default_rope_parameters)

    # tf5 `_init_weights` calls `module.compute_default_rope_parameters` on every
    # rotary module whose `rope_type == "default"`.
    for name in dir(mod):
        obj = getattr(mod, name)
        if isinstance(obj, type) and name.endswith("RotaryEmbedding"):
            if not hasattr(obj, "compute_default_rope_parameters"):
                obj.compute_default_rope_parameters = staticmethod(_compute_default_rope_parameters)


def build_config():
    config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)

    thinker = config.thinker_config
    # Preserve the forced-aligner identity + classify head width.
    assert "forced_aligner" in thinker.model_type, thinker.model_type
    text = thinker.text_config
    audio = thinker.audio_config

    # transformers 5 dropped the generic special-token attributes from the base
    # PretrainedConfig; the remote code still reads config.pad_token_id.
    if not hasattr(thinker, "pad_token_id"):
        thinker.pad_token_id = None

    # --- Text decoder (qwen3_asr_text): shrink transformer, keep tokenizer vocab. ---
    text.hidden_size = 64
    text.head_dim = 16
    text.num_attention_heads = 4
    text.num_key_value_heads = 2
    text.intermediate_size = 128
    text.num_hidden_layers = 2
    text.max_position_embeddings = 2048

    # Keep the interleaved M-RoPE path; rescale mrope_section to sum to head_dim // 2.
    # Original invariant: sum([24, 20, 20]) == 64 == original head_dim(128) // 2.
    rope_scaling = dict(text.rope_scaling)
    assert rope_scaling["mrope_interleaved"] is True, rope_scaling
    assert sum(rope_scaling["mrope_section"]) == 128 // 2, rope_scaling["mrope_section"]
    rope_scaling["mrope_section"] = [4, 2, 2]  # sum == 8 == head_dim(16) // 2
    text.rope_scaling = rope_scaling

    # --- Audio encoder (qwen3_asr_audio_encoder): shrink conformer, keep mel bins. ---
    audio.d_model = 64
    audio.encoder_layers = 2
    audio.num_hidden_layers = 2
    audio.encoder_attention_heads = 4  # head_dim = d_model / heads = 16
    audio.encoder_ffn_dim = 128
    audio.downsample_hidden_size = 32
    # output_dim must match the text hidden size so audio embeddings merge into
    # the decoder embedding space.
    audio.output_dim = 64
    # num_mel_bins stays 128: the processor always emits 128 mel-bin features.

    # Force float32 at every effective precision field.
    for sub in (config, thinker, text, audio):
        sub.dtype = "float32"
        sub.torch_dtype = "float32"

    return config


def patch_saved_thinker_model_type(output_dir: Path) -> None:
    """Restore the forced-aligner thinker model_type in the saved config.json."""
    config_path = output_dir / "config.json"
    with open(config_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    data.setdefault("thinker_config", {})["model_type"] = FORCED_ALIGNER_MODEL_TYPE
    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)
        f.write("\n")


def normalize_tied_weight_keys(model) -> None:
    """Convert legacy list-format ``_tied_weights_keys`` to the tf5 dict format.

    The ``qwen_asr`` remote code declares ``_tied_weights_keys`` as a list (the
    transformers 4.x format), but tf5's ``save_pretrained`` calls ``.keys()`` on
    it. Nothing is actually tied here (``tie_word_embeddings=False`` and the
    forced-aligner ``lm_head`` differs in shape from the embeddings), so an empty
    mapping is correct.
    """
    for module in model.modules():
        keys = getattr(module, "_tied_weights_keys", None)
        if isinstance(keys, list):
            module._tied_weights_keys = {}


def generate_audio_data(seconds: float = 2.0, sample_rate: int = 16000):
    np.random.seed(SEED)
    t = np.linspace(0, 1.0, int(sample_rate * seconds), endpoint=False)
    audio = (0.5 * np.sin(2 * np.pi * 440 * t)).astype(np.float32)
    return audio, sample_rate


def run_inference(output_dir: Path) -> None:
    """Reload the exact saved artifact and run a forward + short generate."""
    reloaded_config = AutoConfig.from_pretrained(output_dir, trust_remote_code=True)
    assert reloaded_config.model_type == "qwen3_asr", reloaded_config.model_type
    assert "forced_aligner" in reloaded_config.thinker_config.model_type, reloaded_config.thinker_config.model_type
    assert reloaded_config.thinker_config.audio_config.model_type == "qwen3_asr_audio_encoder"
    assert reloaded_config.thinker_config.text_config.model_type == "qwen3_asr_text"

    model = Qwen3ASRForConditionalGeneration.from_pretrained(output_dir, trust_remote_code=True).float().eval()

    # The classify head width must equal classify_num for a forced aligner.
    classify_num = reloaded_config.thinker_config.classify_num
    assert model.thinker.lm_head.out_features == classify_num, (
        model.thinker.lm_head.out_features,
        classify_num,
    )

    processor = AutoProcessor.from_pretrained(output_dir, trust_remote_code=True)
    audio_data, sample_rate = generate_audio_data()
    text_prompt = processor.apply_chat_template(
        [
            {"role": "system", "content": ""},
            {"role": "user", "content": [{"type": "audio", "audio": ""}]},
        ],
        add_generation_prompt=True,
        tokenize=False,
    )
    inputs = processor(text=text_prompt, audio=audio_data, sampling_rate=sample_rate, return_tensors="pt")

    with torch.no_grad():
        outputs = model.thinker(
            input_ids=inputs["input_ids"],
            input_features=inputs["input_features"],
            feature_attention_mask=inputs["feature_attention_mask"],
            attention_mask=inputs["attention_mask"],
        )
    logits = outputs.logits
    print(f"  forward ok, logits shape: {tuple(logits.shape)} (last dim == classify_num={classify_num})")
    assert logits.shape[-1] == classify_num

    with torch.no_grad():
        generated = model.generate(
            input_ids=inputs["input_ids"],
            input_features=inputs["input_features"],
            feature_attention_mask=inputs["feature_attention_mask"],
            attention_mask=inputs["attention_mask"],
            max_new_tokens=5,
            do_sample=False,
        )
    if hasattr(generated, "sequences"):
        generated = generated.sequences
    print(f"  generate ok, output shape: {tuple(generated.shape)}")


def main():
    torch.manual_seed(SEED)

    register_default_rope_shim()

    config = build_config()
    model = Qwen3ASRForConditionalGeneration(config).float().eval()

    num_params = sum(p.numel() for p in model.parameters())
    num_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    print(f"params={num_params:,}  bytes={num_bytes:,} ({num_bytes / 1024 / 1024:.2f} MB)")
    if num_params > MAX_PARAMS:
        sys.exit(f"param budget exceeded: {num_params:,} > {MAX_PARAMS:,}")
    if num_bytes > MAX_BYTES:
        sys.exit(f"memory budget exceeded: {num_bytes:,} > {MAX_BYTES:,}")

    normalize_tied_weight_keys(model)

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(OUTPUT_DIR, safe_serialization=True)

    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
    processor.save_pretrained(OUTPUT_DIR)

    patch_saved_thinker_model_type(OUTPUT_DIR)

    print(f"Saved tiny qwen3 forced aligner model to: {OUTPUT_DIR}")
    run_inference(OUTPUT_DIR)
    print("Inference smoke test passed.")


if __name__ == "__main__":
    main()