.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/plot_create_model_random_weights.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_plot_create_model_random_weights.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_plot_create_model_random_weights.py:


Convert a random-weight model with ``create_model``
===================================================

This example shows the core workflow of mbext on a **tiny, randomly
initialised** model. Nothing is downloaded from the Hugging Face hub: a small
:class:`~transformers.LlamaConfig` is built with a single hidden layer and random
weights, converted to ONNX with :func:`modelbuilder.builder.create_model`, and
the ONNX logits are compared to the PyTorch reference.

This is exactly how the fast test suite validates every architecture, which is
why the tests run offline and in seconds (see the *Design* page).

.. GENERATED FROM PYTHON SOURCE LINES 19-25

.. code-block:: Python


    import os
    import tempfile

    import numpy as np








.. GENERATED FROM PYTHON SOURCE LINES 26-30

Build a tiny random-weight model
--------------------------------
A minimal ``LlamaForCausalLM`` with small dimensions and a single hidden layer.
The weights are randomly initialised (seeded for reproducibility).

.. GENERATED FROM PYTHON SOURCE LINES 30-54

.. code-block:: Python

    import torch
    from tokenizers import Tokenizer
    from tokenizers.models import WordLevel
    from transformers import AutoModelForCausalLM, LlamaConfig, PreTrainedTokenizerFast

    num_hidden_layers = 1
    config = LlamaConfig(
        architectures=["LlamaForCausalLM"],
        hidden_size=512,
        intermediate_size=1376,
        max_position_embeddings=2048,
        num_attention_heads=8,
        num_key_value_heads=4,
        num_hidden_layers=num_hidden_layers,
        vocab_size=1024,
        bos_token_id=1,
        eos_token_id=2,
    )
    head_size = config.hidden_size // config.num_attention_heads

    torch.manual_seed(42)
    model = AutoModelForCausalLM.from_config(config)
    model.eval()





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    LlamaForCausalLM(
      (model): LlamaModel(
        (embed_tokens): Embedding(1024, 512)
        (layers): ModuleList(
          (0): LlamaDecoderLayer(
            (self_attn): LlamaAttention(
              (q_proj): Linear(in_features=512, out_features=512, bias=False)
              (k_proj): Linear(in_features=512, out_features=256, bias=False)
              (v_proj): Linear(in_features=512, out_features=256, bias=False)
              (o_proj): Linear(in_features=512, out_features=512, bias=False)
            )
            (mlp): LlamaMLP(
              (gate_proj): Linear(in_features=512, out_features=1376, bias=False)
              (up_proj): Linear(in_features=512, out_features=1376, bias=False)
              (down_proj): Linear(in_features=1376, out_features=512, bias=False)
              (act_fn): SiLUActivation()
            )
            (input_layernorm): LlamaRMSNorm((512,), eps=1e-06)
            (post_attention_layernorm): LlamaRMSNorm((512,), eps=1e-06)
          )
        )
        (norm): LlamaRMSNorm((512,), eps=1e-06)
        (rotary_emb): LlamaRotaryEmbedding()
      )
      (lm_head): Linear(in_features=512, out_features=1024, bias=False)
    )



.. GENERATED FROM PYTHON SOURCE LINES 55-57

A minimal word-level tokenizer is enough for ``create_model`` to write the
processing files next to the ONNX model.

.. GENERATED FROM PYTHON SOURCE LINES 57-64

.. code-block:: Python

    tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=Tokenizer(WordLevel(vocab={"<unk>": 0, "<s>": 1, "</s>": 2}, unk_token="<unk>")),
        bos_token="<s>",
        eos_token="</s>",
        unk_token="<unk>",
    )








.. GENERATED FROM PYTHON SOURCE LINES 65-69

Convert to ONNX with ``create_model``
-------------------------------------
The PyTorch model and tokenizer are saved to a folder, then converted to ONNX
for the CPU execution provider in ``fp32`` precision.

.. GENERATED FROM PYTHON SOURCE LINES 69-92

