Update README.md

9035d80 verified about 2 months ago

10.7 kB

---
license: apache-2.0
---
```python
"""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()
```

	---
	license: apache-2.0
	---

	```python

	"""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()


	```