.. code-block:: Python

    from modelbuilder.builder import create_model  # noqa: E402

    work_dir = tempfile.mkdtemp(prefix="mbext_example_")
    model_dir = os.path.join(work_dir, "model")
    output_dir = os.path.join(work_dir, "onnx")
    cache_dir = os.path.join(work_dir, "cache")

    model.save_pretrained(model_dir)
    tokenizer.save_pretrained(model_dir)

    create_model(
        model_name="tiny-llama",
        input_path=model_dir,
        output_dir=output_dir,
        precision="fp32",
        execution_provider="cpu",
        cache_dir=cache_dir,
        num_hidden_layers=num_hidden_layers,
    )

    onnx_path = os.path.join(output_dir, "model.onnx")
    print("ONNX model written to:", onnx_path)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]
    Writing model shards: 100%|██████████| 1/1 [00:00<00:00, 221.18it/s]
    GroupQueryAttention (GQA) is used in this model.

    Loading weights:   0%|          | 0/12 [00:00<?, ?it/s]
    Loading weights: 100%|██████████| 12/12 [00:00<00:00, 7076.01it/s]
    Reading embedding layer
    Reading layer 0
    Reading final norm
    Reading LM head
    Saving ONNX model in /tmp/mbext_example_vryhgkx5/onnx

    0it [00:00, ?it/s]
    Saving model.layers.0.input_layernorm.weight (f32, [512]):   8%|▊         | 1/12 [00:00<00:00, 3622.02it/s]
    Saving model.layers.0.post_attention_layernorm.weight (f32, [512]):  17%|█▋        | 2/12 [00:00<00:00, 3788.89it/s]
    Saving model.layers.1.final_norm_layernorm.weight (f32, [512]):  25%|██▌       | 3/12 [00:00<00:00, 4281.36it/s]    
    Saving cos_cache (f32, [2048,32]):  33%|███▎      | 4/12 [00:00<00:00, 4672.02it/s]                             
    Saving sin_cache (f32, [2048,32]):  42%|████▏     | 5/12 [00:00<00:00, 4497.43it/s]
    Saving model.layers.0.attn.o_proj.MatMul.weight (f32, [512,512]):  50%|█████     | 6/12 [00:00<00:00, 4547.49it/s]
    Saving model.embed_tokens.weight (f32, [1024,512]):  58%|█████▊    | 7/12 [00:00<00:00, 3407.63it/s]              
    Saving model.layers.0.attn.qkv_proj.MatMul.weight (f32, [512,1024]):  67%|██████▋   | 8/12 [00:00<00:00, 2215.55it/s]
    Saving lm_head.MatMul.weight (f32, [512,1024]):  75%|███████▌  | 9/12 [00:00<00:00, 1602.17it/s]                     
    Saving model.layers.0.mlp.gate_proj.MatMul.weight (f32, [512,1376]):  83%|████████▎ | 10/12 [00:00<00:00, 1436.11it/s]
    Saving model.layers.0.mlp.up_proj.MatMul.weight (f32, [512,1376]):  92%|█████████▏| 11/12 [00:00<00:00, 1221.34it/s]  
    Saving model.layers.0.mlp.down_proj.MatMul.weight (f32, [1376,512]): 100%|██████████| 12/12 [00:00<00:00, 1143.54it/s]
    Saving model.layers.0.mlp.down_proj.MatMul.weight (f32, [1376,512]): 100%|██████████| 12/12 [00:00<00:00, 894.12it/s] 
    Saving GenAI config in /tmp/mbext_example_vryhgkx5/onnx
    Saving processing files in /tmp/mbext_example_vryhgkx5/onnx for GenAI
    Saving VS Code settings in /tmp/mbext_example_vryhgkx5/onnx/.vscode/settings.json
    ONNX model written to: /tmp/mbext_example_vryhgkx5/onnx/model.onnx




.. GENERATED FROM PYTHON SOURCE LINES 93-98

Compare the ONNX and PyTorch logits
-----------------------------------
The exported graph expects ``input_ids``, ``attention_mask``, ``position_ids``
and an (empty) KV cache. We run a single prefill step through OnnxRuntime and
compare the logits with the PyTorch model.

.. GENERATED FROM PYTHON SOURCE LINES 98-126

.. code-block:: Python

    import onnxruntime  # noqa: E402

    sess = onnxruntime.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    onnx_input_names = {i.name for i in sess.get_inputs()}

    batch_size, seq_len = 1, 5
    torch.manual_seed(0)
    input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))

    feed = {
        "input_ids": input_ids.numpy().astype(np.int64),
        "attention_mask": np.ones((batch_size, seq_len), dtype=np.int64),
        "position_ids": np.arange(seq_len, dtype=np.int64).reshape(batch_size, seq_len),
    }
    for i in range(num_hidden_layers):
        feed[f"past_key_values.{i}.key"] = np.zeros((batch_size, config.num_key_value_heads, 0, head_size), dtype=np.float32)
        feed[f"past_key_values.{i}.value"] = np.zeros((batch_size, config.num_key_value_heads, 0, head_size), dtype=np.float32)
    feed = {k: v for k, v in feed.items() if k in onnx_input_names}

    onnx_logits = sess.run(None, feed)[0]

    with torch.no_grad():
        torch_logits = model(input_ids).logits.numpy()

    max_abs_diff = np.abs(onnx_logits - torch_logits).max()
    print("logits shape:", onnx_logits.shape)
    print("max absolute discrepancy (ONNX vs PyTorch):", float(max_abs_diff))





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    logits shape: (1, 5, 1024)
    max absolute discrepancy (ONNX vs PyTorch): 1.3113021850585938e-06




.. GENERATED FROM PYTHON SOURCE LINES 127-130

The discrepancy is tiny, confirming the ONNX model matches the PyTorch
reference. This is the check that runs for every supported architecture in the
fast test suite.


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 5.551 seconds)


.. _sphx_glr_download_auto_examples_plot_create_model_random_weights.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: plot_create_model_random_weights.ipynb <plot_create_model_random_weights.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: plot_create_model_random_weights.py <plot_create_model_random_weights.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: plot_create_model_random_weights.zip <plot_create_model_random_weights.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_