<?xml version="1.0" encoding="utf-8"?><feed xmlns="http://www.w3.org/2005/Atom" ><generator uri="https://jekyllrb.com/" version="3.10.0">Jekyll</generator><link href="https://karlosmid.com/feed.xml" rel="self" type="application/atom+xml" /><link href="https://karlosmid.com/" rel="alternate" type="text/html" /><updated>2026-06-17T10:53:59+00:00</updated><id>https://karlosmid.com/feed.xml</id><title type="html">Karlo Šmid Blog</title><subtitle>karlosmid.github.io</subtitle><entry><title type="html">Recap on fine-tuning to follow instructions</title><link href="https://karlosmid.com/2026/06/recap-on-fine-tuning-to-follow-instructions/" rel="alternate" type="text/html" title="Recap on fine-tuning to follow instructions" /><published>2026-06-17T09:00:00+00:00</published><updated>2026-06-17T09:00:00+00:00</updated><id>https://karlosmid.com/2026/06/recap-on-fine-tuning-to-follow-instructions</id><content type="html" xml:base="https://karlosmid.com/2026/06/recap-on-fine-tuning-to-follow-instructions/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>I am recapping the blog post <a href="/fine-tuning-to-follow-instructions/">Fine-tuning to follow instructions</a> by reviewing the chapter 7 workbook questions from <em>Build an LLM from Scratch</em> and reading Giles’ blog posts on chapter 7.</p>

<h2 id="giles-blog-posts">Giles blog posts</h2>

<p>Giles wrote three blog posts: <a href="https://www.gilesthomas.com/2025/10/llm-from-scratch-24-the-transcript-hack">part24</a>, <a href="https://www.gilesthomas.com/2025/10/llm-from-scratch-25-instruction-fine-tuning">part25</a>, and <a href="https://www.gilesthomas.com/2025/11/llm-from-scratch-26-evaluating-the-fine-tuned-model">part26</a>.</p>

<p>In the part24 post, another sidestep, Giles tries to feed instruction prompts into the GPT-2 models mentioned in the book. The result is that none of those models could give a meaningful answer, but the base Qwen 3 model with 600M parameters gave meaningful instruction results.</p>

<p>In part25, Giles always thinks more deeply about topics presented in books. For example, in relation to the Alpaca prompt format, it was created because interaction with the model was one-shot, due to the limited size of the model context. Today’s models have context sizes in the hundreds of thousands of tokens, so they could accept lengthy conversation transcripts. Then he explains that padding tokens waste GPU cycles, so we have batches of different sizes. The question is why we do not put -100 instead of padding tokens in inputs, but only in targets. Giles explains that this is not valid token vocabulary, but targets are also vocabulary tokens, so I do not understand that reasoning.</p>

<p>part26 is the shortest blog post! It is about evaluating our instruction fine-tuned model using an Ollama 3 model.</p>

<h2 id="workbook-questions">Workbook questions</h2>

<h3 id="quick-questions">Quick questions</h3>

<p>1 What is the primary challange that pretrained LLMs often face regarding instructions?</p>

<p>A Inability to complete sentences<br />
B Difficulty in generating coherent text<br />
C Limited vocabulary size.<br />
D Struggling with specific instructions like grammar correction or voice conversion.</p>

<p>My answer: D</p>

<p>Correct answer: +</p>

<p>2 What is the key component in preparing da dataset for supervised instruction fine-tuning?</p>

<p>A Pre-trained language model.<br />
B Optimization algorithm.<br />
C Instruction-response pairs.<br />
D Tokenization algorithm.</p>

<p>My answer: C</p>

<p>Correct answer: +</p>

<p>3 What data format is commonly used for instruction datasets due to its human and machine redability?</p>

<p>A JSON (JavaScript Object Notation)<br />
B CSV (comma separated values)<br />
C YAML (YAML Ain’t Markup Language)<br />
D XML (Extensible Markup Language).</p>

<p>My answer: A</p>

<p>Correct answer: +</p>

<p>4 What is the purpose of a custom collate function in the context of instruction fine-tuning?</p>

<p>A To optimize the model’s architecture.<br />
B To handle the specific formatting and requirements of the instruction fine-tuning dataset.<br />
C To pre-process the input data.<br />
D To evaluate the model’s performance</p>

<p>My answer: B</p>

<p>Correct answer: +</p>

<p>5 What is the purpose of using the ignore_index parameter (-100) in the custom collate function?</p>

<p>A To mark the end of a sequence.<br />
B To indicate an unknown token.<br />
C To exclude padding tokens from the loss calculation.<br />
D To identify the start of a sequence.</p>

<p>My answer: C</p>

<p>Correct answer: +</p>

<p>6 What is the purpose of saving the fine-tuned model’s state dictionary?</p>

<p>A To compress the model for efficient storage.<br />
B To visualize the model’s architecture.<br />
C To save the model’s parameters for later use or reuse in other projects.<br />
D To improve the model’s performance.</p>

<p>My answer: C</p>

<p>Correct answer: +</p>

<h3 id="71-introduction-to-instruction-fine-tuning">7.1 Introduction to instruction fine-tuning</h3>

<p>1 What is the primary function of a pretrained LLM?</p>

<p>My answer: The primary function of a pretrained LLM is to provide model knoweledge of language in which was trained grammar patterns.</p>

<p>Correct answer: <em>A pretrained LLM is primarily capable of text completion, meaning it can finish sentences or write text paragraphs given a fragment as input.</em></p>

<p>It seems that I still do not know to describe the primary task of pre-trained LLM.</p>

<p>2 What is the challange that pretrained LLMs often face?</p>

<p>My answer: The challange is that pretrained LLMs are not good at resolving instructions tasks.</p>

<p>Correct answer: <em>Pretrained LLMs often struggle with specific instructions, such as grammar correction or voice conversion, requiring further fine-tuning.</em></p>

<p>This one is correct.</p>

<p>3 What is the purpose of instruction fine-tuning?</p>

<p>My answer: The purpose of instruction fine-tuning is to train model for instructions tasks solving using as fine-tuning data instruction set with response how to resolve this task.</p>

<p>Correct answer: <em>Instruction fine-tuning aims to improve an LLM’s ability to follow specific instructions and generate desired responses based on those instructions.</em></p>

<p>My answer is correct.</p>

<p>4 What is a crucial aspect of instruction fine-tuning?</p>

<p>My answer: A crucail aspect of instruction fine-tuning is that we make LLM into chatbot.</p>

<p>Correct answer: <em>Preparing a suitable dataset is a key aspect of instruction fine-tuning, providing the model with examples of instructions and desired responses.</em></p>

<p>Wrong answer.</p>

<h3 id="72-preparing-a-dataset-for-supervised-instruction-fine-tuning">7.2 Preparing a dataset for supervised instruction fine-tuning</h3>

<p>1 What is the purpose of the instruction dataset used for fine-tuning a pretrained LLM?</p>

<p>My answer: Purpose of the instruction dataset is to provide training, validation and testing datasets so we could fine tune LLM to resolve instruction tasks.</p>

<p>Correct answer: <em>The instruction dataset consists of instruction-response pairs that are used to train the LLM to follow instructions and generate appropriate responses based on given inputs.</em></p>

<p>My answer is wrong.</p>

<p>2 Describe the format of the instruction dataset used in this section.</p>

<p>My answer: Each record has three keys, instructions contains instruction task, input is empty and output contains result of instruction task.</p>

<p>Correct answer: The dataset is stored in a JSON file and contains entries that are Python dictionary objects. Each entry includes an ‘instruction’, ‘input’, and ‘output’ field, representing the task, input data, and desired response, respectively._</p>

<p>I forgot to answer that format is JSON.</p>

<p>3 What are the two prompt styles mentioned in the section, and how they differ?</p>

<p>My answer: We have alpaca format and phi3 format. In phi3 we are using user and assistant tags, and in alpaca we use instructions, input and output tags.</p>

<p>Correct answer: <em>The two prompt styles are Alpaca and Phi-3. Alpaca uses a structured format with defined sections for instruction, input, and response, while Phi-3 employs a simpler format with designated &lt;|user|&gt; and &lt;|assistant|&gt; tokens</em></p>

<p>My answer is correct.</p>

<p>4 What is the purpose of <code class="language-plaintext highlighter-rouge">format_input</code> action and how does it work?</p>

<p>My answer: The purpose of <code class="language-plaintext highlighter-rouge">format_input</code> function is to format input dataset to alpaca or phi3 prompt format. It creates a chatbot prompt.</p>

<p>Correct answer: <em>The format_input function converts the entries in the instruction dataset into the Alpaca-style input format. It constructs a formatted string that includes the instruction, input (if available), and a placeholder for the response.</em></p>

<p>My answer is wrong.</p>

<p>5 How is the instruction dataset divided into training, validation, and test sets?</p>

<p>My answer: We randomly divide instruction dataset, 85% is training data, 10% is testing set, and rest is validation set.</p>

<p>Correct answer: <em>The dataset is divided into training, validation, and test sets using a specific ratio. The training set is used to train the model, the validation set is used to evaluate the model’s performance during training, and the test set is used to evaluate the model’s final performance.</em></p>

<p>My answer is correct.</p>

<p>6 Match the terms on the left to their description on the right:</p>

<p><img src="/assets/2026/06/instruction-prompt-styles-table.png" alt="Instruction prompt styles table" /></p>

<p>My answer: Instruction-response pairs =&gt; 3, Prompt styles =&gt; 1, Alpaca prompt style =&gt; 4, Phi3 Prompt style =&gt; 2</p>

<p>Correct answer: +</p>

<h3 id="73-organizing-data-into-training-batches">7.3 Organizing data into training batches</h3>

<p>1 What is the purpose of a custom collate function in instruction fine-tuning?</p>

<p>My answer: Purpose of a custom collate function is to prepare prompts in alpaca or phi3 formats from input data. It also pads input data to to longest batch value, and sets in target data padding tokens to -100 value, except for the one next to last token.</p>

<p>Correct answer: <em>A custom collate function is used to handle the specific requirements and formatting of the instruction fine-tuning dataset. It ensures that training examples are padded to the same length within each batch, allowing for efficient processing by the model.</em></p>

<p>My answer is correct.</p>

<p>2 How does the custom collate function handle padding in instruction fine-tuning?</p>

<p>My answer: It padds inputs to the value of longest input in the batch. Targets padding tokens are replaced with -100 instead of one next to last token.</p>

<table>
  <tbody>
    <tr>
      <td>Correct answer: _The custom collate function pads training examples to the length of the longest example in each batch, using the &lt;</td>
      <td>endoftext</td>
      <td>&gt; token ID (50256). This minimizes unnecessary padding by only extending sequences to match the longest one in each batch._</td>
    </tr>
  </tbody>
</table>

<p>I forgot to answer value of padding token.</p>

<p>3 Explain the role of target token IDs in instruction fine-tuning and how they are generated.</p>

<p>My answer: Target token IDs are pair with the input so we know which token should come next. It is generated from input so that input is shiffted by one token to the left.</p>

<p>Correct answer: <em>Target token IDs represent the desired output sequence that the model should generate. They are created by shifting the input token IDs one position to the right, omitting the first token and appending an end-of- text token. This setup allows the model to learn how to predict the next token in a sequence.</em></p>

<p>One token to the right and I did not answer that we are adding endoftext token when we reach end of input.</p>

<p>4 Why are padding tokens replaced with the -100 placeholder value in the target token IDs?</p>

<p>My answer: We are replacing padding tokens with -100 because -100 tokens are not considered in loss calculation.</p>

<p>Correct answer: <em>Replacing padding tokens with -100 allows the cross entropy loss function to ignore them during training. This ensures that only meaningful data influences model learning and prevents padding tokens from contributing to the loss calculation</em></p>

<p>I did not answer that padding tokens should not contribute to loss function calculation.</p>

<p>5 What is the purpose of retaining one end-of-text token in the target sequence?</p>

<p>My answer: Keeping that token we let know llm model to know where is the end of input instruction.</p>

<p>Correct answer: <em>Retaining one end-of-text token in the target sequence helps the LLM learn to generate end-of-text tokens, which act as an indicator that the generated response is complete.</em></p>

<p>Wrong answer, we are kepping endoftext token in target sequence, not input. That was the question.</p>

<p>6 Pieces of code have been removed from four places in this listing. Which of these terms have been removed and where should they go? Note that a term may appear more than once!</p>

<p>A true B false C format D encode_text! E full_text</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defp</span> <span class="n">encode_example</span><span class="p">(</span><span class="n">entry</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">,</span> <span class="n">prompt_style</span><span class="p">,</span> <span class="n">__1__</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">entry</span>
    <span class="c1"># The model trains on the instruction/input prompt followed by the</span>
    <span class="c1"># expected response target, matching the chapter 7 PyTorch dataset.</span>
    <span class="o">|&gt;</span> <span class="n">__3__</span><span class="p">(</span><span class="n">prompt_style</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">__4__</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">encode_example</span><span class="p">(</span><span class="n">entry</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">,</span> <span class="n">prompt_style</span><span class="p">,</span> <span class="n">__2__</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">token_ids</span> <span class="o">=</span>
      <span class="n">entry</span>
      <span class="o">|&gt;</span> <span class="n">__5__</span><span class="p">(</span><span class="n">prompt_style</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="n">__6__</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="n">prompt_length</span> <span class="o">=</span>
      <span class="n">entry</span>
      <span class="o">|&gt;</span> <span class="n">prompt_text</span><span class="p">(</span><span class="n">prompt_style</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="n">encode_text!</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="p">()</span>

    <span class="p">%{</span><span class="ss">token_ids:</span> <span class="n">token_ids</span><span class="p">,</span> <span class="ss">prompt_length:</span> <span class="n">prompt_length</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">encode_example</span><span class="p">(</span>
         <span class="n">_entry</span><span class="p">,</span>
         <span class="n">_tokenizer</span><span class="p">,</span>
         <span class="n">_opts</span><span class="p">,</span>
         <span class="n">_prompt_style</span><span class="p">,</span>
         <span class="n">mask_out_instructions_in_target</span>
       <span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
          <span class="s2">"expected :mask_out_instructions_in_target to be a boolean, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">mask_out_instructions_in_target</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; B, 2 =&gt; A, 3 =&gt; E, 4 =&gt; D, 5 =&gt; E, 6 =&gt; D</p>

<p>Correct answer: +</p>

<h3 id="74-creating-data-loaders-for-an-instruction-dataset">7.4 Creating data loaders for an instruction dataset</h3>

<p>1 What is the purpose of the <code class="language-plaintext highlighter-rouge">custom_collate_fn</code> function in the context of instruction fine-tuning?</p>

<p>My answer: The purpose is to prepare input data into batches with input/target pairs, to padd data with endoftext tokens and skip loss function calculation token.</p>

<p>Correct answer: <em>The custom_collate_fn function is used to batch the instruction dataset, ensuring that the input and target tensors are moved to the specified device (CPU, GPU, or MPS) before being fed into the LLM for fine-tuning.</em></p>

<p>2 Explain the advantage of moving data to the target device within the <code class="language-plaintext highlighter-rouge">custom_collate_fn</code> instead of the main training loop.</p>

<p>My answer: Advantage is faster calculation for training data.</p>

<p>Correct answer: <em>By performing the device transfer within the custom_collate_fn, the process becomes a background task, preventing it from blocking the GPU during model training and improving efficiency</em></p>

<p>3 How is the <code class="language-plaintext highlighter-rouge">device</code> setting determined and used in the <code class="language-plaintext highlighter-rouge">custom_collate_fn</code>?</p>

<p>My answer: Did not answer as Elixir implmentation is different.</p>

<p>Correct answer: <em>The device setting is determined based on the availability of a GPU or MPS. The partial function from functools is used to create a new version of the custom_collate_fn with the device argument prefilled, ensuring that the function uses the correct device for data transfer.</em></p>

<p>4 What is the purpose of the <code class="language-plaintext highlighter-rouge">allowed_max_length</code> parameter in the <code class="language-plaintext highlighter-rouge">customized_collate_fn</code>?</p>

<p>My answer: We truncate inputs and targets to the value of this parameter. If value is nil, we keep original size of inputs/targets.</p>

<p>Correct answer: <em>The allowed_max_length parameter is used to truncate the data to the maximum context length supported by the LLM model being fine-tuned, in this case, the GPT-2 model.</em></p>

<p>5 Describe the process of creating data loaders for the training, validation, and test sets using the <code class="language-plaintext highlighter-rouge">DataLoader</code> calss.</p>

<p>My answer: We randonly split input data set to the precentage of total input data set: training is 85%, validation is 10 % and rest is test set.</p>

<p>Correct answer: The DataLoader class is used to create data loaders for each set (training, validation, and test). The batch_size, collate_fn, shuffle, drop_last, and num_workers parameters are configured to control the batching process, shuffling, and data loading behavior_</p>

<p>6 Pieces of the code have been removed from four places in the listing. Which of these terms have been removed and where should they go?</p>

<p>A inputs B targets C maybe_truncate</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">def</span> <span class="n">custom_collate</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">,</span> <span class="n">ignore_index</span><span class="p">,</span> <span class="n">allowed_max_length</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">examples</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">normalize_example</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>

    <span class="c1"># Find the longest sequence length after the one extra pad token that the</span>
    <span class="c1"># Python collate function appends to every item.</span>
    <span class="n">batch_max_length</span> <span class="o">=</span>
      <span class="n">examples</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="nv">&amp;1</span><span class="o">.</span><span class="n">token_ids</span><span class="p">)</span> <span class="o">+</span> <span class="mi">1</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="mi">0</span> <span class="k">end</span><span class="p">)</span>

    <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">examples</span><span class="p">,</span> <span class="k">fn</span> <span class="p">%{</span><span class="ss">token_ids:</span> <span class="n">item</span><span class="p">,</span> <span class="ss">prompt_length:</span> <span class="n">prompt_length</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="c1"># Elixir data is immutable, so this builds the equivalent of</span>
        <span class="c1"># `new_item = item.copy(); new_item += [pad_token_id]`.</span>
        <span class="n">new_item</span> <span class="o">=</span> <span class="n">item</span> <span class="o">++</span> <span class="p">[</span><span class="n">pad_token_id</span><span class="p">]</span>

        <span class="c1"># Pad the copied item to the longest sequence length in the batch.</span>
        <span class="n">padded</span> <span class="o">=</span> <span class="n">pad_to_length</span><span class="p">(</span><span class="n">new_item</span><span class="p">,</span> <span class="n">batch_max_length</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">)</span>

        <span class="c1"># Match `padded[:-1]` for inputs and `padded[1:]` for next-token</span>
        <span class="c1"># targets.</span>
        <span class="n">inputs</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">padded</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">batch_max_length</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>

        <span class="n">targets</span> <span class="o">=</span>
          <span class="n">padded</span>
          <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">batch_max_length</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
          <span class="o">|&gt;</span> <span class="n">mask_extra_padding_targets</span><span class="p">(</span><span class="n">pad_token_id</span><span class="p">,</span> <span class="n">ignore_index</span><span class="p">)</span>
          <span class="o">|&gt;</span> <span class="n">mask_prompt_targets</span><span class="p">(</span><span class="n">prompt_length</span><span class="p">,</span> <span class="n">ignore_index</span><span class="p">)</span>

        <span class="c1"># Optionally cap both rows to the model context length.</span>
        <span class="p">{</span>
          <span class="n">maybe_truncate</span><span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="n">allowed_max_length</span><span class="p">),</span>
          <span class="n">__1__</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="n">allowed_max_length</span><span class="p">)</span>
        <span class="p">}</span>
      <span class="k">end</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">unzip</span><span class="p">()</span>

    <span class="p">{</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">__2__</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}),</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
    <span class="p">}</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; C, 2 =&gt; A</p>

<p>Correct answer: +</p>

<p>7 Match the terms on the left with their description on the right:</p>

<table>
  <thead>
    <tr>
      <th>Term</th>
      <th> </th>
      <th>Description</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>Batching<br />Process</td>
      <td> </td>
      <td>A function that defines how to combine individual data samples into a batch, specifically tailored for the requirements of instruction fine-tuning.</td>
    </tr>
    <tr>
      <td>Custom<br />Collate<br />Function</td>
      <td> </td>
      <td>The process of organizing training data into groups of samples, called batches, to improve training efficiency.</td>
    </tr>
    <tr>
      <td>Padding<br />Tokens</td>
      <td> </td>
      <td>Special tokens added to the end of sequences to ensure all inputs in a batch have the same length, allowing for efficient processing by the model.</td>
    </tr>
    <tr>
      <td>Ignore<br />Index</td>
      <td> </td>
      <td>A special value, typically -100, used to indicate padding tokens in the target sequence, preventing them from contributing to the loss calculation during training.</td>
    </tr>
  </tbody>
</table>

<p>My answer: 1 =&gt; 2, 2 =&gt; 1, 3 =&gt; 3, 4 =&gt; 4</p>

<p>Correct answer: +</p>

<h3 id="75-loading-a-pretrained-llm">7.5 Loading a pretrained LLM</h3>

<p>1 Why is larger pretrained model, like gpt2-medium, prefered for instruction fine-tuning compared to the smaller gpt2-small model?</p>

<p>My answer: Because pretrained medium model shows better performace for instruction fine-tuning.</p>

<p>Correct answer: <em>The smaller model lacks the capacity to learn and retain the complex patterns and nuanced behaviors required for high-quality instruction- following tasks. The larger model has more parameters, allowing it to handle more intricate instructions and generate more accurate responses.</em></p>

<p>Did not explain why.</p>

<p>2 What is the purpose of loading a pretrained LLM before starting instruction fine-tuning?</p>

<p>My answer: Purpose is that we reuse step of training the LLM model.</p>

<p>Correct answer: <em>Loading a pretrained LLM provides a foundation for the fine-tuning process. It allows the model to leverage existing knowledge and patterns learned during pretraining, enabling it to learn new tasks more efficiently and effectively.</em></p>

<p>Did not explain why.</p>

<p>3 How does the code for loading a pretrained model differ from the code used for pretraining of classification fine-tuning?</p>

<p>My answer: ?</p>

<p>Correct answer: <em>The code remains largely the same, but instead of specifying the ‘gpt2- small’ model, we now specify ‘gpt2-medium’ to load the larger model with 355 million parameters. This change reflects the choice of a more capable model for instruction fine-tuning.</em></p>

<p>Did not think of this simple answer.</p>

<p>4 What is the purpose of evaluating the pretrained LLM’s performance on a validation task before fine-tuning?</p>

<p>My answer: Purpose is to avoid oveffiting LLM model in pretrained phase.</p>

<p>Correct answer: <em>Evaluating the pretrained LLM’s performance provides a baseline understanding of its capabilities before fine-tuning. This allows us to assess the impact of fine-tuning on the model’s ability to follow instructions and generate accurate responses.</em></p>

<p>Wrong answer.</p>

<p>5 How is the model’s generated response isolated from the input instruction in the provided code?</p>

<p>My answer: It is isolated with dedicated tag: generated_response.</p>

<p>Correct answer:  <em>The code subtracts the length of the input instruction from the start of the generated text, effectively removing the input text and leaving only the model’s generated response. The strip() function is then applied to remove any extra whitespace.</em></p>

<p>I did not answer that we need to remove instructions part from the response.</p>

<h3 id="76-fine-tuning-the-llm-on-instruction-data">7.6 Fine-tuning the LLM on instruction data</h3>

<p>1 What is the purpose of fine-tuning an LLM on instruction data?</p>

<p>My Answer: The purpose of fine-tuning an LLM on instruction data is to train LLM to understands the patterns for resolving instruction tasks.</p>

<p>Correct answer: <em>Fine-tuning an LLM on instruction data aims to improve its ability to understand and follow instructions, leading to more accurate and relevant responses to user prompts.</em></p>

<p>2 Describe the process of fine-tuning an LLM on instruction data, highlighting the key steps involved.</p>

<p>My answer: We first format instructions input dataset in form of prompt. Then we split input data into batches where we padd each input to the longest input in particular batch. Every batch could have different max length. Padding token is |endoftext| token. Then we create input/target pairs where target is input shiffted for one token to the right. If we reached last input token, then we add <code class="language-plaintext highlighter-rouge">endoftext</code> token. We first padd targets then replace all padded tokens with -100, except for the first one that is closest to last token. -100 value are skipped while calculating the loss function. Then we train the model on that data using same algorithm as for pretraining step.</p>

<p>Correct answer: <em>The process involves loading a pretrained LLM, preparing an instruction dataset, and training the model on this dataset using a suitable loss function and optimizer. The training process aims to minimize the loss, indicating improved performance in following instructions.</em></p>

<p>My answer is wrong because I did not mention loading pre-trained model.</p>

<p>3 What are some potential challanges encountered during fine-tuning an LLM on instruction data, and how can these challanges be addressed?</p>

<p>My answer: ?</p>

<p>Correct answer: <em>Challenges include hardware limitations, such as memory constraints, which can be addressed by using a smaller model, reducing the batch size, or utilizing a GPU for faster training. Additionally, managing the length of input sequences can be crucial for efficient training.</em></p>

<p>4 How is the effectivness of fine-tuning evaluated during the training process?</p>

<p>My answer: Input set is randomly split in training, validation and test sets. We are calculating entropy loss function for those steps.</p>

<p>Correct answer: <em>The effectiveness is evaluated by monitoring the training and validation losses. A decrease in these losses indicates that the model is learning to follow instructions better. Additionally, inspecting the generated responses during training provides qualitative insights into the model’s progress.</em></p>

<p>5 What is the significance of the Alpaca dataset in the context of fine-tuning LLMs?</p>

<p>My answer: Alpaca dataset is in form of prompt, setting first the context for instruction that should be resolved.</p>

<p>Correct answer: <em>The Alpaca dataset is a valuable resource for fine-tuning LLMs on instruction data. It provides a large collection of diverse instructions and corresponding responses, enabling the model to learn a wider range of task-specific behaviors.</em></p>

<p>Its a larger instruction training dataset.</p>

<h3 id="77-extracting-and-saving-responses">7.7 Extracting and saving responses</h3>

<p>1 Describe the process of evaluating the performance of an instruction-fine-tuned Large Language Model (LLM) after training.</p>

<p>My answer: We have two options. First is that humans checks the quality of resolving instruction tasks. This technique does not scale, so we have second option, which is to use another llm for instruction tasks resloving. We asked for example ollama3 model to score the instruction task model resolution.</p>

<p>Correct answer: <em>Evaluation involves extracting model-generated responses from a held- out test set, manually analyzing them, and then quantifying the response quality using various methods such as benchmarks, human comparison, or automated metrics.</em></p>

<p>2 What are the different methods for evaluating the performance of instruction-fine-tuned LLMs, and what are their relative strengths ans weaknesses?</p>

<p>My answer: Human evaluation and another llm score evaluation. Human evaluation does not scale well.</p>

<p>Correct answer: <em>Methods include short-answer benchmarks (e.g., MMLU), human preference comparisons, and automated conversational benchmarks (e.g., AlpacaEval). Human evaluation provides valuable insights but is time-consuming, while automated methods are efficient but may lack the nuance of human judgment.</em></p>

<p>3 What is meant by conversational performance in the context of LLMs, and why is it important?</p>

<p>My answer: ?</p>

<p>Correct answer: <em>Conversational performance refers to an LLM’s ability to engage in human-like communication, understanding context, nuance, and intent.  It’s crucial for applications like chatbots where natural and coherent interaction is essential.</em></p>

<p>Understanding context, naunce and intent.</p>

<p>4 How can the response generated by an LLM be automatically evaluated using another LLM, and what are the advantages of this approach?</p>

<p>My answer: We construct a prompt that askes from another llm to score the response of our model. We send to model instructions and correct response.</p>

<p>Correct answer: <em>An approach similar to AlpacaEval can be used, employing another LLM to evaluate the responses. This automated method is efficient, saving time and resources compared to manual human evaluation while still providing meaningful performance indicators</em></p>

<p>5 Explain the process of appending generated model responses to a test set and saving the updated data for later analysis.</p>

<p>My answer: ?</p>

<p>Correct answer: <em>The generated model responses are added to a dictionary containing the test data. This updated data is then saved as a JSON file (e.g., ‘instruction-data-with-response.json’) for easy access and analysis in future sessions.</em></p>

<p>Saving model responses to json file.</p>

<h3 id="78-evaluating-the-fine-tuned-llm">7.8 Evaluating the fine-tuned LLM</h3>

<p>1 What is the purpose of using a larger LLM to evaluate the responses of a fine-tuned model?</p>

<p>My answer: Purpose is to automate evaluation process of our fine-tuned llm.</p>

<p>Correct answer: <em>A larger LLM, like Llama 3 or GPT-4, can be used to automatically assess the quality of responses generated by a fine-tuned model. This provides a more objective and scalable method for evaluating model performance compared to manually reviewing a few examples</em></p>

<p>2 Describe the process of using Ollama to evaluate the responses of a fine-tuned model.</p>

<p>My answer: We use testing instruction dataset. In json file we store instruction, model response, and correnct response (output). Then we send to ollama3 input, output and model response to evaluate model response. For that we create a prompt where we are asing ollama3 to score (0-100) model response.</p>

<p>Correct answer: <em>Ollama, an open-source application, allows you to run LLMs locally. You can use the query_model function to send prompts to the LLM, such as asking it to score a model’s response on a scale of 0 to 100. The LLM’s evaluation can then be used to assess the overall performance of the fine-tuned model.</em></p>

<p>3 What aer some alternative LLMs that chan be used for evaluating model responses, besides the 8-billion-parameter Llama 3 model?</p>

<p>My answer: We could use Qwen 8B model.</p>

<p>Correct answer: <em>Other LLMs, such as the 3.8-billion-parameter phi3 model or the larger 70-billion-parameter Llama 3 model, can be used with Ollama. The choice of model depends on the available computational resources and the desired level of performance.</em></p>

<p>4 How can the generate_model_score function be used to assess the performance of a fine-tuned model?</p>

<p>My answer: Each entry from test instructions dataset is sent to ollama3 evaluation in the form of prompt that request response score 1-100 based on input, expected output and output given by the model.</p>

<p>Correct answer: <em>The generate_model_scores function iterates through a set of test data, sending prompts to the LLM to evaluate each model response. It then calculates the average score across all responses, providing a quantitative measure of the model’s performance.</em></p>

<p>5 What are some strategies for improving the performance of a fine-tuned model?</p>

<p>My answer: We could implement Lora for model linear layers. Or we could do training on broader instructions data set, like alpaca data set.</p>

<p>Correct answer: <em>Strategies for improving model performance include adjusting hyperparameters during fine-tuning, increasing the size or diversity of the training dataset, experimenting with different prompts or instruction formats, and using a larger pretrained model.</em></p>

<p>6 Match the terms on the left with their description on the right:</p>

<table>
  <thead>
    <tr>
      <th> </th>
      <th> </th>
      <th> </th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>Test Set</td>
      <td> </td>
      <td>A method of evaluating the conversational performance of a language model using another language model to assess the quality of responses.</td>
    </tr>
    <tr>
      <td>Conversational<br />Performance</td>
      <td> </td>
      <td>The ability of a language model to engage in human-like communication, understanding context, nuance, and intent.</td>
    </tr>
    <tr>
      <td>Automated<br />Conversational<br />Benchmarks</td>
      <td> </td>
      <td>A portion of data that is held out from the training process and used to evaluate the performance of a trained model.</td>
    </tr>
  </tbody>
</table>

<p>My answer: 1 =&gt; 3, 2 =&gt; 2, 3 =&gt; 1</p>

<p>Correct answer: +</p>]]></content><author><name></name></author><category term="llm-from-scratch" /><category term="gpt" /><category term="tutorials" /><category term="llm" /><category term="gpt" /><category term="workbook" /><summary type="html"><![CDATA[Workbook notes on instruction fine-tuning a pretrained GPT-2 model in Elixir.]]></summary></entry><entry><title type="html">Fine-tuning to follow instructions</title><link href="https://karlosmid.com/2026/06/fine-tuning-to-follow-instructions/" rel="alternate" type="text/html" title="Fine-tuning to follow instructions" /><published>2026-06-15T09:00:00+00:00</published><updated>2026-06-15T09:00:00+00:00</updated><id>https://karlosmid.com/2026/06/fine-tuning-to-follow-instructions</id><content type="html" xml:base="https://karlosmid.com/2026/06/fine-tuning-to-follow-instructions/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>This post fine-tunes a pretrained GPT-2 355M model on instruction data in Elixir. It prepares an Alpaca-style dataset, builds custom batching and target masking, trains the model for two epochs, saves checkpoints and loss metrics, extracts test responses, and evaluates them with a local Ollama model. The fine-tuned Alpaca-style model reaches an average Ollama score of about 50, while the Phi-3 prompt style and instruction/input masking variants perform slightly worse in these experiments.</p>

<p>This post follows <a href="/fine-tuning-for-classification/">Fine-tuning for classification</a> and uses the source code from my <a href="https://github.com/karlosmid/llm-from-scratch">GitHub repository</a>.</p>

<h2 id="71-introduction-to-instruction-fine-tuning">7.1 Introduction to instruction fine-tuning</h2>

<p>What does instruction fine-tuning mean? Pretraining teaches an LLM how to continue text from the input by predicting the most probable next token. With that objective, an LLM can finish sentences or create paragraphs. After instruction fine-tuning, the LLM can learn to follow commands, such as translating text into Croatian or correcting English grammar in a given input. As with classification, we need to go through several steps: download and prepare the dataset, organize it into batches, create data loaders, load a pretrained model, instruction fine-tune that model, inspect the modeling loss, extract responses, perform qualitative evaluation, and score the responses.</p>

<h2 id="72-preparing-a-dataset-for-supervised-instruction-fine-tuning">7.2 Preparing a dataset for supervised instruction fine-tuning</h2>

<p>The first step is to get the training data and prepare it for instruction training with our pretrained OpenAI GPT-2 model. We are extending the <code class="language-plaintext highlighter-rouge">FineTuneDataLoader</code> module with a function that downloads and saves instruction training data in JSON format. The instruction training data was prepared by Sebastian, the book author:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@instruction_data_url</span> <span class="s2">"https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch07/01_main-chapter-code/instruction-data.json"</span>

<span class="nv">@type</span> <span class="n">instruction_record</span> <span class="p">::</span> <span class="p">%{</span>
          <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span> <span class="o">=&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Downloads and loads the instruction fine-tuning dataset.

  If `file_path` already exists, the local file is reused. Otherwise the JSON
  file is downloaded from `url`, saved to `file_path`, and then decoded.

  ## Input Parameters

    * `file_path` - local JSON cache path. Defaults to
      `"instruction-data.json"`.
    * `url` - URL of the instruction dataset JSON. Defaults to the chapter 7
      instruction-data file from the LLMs-from-scratch repository.

  ## Output

  Returns the decoded JSON as a list of maps with string keys. Each decoded
  instruction map contains:

    * `"instruction"` - the task description or question.
    * `"input"` - optional context for the task. This is an empty string when
      no extra input is provided.
    * `"output"` - the expected response used as the supervised fine-tuning
      target.
  """</span>
  <span class="nv">@spec</span> <span class="n">download_and_load_instructions_file</span><span class="p">(</span><span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="p">[</span><span class="n">instruction_record</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">download_and_load_instructions_file</span><span class="p">(</span>
        <span class="n">file_path</span> <span class="p">\\</span> <span class="nv">@instruction_data_path</span><span class="p">,</span>
        <span class="n">url</span> <span class="p">\\</span> <span class="nv">@instruction_data_url</span>
      <span class="p">)</span> <span class="k">do</span>
    <span class="k">unless</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">file_path</span><span class="p">)</span> <span class="k">do</span>
      <span class="no">File</span><span class="o">.</span><span class="n">mkdir_p!</span><span class="p">(</span><span class="no">Path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">file_path</span><span class="p">))</span>
      <span class="n">download_file!</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="n">file_path</span><span class="p">)</span>
    <span class="k">end</span>

    <span class="n">file_path</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Jason</span><span class="o">.</span><span class="n">decode!</span><span class="p">()</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>The next function we need formats the training data into the Alpaca prompt format that we feed into our LLM as the training dataset:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="nv">@doc</span> <span class="sd">"""
  Formats an instruction dataset entry as an Alpaca-style model input.

  This mirrors the chapter 7 prompt template:

      Below is an instruction that describes a task. Write a response that
      appropriately completes the request.

      ### Instruction:
      ...

      ### Input:
      ...

  The input section is included only when the entry's `"input"` field is not an
  empty string.

  ## Input Parameters

    * `entry` - map containing string keys `"instruction"` and `"input"`, such
      as records returned by `download_and_load_instructions_file/2`.

  ## Output

  Returns the formatted prompt string without the desired response.
  """</span>
  <span class="nv">@spec</span> <span class="n">format_input</span><span class="p">(</span><span class="n">instruction_record</span><span class="p">())</span> <span class="p">::</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">format_input</span><span class="p">(%{</span><span class="s2">"instruction"</span> <span class="o">=&gt;</span> <span class="n">instruction</span><span class="p">,</span> <span class="s2">"input"</span> <span class="o">=&gt;</span> <span class="n">input</span><span class="p">})</span>
      <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">instruction</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">input</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">instruction_text</span> <span class="o">=</span>
      <span class="s2">"Below is an instruction that describes a task. "</span> <span class="o">&lt;&gt;</span>
        <span class="s2">"Write a response that appropriately completes the request."</span> <span class="o">&lt;&gt;</span>
        <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Instruction:</span><span class="se">\n</span><span class="si">#{</span><span class="n">instruction</span><span class="si">}</span><span class="s2">"</span>

    <span class="n">input_text</span> <span class="o">=</span>
      <span class="k">if</span> <span class="n">input</span> <span class="o">==</span> <span class="s2">""</span> <span class="k">do</span>
        <span class="s2">""</span>
      <span class="k">else</span>
        <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Input:</span><span class="se">\n</span><span class="si">#{</span><span class="n">input</span><span class="si">}</span><span class="s2">"</span>
      <span class="k">end</span>

    <span class="n">instruction_text</span> <span class="o">&lt;&gt;</span> <span class="n">input_text</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>We split the input data into training, validation, and test sets for the same reason as in classification training. Here is the test that uses those two new functions:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="k">defmodule</span> <span class="no">LlmFromScratch7Test</span> <span class="k">do</span>
  <span class="kn">use</span> <span class="no">ExUnit</span><span class="o">.</span><span class="no">Case</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">FineTuneDataLoader</span>

  <span class="nv">@instruction_data_url</span> <span class="s2">"https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch07/01_main-chapter-code/instruction-data.json"</span>

  <span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="n">test</span> <span class="s2">"7.1 downloads and loads instruction data"</span> <span class="k">do</span>
    <span class="n">file_path</span> <span class="o">=</span>
      <span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"llm_scratch_instruction_data"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"instruction-data.json"</span><span class="p">)</span>

    <span class="n">data</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">download_and_load_instructions_file</span><span class="p">(</span>
        <span class="n">file_path</span><span class="p">,</span>
        <span class="nv">@instruction_data_url</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1_100</span>

    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="mi">50</span><span class="p">)</span> <span class="o">==</span> <span class="p">%{</span>
             <span class="s2">"instruction"</span> <span class="o">=&gt;</span> <span class="s2">"Identify the correct spelling of the following word."</span><span class="p">,</span>
             <span class="s2">"input"</span> <span class="o">=&gt;</span> <span class="s2">"Ocassion"</span><span class="p">,</span>
             <span class="s2">"output"</span> <span class="o">=&gt;</span> <span class="s2">"The correct spelling is 'Occasion.'"</span>
           <span class="p">}</span>

    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="mi">999</span><span class="p">)</span> <span class="o">==</span> <span class="p">%{</span>
             <span class="s2">"instruction"</span> <span class="o">=&gt;</span> <span class="s2">"What is an antonym of 'complicated'?"</span><span class="p">,</span>
             <span class="s2">"input"</span> <span class="o">=&gt;</span> <span class="s2">""</span><span class="p">,</span>
             <span class="s2">"output"</span> <span class="o">=&gt;</span> <span class="s2">"An antonym of 'complicated' is 'simple'."</span>
           <span class="p">}</span>

    <span class="n">entry_50</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="mi">50</span><span class="p">)</span>
    <span class="n">model_input_50</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="n">entry_50</span><span class="p">)</span>
    <span class="n">desired_response_50</span> <span class="o">=</span> <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Response:</span><span class="se">\n</span><span class="si">#{</span><span class="n">entry_50</span><span class="p">[</span><span class="s2">"output"</span><span class="p">]</span><span class="si">}</span><span class="s2">"</span>

    <span class="n">assert</span> <span class="n">model_input_50</span> <span class="o">&lt;&gt;</span> <span class="n">desired_response_50</span> <span class="o">==</span>
             <span class="s2">"Below is an instruction that describes a task. "</span> <span class="o">&lt;&gt;</span>
               <span class="s2">"Write a response that appropriately completes the request."</span> <span class="o">&lt;&gt;</span>
               <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Instruction:</span><span class="se">\n</span><span class="s2">Identify the correct spelling of the following word."</span> <span class="o">&lt;&gt;</span>
               <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Input:</span><span class="se">\n</span><span class="s2">Ocassion"</span> <span class="o">&lt;&gt;</span>
               <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Response:</span><span class="se">\n</span><span class="s2">The correct spelling is 'Occasion.'"</span>

    <span class="n">entry_999</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="mi">999</span><span class="p">)</span>
    <span class="n">model_input_999</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="n">entry_999</span><span class="p">)</span>
    <span class="n">desired_response_999</span> <span class="o">=</span> <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Response:</span><span class="se">\n</span><span class="si">#{</span><span class="n">entry_999</span><span class="p">[</span><span class="s2">"output"</span><span class="p">]</span><span class="si">}</span><span class="s2">"</span>

    <span class="n">assert</span> <span class="n">model_input_999</span> <span class="o">&lt;&gt;</span> <span class="n">desired_response_999</span> <span class="o">==</span>
             <span class="s2">"Below is an instruction that describes a task. "</span> <span class="o">&lt;&gt;</span>
               <span class="s2">"Write a response that appropriately completes the request."</span> <span class="o">&lt;&gt;</span>
               <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Instruction:</span><span class="se">\n</span><span class="s2">What is an antonym of 'complicated'?"</span> <span class="o">&lt;&gt;</span>
               <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Response:</span><span class="se">\n</span><span class="s2">An antonym of 'complicated' is 'simple'."</span>

    <span class="n">train_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.85</span><span class="p">)</span>
    <span class="n">test_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">val_portion</span> <span class="o">=</span> <span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">-</span> <span class="n">train_portion</span> <span class="o">-</span> <span class="n">test_portion</span>

    <span class="n">train_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">train_portion</span><span class="p">)</span>
    <span class="n">test_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span><span class="p">,</span> <span class="n">test_portion</span><span class="p">)</span>
    <span class="n">val_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span> <span class="o">+</span> <span class="n">test_portion</span><span class="p">,</span> <span class="n">val_portion</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_data</span><span class="p">)</span> <span class="o">==</span> <span class="mi">935</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">val_data</span><span class="p">)</span> <span class="o">==</span> <span class="mi">55</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">test_data</span><span class="p">)</span> <span class="o">==</span> <span class="mi">110</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<h2 id="73-organizing-data-into-training-batches">7.3 Organizing data into training batches</h2>

<p>The process of organizing instruction data into training batches is more complicated than organizing classification data or pretraining data into batches. The reason is that instruction data uses a prompt format. Sebastian breaks this process into five steps:</p>

<ul>
  <li>use prompt template</li>
  <li>tokenize prompt templates</li>
  <li>pad inputs to the same length</li>
  <li>create target token IDs</li>
  <li>replace padding tokens with placeholders</li>
</ul>

<p>Here are the first two steps implemented in the <code class="language-plaintext highlighter-rouge">InstructionDataset</code> module, along with a test for this new functionality:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code> <span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">InstructionDataset</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Pre-tokenized instruction fine-tuning dataset.

  This mirrors the PyTorch `InstructionDataset` example from chapter 7. A
  dataset is built from decoded instruction maps, formats each map as an
  Alpaca-style prompt plus response, and pre-tokenizes the full text during
  construction.

  ## Examples

      iex&gt; data = [
      ...&gt;   %{
      ...&gt;     "instruction" =&gt; "Name the capital of France.",
      ...&gt;     "input" =&gt; "",
      ...&gt;     "output" =&gt; "Paris."
      ...&gt;   }
      ...&gt; ]
      iex&gt; dataset = LlmScratch.InstructionDataset.new(data, "code-davinci-002")
      iex&gt; LlmScratch.InstructionDataset.length(dataset)
      1
      iex&gt; LlmScratch.InstructionDataset.get(dataset, 0)
      [21106, 318, 281, 12064, 326, 8477, 257, 4876, 13, 19430, 257,
       2882, 326, 20431, 32543, 262, 2581, 13, 198, 198, 21017, 46486,
       25, 198, 5376, 262, 3139, 286, 4881, 13, 198, 198, 21017, 18261,
       25, 198, 40313, 13]
  """</span>

  <span class="nv">@enforce_keys</span> <span class="p">[</span><span class="ss">:data</span><span class="p">,</span> <span class="ss">:encoded_texts</span><span class="p">]</span>
  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:data</span><span class="p">,</span> <span class="ss">:encoded_texts</span><span class="p">]</span>

  <span class="nv">@end_of_text</span> <span class="s2">"&lt;|endoftext|&gt;"</span>

  <span class="nv">@type</span> <span class="n">instruction_record</span> <span class="p">::</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">instruction_record</span><span class="p">()</span>
  <span class="nv">@type</span> <span class="n">tokenizer</span> <span class="p">::</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">data:</span> <span class="p">[</span><span class="n">instruction_record</span><span class="p">()],</span>
          <span class="ss">encoded_texts:</span> <span class="p">[[</span><span class="n">integer</span><span class="p">()]]</span>
        <span class="p">}</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Creates an instruction dataset from decoded instruction records.

  Each record is formatted with `LlmScratch.FineTuneDataLoader.format_input/1`,
  then the target response is appended as:

      \\n\\n### Response:\\n...

  The resulting full text is encoded immediately and stored in
  `dataset.encoded_texts`.

  ## Input Parameters

    * `data` - list of decoded instruction maps. Each map must contain string
      keys `"instruction"`, `"input"`, and `"output"`.
    * `tokenizer` - Tiktoken model name accepted by `Tiktoken.encode/3`.
    * `opts` - optional keyword list controlling tokenization.

  ## Options

    * `:allowed_special` - special tokens allowed when `tokenizer` is a
      Tiktoken model name. Defaults to `["&lt;|endoftext|&gt;"]`.

  ## Output

  Returns a `%LlmScratch.InstructionDataset{}`. The `:data` field contains the
  original records, and `:encoded_texts` contains one pre-tokenized token-id
  list per record.

  ## Examples

      iex&gt; data = [
      ...&gt;   %{
      ...&gt;     "instruction" =&gt; "Classify the sentiment.",
      ...&gt;     "input" =&gt; "I loved it.",
      ...&gt;     "output" =&gt; "Positive."
      ...&gt;   }
      ...&gt; ]
      iex&gt; dataset = LlmScratch.InstructionDataset.new(data, "code-davinci-002")
      iex&gt; dataset.data
      [
        %{
          "input" =&gt; "I loved it.",
          "instruction" =&gt; "Classify the sentiment.",
          "output" =&gt; "Positive."
        }
      ]
      iex&gt; dataset.encoded_texts
      [
        [21106, 318, 281, 12064, 326, 8477, 257, 4876, 13, 19430, 257,
         2882, 326, 20431, 32543, 262, 2581, 13, 198, 198, 21017, 46486,
         25, 198, 9487, 1958, 262, 15598, 13, 198, 198, 21017, 23412, 25,
         198, 40, 6151, 340, 13, 198, 198, 21017, 18261, 25, 198, 21604,
         1800, 13]
      ]
  """</span>
  <span class="nv">@spec</span> <span class="n">new</span><span class="p">([</span><span class="n">instruction_record</span><span class="p">()],</span> <span class="n">tokenizer</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="k">do</span>
    <span class="c1"># Pre-tokenize each complete training example once so repeated dataset</span>
    <span class="c1"># access does not rebuild the prompt or call the tokenizer again.</span>
    <span class="n">encoded_texts</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="k">fn</span> <span class="n">entry</span> <span class="o">-&gt;</span>
        <span class="n">entry</span>
        <span class="c1"># The model trains on the instruction/input prompt followed by the</span>
        <span class="c1"># expected response target, matching the chapter 7 PyTorch dataset.</span>
        <span class="o">|&gt;</span> <span class="n">full_text</span><span class="p">()</span>
        <span class="o">|&gt;</span> <span class="n">encode_text!</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="c1"># Keep the decoded records for inspection while serving encoded examples</span>
    <span class="c1"># from `encoded_texts`, which mirrors the PyTorch dataset fields.</span>
    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">data:</span> <span class="n">data</span><span class="p">,</span>
      <span class="ss">encoded_texts:</span> <span class="n">encoded_texts</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns one encoded instruction example by index.

  ## Input Parameters

    * `dataset` - `%LlmScratch.InstructionDataset{}` returned by `new/3`.
    * `index` - zero-based row index.

  ## Output

  Returns the pre-tokenized token-id list for the requested record.
  """</span>
  <span class="nv">@spec</span> <span class="n">get</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="n">non_neg_integer</span><span class="p">())</span> <span class="p">::</span> <span class="p">[</span><span class="n">integer</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">get</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">dataset</span><span class="p">,</span> <span class="n">index</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">index</span><span class="p">)</span> <span class="ow">and</span> <span class="n">index</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="no">Enum</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span> <span class="n">index</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns the number of records in the dataset.

  ## Input Parameters

    * `dataset` - `%LlmScratch.InstructionDataset{}` returned by `new/3`.

  ## Output

  Returns the number of instruction records as a non-negative integer.
  """</span>
  <span class="nv">@spec</span> <span class="n">length</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">length</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">dataset</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">dataset</span><span class="o">.</span><span class="n">data</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">full_text</span><span class="p">(%{</span><span class="s2">"output"</span> <span class="o">=&gt;</span> <span class="n">output</span><span class="p">}</span> <span class="o">=</span> <span class="n">entry</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">output</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">instruction_plus_input</span> <span class="o">=</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="n">entry</span><span class="p">)</span>
    <span class="n">response_text</span> <span class="o">=</span> <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Response:</span><span class="se">\n</span><span class="si">#{</span><span class="n">output</span><span class="si">}</span><span class="s2">"</span>

    <span class="n">instruction_plus_input</span> <span class="o">&lt;&gt;</span> <span class="n">response_text</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">encode_text!</span><span class="p">(</span><span class="n">text</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">allowed_special</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:allowed_special</span><span class="p">,</span> <span class="p">[</span><span class="nv">@end_of_text</span><span class="p">])</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">token_ids</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="n">text</span><span class="p">,</span> <span class="n">allowed_special</span><span class="p">)</span>
    <span class="n">token_ids</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>The next step is to pad inputs to the same length. Note that each batch can have a different length. We pad to the length of the longest input in the batch.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@pad_token_id</span> <span class="mi">50_256</span>
<span class="nv">@doc</span> <span class="sd">"""
  Pads a batch of token-id sequences and stacks them into an input tensor.

  This mirrors the first draft of the chapter 7 Python collate function:

      inputs_1 = [0, 1, 2, 3, 4]
      inputs_2 = [5, 6]
      inputs_3 = [7, 8, 9]
      batch = (inputs_1, inputs_2, inputs_3)
      print(custom_collate_draft_1(batch))

  The function finds the longest sequence length in the batch after adding one
  extra pad token, pads every item to that length, then drops the final token
  from each row before stacking the rows. This draft returns only the input
  tensor; it does not build shifted targets yet.

  ## Input Parameters

    * `batch` - tuple or list of token-id lists.
    * `pad_token_id` - token id used for padding. Defaults to GPT-2's
      end-of-text token id, `50256`.
    * `device` - included for parity with the Python example. Nx places the
      tensor on the active backend, so this argument is currently informational.

  ## Output

  Returns a signed 64-bit Nx tensor with shape
  `{batch_size, max_sequence_length}`.

  ## Examples

      iex&gt; batch = {[0, 1, 2, 3, 4], [5, 6], [7, 8, 9]}
      iex&gt; inputs = LlmScratch.InstructionDataset.custom_collate_draft_1(batch)
      iex&gt; inputs
      #Nx.Tensor&lt;
        s64[3][5]
        [
          [0, 1, 2, 3, 4],
          [5, 6, 50256, 50256, 50256],
          [7, 8, 9, 50256, 50256]
        ]
      &gt;
  """</span>
  <span class="nv">@spec</span> <span class="n">custom_collate_draft_1</span><span class="p">(</span><span class="n">tuple</span><span class="p">()</span> <span class="o">|</span> <span class="p">[[</span><span class="n">integer</span><span class="p">()]],</span> <span class="n">integer</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span>
          <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">custom_collate_draft_1</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">pad_token_id</span> <span class="p">\\</span> <span class="nv">@pad_token_id</span><span class="p">,</span> <span class="n">device</span> <span class="p">\\</span> <span class="s2">"cpu"</span><span class="p">)</span>

  <span class="k">def</span> <span class="n">custom_collate_draft_1</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_tuple</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">batch</span>
    <span class="o">|&gt;</span> <span class="no">Tuple</span><span class="o">.</span><span class="n">to_list</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">custom_collate_draft_1</span><span class="p">(</span><span class="n">pad_token_id</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">def</span> <span class="n">custom_collate_draft_1</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">,</span> <span class="n">_device</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="c1"># Find the longest sequence length after the one extra pad token that the</span>
    <span class="c1"># Python draft appends to every item.</span>
    <span class="n">batch_max_length</span> <span class="o">=</span>
      <span class="n">batch</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">)</span> <span class="o">+</span> <span class="mi">1</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="mi">0</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">inputs</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="k">fn</span> <span class="n">item</span> <span class="o">-&gt;</span>
        <span class="c1"># Elixir data is immutable, so this builds the equivalent of</span>
        <span class="c1"># `new_item = item.copy(); new_item += [pad_token_id]`.</span>
        <span class="n">new_item</span> <span class="o">=</span> <span class="n">item</span> <span class="o">++</span> <span class="p">[</span><span class="n">pad_token_id</span><span class="p">]</span>

        <span class="c1"># Pad the copied item to the longest sequence length in the batch.</span>
        <span class="n">padded</span> <span class="o">=</span> <span class="n">pad_to_length</span><span class="p">(</span><span class="n">new_item</span><span class="p">,</span> <span class="n">batch_max_length</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">)</span>

        <span class="c1"># Match `torch.tensor(padded[:-1])` by dropping the final token before</span>
        <span class="c1"># stacking all rows into a single tensor.</span>
        <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">padded</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">batch_max_length</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>

  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">pad_to_length</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">max_length</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">token_ids</span> <span class="o">++</span> <span class="no">List</span><span class="o">.</span><span class="n">duplicate</span><span class="p">(</span><span class="n">pad_token_id</span><span class="p">,</span> <span class="n">max_length</span> <span class="o">-</span> <span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">token_ids</span><span class="p">))</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>And here is a test that uses this function on simple example data:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"7.2 custom collate draft pads batch inputs"</span> <span class="k">do</span>
    <span class="n">inputs_1</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span>
    <span class="n">inputs_2</span> <span class="o">=</span> <span class="p">[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">]</span>
    <span class="n">inputs_3</span> <span class="o">=</span> <span class="p">[</span><span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">]</span>
    <span class="n">batch</span> <span class="o">=</span> <span class="p">{</span><span class="n">inputs_1</span><span class="p">,</span> <span class="n">inputs_2</span><span class="p">,</span> <span class="n">inputs_3</span><span class="p">}</span>

    <span class="n">inputs_tensor</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">custom_collate_draft_1</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">inputs_tensor</span> <span class="o">==</span>
             <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span>
               <span class="p">[</span>
                 <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">]</span>
               <span class="p">],</span>
               <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
             <span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Each batch needs corresponding target tokens along with the inputs. The trick is the same as before: we get the input targets by shifting the input to the right by one token. If an input does not have a token on the right, we add a padding token:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@doc</span> <span class="sd">"""
  Pads a batch of token-id sequences and creates input/target tensors.

  This mirrors the second draft of the chapter 7 Python collate function:

      inputs_1 = [0, 1, 2, 3, 4]
      inputs_2 = [5, 6]
      inputs_3 = [7, 8, 9]
      batch = (inputs_1, inputs_2, inputs_3)
      inputs, targets = custom_collate(batch)
      print(inputs)
      print(targets)

  The function finds the longest sequence length in the batch after adding one
  extra pad token, pads every item to that length, then creates next-token
  prediction pairs. Inputs drop the final token from each padded row; targets
  drop the first token.

  ## Input Parameters

    * `batch` - tuple or list of token-id lists.
    * `pad_token_id` - token id used for padding. Defaults to GPT-2's
      end-of-text token id, `50256`.
    * `device` - included for parity with the Python example. Nx places the
      tensor on the active backend, so this argument is currently informational.

  ## Output

  Returns `{inputs, targets}` where both values are signed 64-bit Nx tensors
  with shape `{batch_size, max_sequence_length}`.

  ## Examples

      iex&gt; batch = {[0, 1, 2, 3, 4], [5, 6], [7, 8, 9]}
      iex&gt; {inputs, targets} = LlmScratch.InstructionDataset.custom_collate(batch)
      iex&gt; inputs
      #Nx.Tensor&lt;
        s64[3][5]
        [
          [0, 1, 2, 3, 4],
          [5, 6, 50256, 50256, 50256],
          [7, 8, 9, 50256, 50256]
        ]
      &gt;
      iex&gt; targets
      #Nx.Tensor&lt;
        s64[3][5]
        [
          [1, 2, 3, 4, 50256],
          [6, 50256, 50256, 50256, 50256],
          [8, 9, 50256, 50256, 50256]
        ]
      &gt;
  """</span>
  <span class="nv">@spec</span> <span class="n">custom_collate</span><span class="p">(</span><span class="n">tuple</span><span class="p">()</span> <span class="o">|</span> <span class="p">[[</span><span class="n">integer</span><span class="p">()]],</span> <span class="n">integer</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span>
          <span class="p">{</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()}</span>
  <span class="k">def</span> <span class="n">custom_collate</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">pad_token_id</span> <span class="p">\\</span> <span class="nv">@pad_token_id</span><span class="p">,</span> <span class="n">device</span> <span class="p">\\</span> <span class="s2">"cpu"</span><span class="p">)</span>

  <span class="k">def</span> <span class="n">custom_collate</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_tuple</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">batch</span>
    <span class="o">|&gt;</span> <span class="no">Tuple</span><span class="o">.</span><span class="n">to_list</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">custom_collate</span><span class="p">(</span><span class="n">pad_token_id</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">def</span> <span class="n">custom_collate</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">,</span> <span class="n">_device</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="c1"># Find the longest sequence length after the one extra pad token that the</span>
    <span class="c1"># Python collate function appends to every item.</span>
    <span class="n">batch_max_length</span> <span class="o">=</span>
      <span class="n">batch</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">)</span> <span class="o">+</span> <span class="mi">1</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="mi">0</span> <span class="k">end</span><span class="p">)</span>

    <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="k">fn</span> <span class="n">item</span> <span class="o">-&gt;</span>
        <span class="c1"># Elixir data is immutable, so this builds the equivalent of</span>
        <span class="c1"># `new_item = item.copy(); new_item += [pad_token_id]`.</span>
        <span class="n">new_item</span> <span class="o">=</span> <span class="n">item</span> <span class="o">++</span> <span class="p">[</span><span class="n">pad_token_id</span><span class="p">]</span>

        <span class="c1"># Pad the copied item to the longest sequence length in the batch.</span>
        <span class="n">padded</span> <span class="o">=</span> <span class="n">pad_to_length</span><span class="p">(</span><span class="n">new_item</span><span class="p">,</span> <span class="n">batch_max_length</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">)</span>

        <span class="c1"># Match `padded[:-1]` for inputs and `padded[1:]` for next-token</span>
        <span class="c1"># targets before stacking all rows into tensors.</span>
        <span class="p">{</span>
          <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">padded</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">batch_max_length</span> <span class="o">-</span> <span class="mi">1</span><span class="p">),</span>
          <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">padded</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">batch_max_length</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
        <span class="p">}</span>
      <span class="k">end</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">unzip</span><span class="p">()</span>

    <span class="p">{</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}),</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
    <span class="p">}</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>And here is a test that uses <code class="language-plaintext highlighter-rouge">custom_collate</code>:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="n">test</span> <span class="s2">"7.3 custom collate creates shifted targets"</span> <span class="k">do</span>
    <span class="n">inputs_1</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span>
    <span class="n">inputs_2</span> <span class="o">=</span> <span class="p">[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">]</span>
    <span class="n">inputs_3</span> <span class="o">=</span> <span class="p">[</span><span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">]</span>
    <span class="n">batch</span> <span class="o">=</span> <span class="p">{</span><span class="n">inputs_1</span><span class="p">,</span> <span class="n">inputs_2</span><span class="p">,</span> <span class="n">inputs_3</span><span class="p">}</span>

    <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">custom_collate</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">inputs</span> <span class="o">==</span>
             <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span>
               <span class="p">[</span>
                 <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">]</span>
               <span class="p">],</span>
               <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
             <span class="p">)</span>

    <span class="n">assert</span> <span class="n">targets</span> <span class="o">==</span>
             <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span>
               <span class="p">[</span>
                 <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">6</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">]</span>
               <span class="p">],</span>
               <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
             <span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Now we are at the final step. In the targets, we need to replace padding tokens with <code class="language-plaintext highlighter-rouge">-100</code>. We keep the last padding token in the targets so the model can learn when to end the response. First, here is the implementation, followed by an explanation of why we do this:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  diff --git a/lib/llm_scratch/instruction_dataset.ex b/lib/llm_scratch/instruction_dataset.ex
<span class="gh">index 4e91dfb..c161996 100644
</span><span class="gd">--- a/lib/llm_scratch/instruction_dataset.ex
</span><span class="gi">+++ b/lib/llm_scratch/instruction_dataset.ex
</span><span class="p">@@ -152,26 +152,32 @@</span> defmodule LlmScratch.InstructionDataset do
   @doc """
   Pads a batch of token-id sequences and creates input/target tensors.
 
<span class="gd">-  This mirrors the second draft of the chapter 7 Python collate function:
</span><span class="gi">+  This mirrors the chapter 7 Python collate function:
</span> 
       inputs_1 = [0, 1, 2, 3, 4]
       inputs_2 = [5, 6]
       inputs_3 = [7, 8, 9]
       batch = (inputs_1, inputs_2, inputs_3)
<span class="gd">-      inputs, targets = custom_collate(batch)
</span><span class="gi">+      inputs, targets = custom_collate_fn(batch)
</span>       print(inputs)
       print(targets)
 
   The function finds the longest sequence length in the batch after adding one
   extra pad token, pads every item to that length, then creates next-token
   prediction pairs. Inputs drop the final token from each padded row; targets
<span class="gd">-  drop the first token.
</span><span class="gi">+  drop the first token. In target rows, all padding tokens after the first one
+  are replaced with `ignore_index` so they do not contribute to the loss.
</span> 
   ## Input Parameters
 
     * `batch` - tuple or list of token-id lists.
     * `pad_token_id` - token id used for padding. Defaults to GPT-2's
       end-of-text token id, `50256`.
<span class="gi">+    * `ignore_index` - label value used for ignored target positions. Defaults
+      to `-100`.
+    * `allowed_max_length` - optional maximum row length after inputs and
+      targets are created. Defaults to `nil`, which keeps the full batch
+      length.
</span>     * `device` - included for parity with the Python example. Nx places the
       tensor on the active backend, so this argument is currently informational.
 
<span class="p">@@ -198,22 +204,36 @@</span> defmodule LlmScratch.InstructionDataset do
         s64[3][5]
         [
           [1, 2, 3, 4, 50256],
<span class="gd">-          [6, 50256, 50256, 50256, 50256],
-          [8, 9, 50256, 50256, 50256]
</span><span class="gi">+          [6, 50256, -100, -100, -100],
+          [8, 9, 50256, -100, -100]
</span>         ]
       &gt;
   """
<span class="gd">-  @spec custom_collate(tuple() | [[integer()]], integer(), String.t()) ::
</span><span class="gi">+  @spec custom_collate(
+          tuple() | [[integer()]],
+          integer(),
+          integer(),
+          nil | pos_integer(),
+          String.t()
+        ) ::
</span>           {Nx.Tensor.t(), Nx.Tensor.t()}
<span class="gd">-  def custom_collate(batch, pad_token_id \\ @pad_token_id, device \\ "cpu")
-
-  def custom_collate(batch, pad_token_id, device) when is_tuple(batch) do
</span><span class="gi">+  def custom_collate(
+        batch,
+        pad_token_id \\ @pad_token_id,
+        ignore_index \\ -100,
+        allowed_max_length \\ nil,
+        device \\ "cpu"
+      )
+
+  def custom_collate(batch, pad_token_id, ignore_index, allowed_max_length, device)
+      when is_tuple(batch) do
</span>     batch
     |&gt; Tuple.to_list()
<span class="gd">-    |&gt; custom_collate(pad_token_id, device)
</span><span class="gi">+    |&gt; custom_collate(pad_token_id, ignore_index, allowed_max_length, device)
</span>   end
 
<span class="gd">-  def custom_collate(batch, pad_token_id, _device) when is_list(batch) do
</span><span class="gi">+  def custom_collate(batch, pad_token_id, ignore_index, allowed_max_length, _device)
+      when is_list(batch) do
</span>     # Find the longest sequence length after the one extra pad token that the
     # Python collate function appends to every item.
     batch_max_length =
<span class="p">@@ -231,10 +251,18 @@</span> defmodule LlmScratch.InstructionDataset do
         padded = pad_to_length(new_item, batch_max_length, pad_token_id)
 
         # Match `padded[:-1]` for inputs and `padded[1:]` for next-token
<span class="gd">-        # targets before stacking all rows into tensors.
</span><span class="gi">+        # targets.
+        inputs = Enum.slice(padded, 0, batch_max_length - 1)
+
+        targets =
+          padded
+          |&gt; Enum.slice(1, batch_max_length - 1)
+          |&gt; mask_extra_padding_targets(pad_token_id, ignore_index)
+
+        # Optionally cap both rows to the model context length.
</span>         {
<span class="gd">-          Enum.slice(padded, 0, batch_max_length - 1),
-          Enum.slice(padded, 1, batch_max_length - 1)
</span><span class="gi">+          maybe_truncate(inputs, allowed_max_length),
+          maybe_truncate(targets, allowed_max_length)
</span>         }
       end)
       |&gt; Enum.unzip()
<span class="p">@@ -261,4 +289,23 @@</span> defmodule LlmScratch.InstructionDataset do
   defp pad_to_length(token_ids, max_length, pad_token_id) do
     token_ids ++ List.duplicate(pad_token_id, max_length - Kernel.length(token_ids))
   end
<span class="gi">+
+  defp mask_extra_padding_targets(targets, pad_token_id, ignore_index) do
+    {masked_targets, _seen_first_pad?} =
+      Enum.map_reduce(targets, false, fn
+        ^pad_token_id, false -&gt;
+          {pad_token_id, true}
+
+        ^pad_token_id, true -&gt;
+          {ignore_index, true}
+
+        token_id, seen_first_pad? -&gt;
+          {token_id, seen_first_pad?}
+      end)
+
+    masked_targets
+  end
+
+  defp maybe_truncate(token_ids, nil), do: token_ids
+  defp maybe_truncate(token_ids, max_length), do: Enum.take(token_ids, max_length)
</span> end
</code></pre></div></div>

<p>And here is a test as an example:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code> <span class="n">test</span> <span class="s2">"7.3 custom collate creates shifted targets with ignored padding"</span> <span class="k">do</span>
    <span class="n">inputs_1</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span>
    <span class="n">inputs_2</span> <span class="o">=</span> <span class="p">[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">]</span>
    <span class="n">inputs_3</span> <span class="o">=</span> <span class="p">[</span><span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">]</span>
    <span class="n">batch</span> <span class="o">=</span> <span class="p">{</span><span class="n">inputs_1</span><span class="p">,</span> <span class="n">inputs_2</span><span class="p">,</span> <span class="n">inputs_3</span><span class="p">}</span>

    <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">custom_collate</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">inputs</span> <span class="o">==</span>
             <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span>
               <span class="p">[</span>
                 <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">]</span>
               <span class="p">],</span>
               <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
             <span class="p">)</span>

    <span class="n">assert</span> <span class="n">targets</span> <span class="o">==</span>
             <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span>
               <span class="p">[</span>
                 <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">6</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="o">-</span><span class="mi">100</span><span class="p">,</span> <span class="o">-</span><span class="mi">100</span><span class="p">,</span> <span class="o">-</span><span class="mi">100</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">,</span> <span class="o">-</span><span class="mi">100</span><span class="p">,</span> <span class="o">-</span><span class="mi">100</span><span class="p">]</span>
               <span class="p">],</span>
               <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
             <span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>What is so special about the <code class="language-plaintext highlighter-rouge">-100</code> value? PyTorch’s cross entropy loss function ignores that value. We could also mask out all instruction tokens from the target, because some researchers claim this can improve training. The last thing to implement is a refactor of our <code class="language-plaintext highlighter-rouge">cross_entropy_loss</code> function so it ignores <code class="language-plaintext highlighter-rouge">-100</code>:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  diff --git a/lib/llm_scratch/loss_utils.ex b/lib/llm_scratch/loss_utils.ex
<span class="gh">index 59a9edc..70495f6 100644
</span><span class="gd">--- a/lib/llm_scratch/loss_utils.ex
</span><span class="gi">+++ b/lib/llm_scratch/loss_utils.ex
</span><span class="p">@@ -46,6 +46,9 @@</span> defmodule LlmScratch.LossUtils do
   `logits` should be shaped `{batch_size, seq_len, vocab_size}` and `targets`
   should be shaped `{batch_size, seq_len}`.
 
<span class="gi">+  Target positions equal to `ignore_index` are excluded from the mean loss.
+  The default `ignore_index` is `-100`, matching PyTorch's cross entropy loss.
+
</span>   ## Examples
 
       iex&gt; logits = Nx.log(Nx.tensor([[[0.5, 0.25, 0.25], [0.125, 0.375, 0.5]]]))
<span class="p">@@ -54,52 +57,69 @@</span> defmodule LlmScratch.LossUtils do
       iex&gt; Float.round(Nx.to_number(loss), 6)
       0.693147
   """
<span class="gd">-  @spec cross_entropy_loss(Nx.Tensor.t(), Nx.Tensor.t()) :: Nx.Tensor.t()
-  def cross_entropy_loss(%Nx.Tensor{} = logits, %Nx.Tensor{} = targets) do
-    cross_entropy_loss_defn(logits, targets)
</span><span class="gi">+  @spec cross_entropy_loss(Nx.Tensor.t(), Nx.Tensor.t(), integer()) :: Nx.Tensor.t()
+  def cross_entropy_loss(%Nx.Tensor{} = logits, %Nx.Tensor{} = targets, ignore_index \\ -100) do
+    cross_entropy_loss_defn(logits, targets, ignore_index)
</span>   end
 
   @doc """
   Defn-compatible mean cross entropy loss from logits and target token ids.
   """
<span class="gd">-  defn cross_entropy_loss_defn(logits, targets) do
</span><span class="gi">+  defn cross_entropy_loss_defn(logits, targets, ignore_index \\ -100) do
</span>     if Nx.rank(logits) == 3 do
       # Language-model flow:
       # logits shape is {batch_size, seq_len, vocab_size}
       # targets shape is {batch_size, seq_len}
       vocab_size = Nx.axis_size(logits, 2)
<span class="gi">+      valid_targets = Nx.not_equal(targets, ignore_index)
+      safe_targets = Nx.select(valid_targets, targets, 0)
</span> 
       # step1: logits
<span class="gd">-      logits
-      # step2: probabilities
-      |&gt; Axon.Activations.softmax(axis: -1)
-      # step3: target probabilities
-      |&gt; Nx.reshape({:auto, vocab_size})
-      |&gt; Nx.take_along_axis(Nx.reshape(targets, {:auto, 1}), axis: 1)
-      |&gt; Nx.squeeze(axes: [1])
-      # step4: logarithmic probabilities
-      |&gt; Nx.log()
-      # step6: negative average log probabilities
-      |&gt; Nx.negate()
-      # step5: average logarithmic probabilities
-      |&gt; Nx.mean()
</span><span class="gi">+      losses =
+        logits
+        # step2: probabilities
+        |&gt; Axon.Activations.softmax(axis: -1)
+        # step3: target probabilities
+        |&gt; Nx.reshape({:auto, vocab_size})
+        |&gt; Nx.take_along_axis(Nx.reshape(safe_targets, {:auto, 1}), axis: 1)
+        |&gt; Nx.squeeze(axes: [1])
+        # step4: logarithmic probabilities
+        |&gt; Nx.log()
+        # step6: negative log probabilities
+        |&gt; Nx.negate()
+
+      valid_targets = Nx.reshape(valid_targets, {:auto})
+
+      valid_targets
+      |&gt; Nx.select(losses, 0.0)
+      |&gt; Nx.sum()
+      # step5: average only over target positions that are not ignored
+      |&gt; Nx.divide(valid_targets |&gt; Nx.as_type({:f, 32}) |&gt; Nx.sum())
</span>     else
       # Classification flow:
       # logits shape is {batch_size, num_classes}
       # targets shape is {batch_size}
<span class="gi">+      valid_targets = Nx.not_equal(targets, ignore_index)
+      safe_targets = Nx.select(valid_targets, targets, 0)
+
</span>       # step1: logits
<span class="gd">-      logits
-      # step2: probabilities
-      |&gt; Axon.Activations.softmax(axis: -1)
-      # step3: target probabilities
-      |&gt; Nx.take_along_axis(Nx.reshape(targets, {:auto, 1}), axis: 1)
-      |&gt; Nx.squeeze(axes: [1])
-      # step4: logarithmic probabilities
-      |&gt; Nx.log()
-      # step6: negative average log probabilities
-      |&gt; Nx.negate()
-      # step5: average logarithmic probabilities
-      |&gt; Nx.mean()
</span><span class="gi">+      losses =
+        logits
+        # step2: probabilities
+        |&gt; Axon.Activations.softmax(axis: -1)
+        # step3: target probabilities
+        |&gt; Nx.take_along_axis(Nx.reshape(safe_targets, {:auto, 1}), axis: 1)
+        |&gt; Nx.squeeze(axes: [1])
+        # step4: logarithmic probabilities
+        |&gt; Nx.log()
+        # step6: negative log probabilities
+        |&gt; Nx.negate()
+
+      valid_targets
+      |&gt; Nx.select(losses, 0.0)
+      |&gt; Nx.sum()
+      # step5: average only over target positions that are not ignored
+      |&gt; Nx.divide(valid_targets |&gt; Nx.as_type({:f, 32}) |&gt; Nx.sum())
</span>     end
   end
 
<span class="p">@@ -128,6 +148,8 @@</span> defmodule LlmScratch.LossUtils do
       model output is sliced to `model(input_batch)[:, -1, :]` before computing
       cross entropy. Use `:first_token` to instead classify from
       `model(input_batch)[:, 0, :]`.
<span class="gi">+    * `:ignore_index` - target value excluded from the mean cross entropy.
+      Defaults to `-100`.
</span>   """
   @spec calc_loss_batch(
           Nx.Tensor.t(),
<span class="p">@@ -141,10 +163,12 @@</span> defmodule LlmScratch.LossUtils do
     input_batch = maybe_transfer(input_batch, device)
     target_batch = maybe_transfer(target_batch, device)
 
<span class="gi">+    ignore_index = Keyword.get(opts, :ignore_index, -100)
+
</span>     model
     |&gt; forward_model(input_batch)
     |&gt; select_loss_logits(opts)
<span class="gd">-    |&gt; cross_entropy_loss(target_batch)
</span><span class="gi">+    |&gt; cross_entropy_loss(target_batch, ignore_index)
</span>   end
 
   @doc """
</code></pre></div></div>

<p>And here is a test from the book that checks whether it really ignores that value:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"7.3 cross entropy ignores masked instruction targets"</span> <span class="k">do</span>
    <span class="n">logits_1</span> <span class="o">=</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span>
        <span class="p">[</span><span class="o">-</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">],</span>
        <span class="p">[</span><span class="o">-</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">1.5</span><span class="p">]</span>
      <span class="p">])</span>

    <span class="n">targets_1</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
    <span class="n">loss_1</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">cross_entropy_loss</span><span class="p">(</span><span class="n">logits_1</span><span class="p">,</span> <span class="n">targets_1</span><span class="p">)</span>

    <span class="n">logits_2</span> <span class="o">=</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span>
        <span class="p">[</span><span class="o">-</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">],</span>
        <span class="p">[</span><span class="o">-</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">1.5</span><span class="p">],</span>
        <span class="p">[</span><span class="o">-</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">1.5</span><span class="p">]</span>
      <span class="p">])</span>

    <span class="n">targets_2</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
    <span class="n">loss_2</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">cross_entropy_loss</span><span class="p">(</span><span class="n">logits_2</span><span class="p">,</span> <span class="n">targets_2</span><span class="p">)</span>

    <span class="n">targets_3</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">100</span><span class="p">],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
    <span class="n">loss_3</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">cross_entropy_loss</span><span class="p">(</span><span class="n">logits_2</span><span class="p">,</span> <span class="n">targets_3</span><span class="p">)</span>

    <span class="n">assert_in_delta</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">(</span><span class="n">loss_1</span><span class="p">),</span> <span class="mf">1.1269</span><span class="p">,</span> <span class="mf">1.0e-4</span>
    <span class="n">assert_in_delta</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">(</span><span class="n">loss_2</span><span class="p">),</span> <span class="mf">0.7936</span><span class="p">,</span> <span class="mf">1.0e-4</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">loss_1</span><span class="p">,</span> <span class="n">loss_3</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span> <span class="o">==</span> <span class="mi">1</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="74-creating-data-loaders-for-an-instruction-dataset">7.4 Creating data loaders for an instruction dataset</h2>

<p>We are ready to load our training data into memory using the <code class="language-plaintext highlighter-rouge">DataLoader</code> module. For that, we first need to refactor our <code class="language-plaintext highlighter-rouge">DataLoader</code> module to support a custom collate function:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="gh">diff --git a/lib/llm_scratch/data_loader.ex b/lib/llm_scratch/data_loader.ex
index 2e05c93..350c863 100644
</span><span class="gd">--- a/lib/llm_scratch/data_loader.ex
</span><span class="gi">+++ b/lib/llm_scratch/data_loader.ex
</span><span class="p">@@ -10,24 +10,30 @@</span> defmodule LlmScratch.DataLoader do
   - `:shuffle` - whether batches are shuffled
   - `:drop_last` - whether incomplete batches are dropped
   - `:num_workers` - concurrency used by `iterate/2`
<span class="gi">+  - `:collate_fn` - function used to transform each list of examples into a
+    batch. Defaults to keeping the list unchanged.
</span> 
   ## Options
   - `:batch_size` - number of samples per batch (default: `32`)
   - `:shuffle` - shuffles dataset once before cycling (default: `true`)
   - `:drop_last` - drops batches smaller than `:batch_size` (default: `false`)
   - `:num_workers` - parallel workers for iteration (default: `0`)
<span class="gi">+  - `:collate_fn` - one-argument function called for each batch after
+    chunking and `:drop_last` filtering. Defaults to identity.
</span>   """
   def new(dataset, opts \\ []) when is_list(dataset) do
     batch_size = Keyword.get(opts, :batch_size, 32)
     shuffle = Keyword.get(opts, :shuffle, true)
     drop_last = Keyword.get(opts, :drop_last, false)
     num_workers = Keyword.get(opts, :num_workers, 0)
<span class="gi">+    collate_fn = Keyword.get(opts, :collate_fn, &amp;identity/1)
</span> 
     batches =
       dataset
       |&gt; prepare_dataset(shuffle)
       |&gt; Stream.chunk_every(batch_size)
       |&gt; filter_incomplete_batches(drop_last, batch_size)
<span class="gi">+      |&gt; Stream.map(collate_fn)
</span>       |&gt; Enum.to_list()
 
     stream = Stream.cycle(batches)
<span class="p">@@ -39,7 +45,8 @@</span> defmodule LlmScratch.DataLoader do
       batch_size: batch_size,
       shuffle: shuffle,
       drop_last: drop_last,
<span class="gd">-      num_workers: num_workers
</span><span class="gi">+      num_workers: num_workers,
+      collate_fn: collate_fn
</span>     }
   end
 
<span class="p">@@ -76,4 +83,6 @@</span> defmodule LlmScratch.DataLoader do
   defp filter_incomplete_batches(stream, false, _batch_size) do
     stream
   end
<span class="gi">+
+  defp identity(batch), do: batch
</span> end
</code></pre></div></div>

<p>And here is a test that loads instruction testing data:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code> <span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="n">test</span> <span class="s2">"7.4 creates instruction data loaders with custom collate function"</span> <span class="k">do</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>
    <span class="n">batch_size</span> <span class="o">=</span> <span class="mi">8</span>
    <span class="n">num_workers</span> <span class="o">=</span> <span class="mi">3</span>
    <span class="ss">:rand</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="ss">:exsss</span><span class="p">,</span> <span class="p">{</span><span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">})</span>

    <span class="n">file_path</span> <span class="o">=</span>
      <span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"llm_scratch_instruction_data"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"instruction-data.json"</span><span class="p">)</span>

    <span class="n">data</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">download_and_load_instructions_file</span><span class="p">(</span>
        <span class="n">file_path</span><span class="p">,</span>
        <span class="nv">@instruction_data_url</span>
      <span class="p">)</span>

    <span class="n">train_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.85</span><span class="p">)</span>
    <span class="n">test_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">val_portion</span> <span class="o">=</span> <span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">-</span> <span class="n">train_portion</span> <span class="o">-</span> <span class="n">test_portion</span>

    <span class="n">train_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">train_portion</span><span class="p">)</span>
    <span class="n">test_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span><span class="p">,</span> <span class="n">test_portion</span><span class="p">)</span>
    <span class="n">val_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span> <span class="o">+</span> <span class="n">test_portion</span><span class="p">,</span> <span class="n">val_portion</span><span class="p">)</span>

    <span class="n">train_dataset</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">train_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>
    <span class="n">val_dataset</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">val_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>
    <span class="n">test_dataset</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">test_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">train_dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">collate_fn:</span> <span class="o">&amp;</span><span class="no">InstructionDataset</span><span class="o">.</span><span class="n">custom_collate</span><span class="o">/</span><span class="mi">1</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span>
      <span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">val_dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">collate_fn:</span> <span class="o">&amp;</span><span class="no">InstructionDataset</span><span class="o">.</span><span class="n">custom_collate</span><span class="o">/</span><span class="mi">1</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span>
      <span class="p">)</span>

    <span class="n">test_loader</span> <span class="o">=</span>
      <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">test_dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">collate_fn:</span> <span class="o">&amp;</span><span class="no">InstructionDataset</span><span class="o">.</span><span class="n">custom_collate</span><span class="o">/</span><span class="mi">1</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">train_dataset</span><span class="p">)</span> <span class="o">==</span> <span class="mi">935</span>
    <span class="n">assert</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">val_dataset</span><span class="p">)</span> <span class="o">==</span> <span class="mi">55</span>
    <span class="n">assert</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">test_dataset</span><span class="p">)</span> <span class="o">==</span> <span class="mi">110</span>

    <span class="n">assert</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">==</span> <span class="mi">116</span>
    <span class="n">assert</span> <span class="n">val_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">==</span> <span class="mi">7</span>
    <span class="n">assert</span> <span class="n">test_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">==</span> <span class="mi">14</span>
    <span class="n">assert</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">num_workers</span> <span class="o">==</span> <span class="mi">3</span>
    <span class="n">assert</span> <span class="n">val_loader</span><span class="o">.</span><span class="n">num_workers</span> <span class="o">==</span> <span class="mi">3</span>
    <span class="n">assert</span> <span class="n">test_loader</span><span class="o">.</span><span class="n">num_workers</span> <span class="o">==</span> <span class="mi">3</span>

    <span class="n">assert</span> <span class="p">[{</span><span class="n">train_inputs</span><span class="p">,</span> <span class="n">train_targets</span><span class="p">}</span> <span class="o">|</span> <span class="n">_</span><span class="p">]</span> <span class="o">=</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">batches</span>
    <span class="n">assert</span> <span class="p">[{</span><span class="n">val_inputs</span><span class="p">,</span> <span class="n">val_targets</span><span class="p">}</span> <span class="o">|</span> <span class="n">_</span><span class="p">]</span> <span class="o">=</span> <span class="n">val_loader</span><span class="o">.</span><span class="n">batches</span>
    <span class="n">assert</span> <span class="p">[{</span><span class="n">test_inputs</span><span class="p">,</span> <span class="n">test_targets</span><span class="p">}</span> <span class="o">|</span> <span class="n">_</span><span class="p">]</span> <span class="o">=</span> <span class="n">test_loader</span><span class="o">.</span><span class="n">batches</span>


    <span class="c1"># As we randomly shuffle data before creating batches, we get a different result than in the book</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">train_inputs</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">8</span><span class="p">,</span> <span class="mi">91</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">train_targets</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">8</span><span class="p">,</span> <span class="mi">91</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">val_inputs</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">8</span><span class="p">,</span> <span class="mi">74</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">val_targets</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">8</span><span class="p">,</span> <span class="mi">74</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">test_inputs</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">8</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">test_targets</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">8</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="75-loading-a-pretrained-llm">7.5 Loading a pretrained LLM</h2>

<p>We are going to load the OpenAI GPT model with 355M parameters. We need a stronger model for instruction training. After we load the model, we will try the first validation instruction from the validation set as proof that the 355M model does not understand instructions yet, because it has not been trained for that task:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">3_600_000</span>
  <span class="n">test</span> <span class="s2">"7.5 generates response for validation instruction with OpenAI GPT-2 355M"</span> <span class="k">do</span>
    <span class="n">use_accelerated_backend</span><span class="p">()</span>

    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="s2">"355M"</span><span class="p">,</span> <span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>

    <span class="n">file_path</span> <span class="o">=</span>
      <span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"llm_scratch_instruction_data"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"instruction-data.json"</span><span class="p">)</span>

    <span class="n">data</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">download_and_load_instructions_file</span><span class="p">(</span>
        <span class="n">file_path</span><span class="p">,</span>
        <span class="nv">@instruction_data_url</span>
      <span class="p">)</span>

    <span class="n">train_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.85</span><span class="p">)</span>
    <span class="n">test_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">val_portion</span> <span class="o">=</span> <span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">-</span> <span class="n">train_portion</span> <span class="o">-</span> <span class="n">test_portion</span>
    <span class="n">val_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span> <span class="o">+</span> <span class="n">test_portion</span><span class="p">,</span> <span class="n">val_portion</span><span class="p">)</span>

    <span class="n">input_text</span> <span class="o">=</span>
      <span class="n">val_data</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">()</span>

    <span class="n">token_ids</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">input_text</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">),</span>
        <span class="mi">35</span><span class="p">,</span>
        <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="mf">0.0</span><span class="p">,</span>
        <span class="no">nil</span><span class="p">,</span>
        <span class="mi">50_256</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">generated_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">response_text</span> <span class="o">=</span>
      <span class="n">generated_text</span> <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">input_text</span><span class="p">)</span><span class="o">..-</span><span class="mi">1</span><span class="o">//</span><span class="mi">1</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">trim</span><span class="p">()</span>

    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span> <span class="o">==</span> <span class="mi">1024</span>
    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">==</span> <span class="mi">1024</span>
    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">==</span> <span class="mi">24</span>
    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">==</span> <span class="mi">16</span>
    <span class="n">assert</span> <span class="n">input_text</span> <span class="o">==</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">(</span><span class="n">val_data</span><span class="p">))</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">starts_with?</span><span class="p">(</span><span class="n">generated_text</span><span class="p">,</span> <span class="n">input_text</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">response_text</span> <span class="o">==</span>
             <span class="s2">"### Response:</span><span class="se">\n\n</span><span class="s2">The chef cooks the meal every day.</span><span class="se">\n\n</span><span class="s2">### Instruction:</span><span class="se">\n\n</span><span class="s2">Convert the active sentence to passive: 'The chef cooks the"</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="76-fine-tuning-the-llm-on-instruction-data">7.6 Fine-tuning the LLM on instruction data</h2>

<p>It is time to fine-tune our GPT-2 355M model on the provided instruction set. We use the same training function as we used for model pretraining. First, let’s calculate the loss on a model that has not been fine-tuned on the instruction set. We can see that the loss is high, but after just two epochs the model starts to understand instructions. At the end, we save a model checkpoint and the data that will be used for plotting the loss values. Training took 687 seconds:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"7.5 and 7.6 evaluates and trains OpenAI GPT-2 355M on instruction data"</span> <span class="k">do</span>
    <span class="n">device</span> <span class="o">=</span> <span class="n">use_accelerated_backend</span><span class="p">()</span>

    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>
    <span class="n">batch_size</span> <span class="o">=</span> <span class="mi">8</span>
    <span class="n">num_workers</span> <span class="o">=</span> <span class="mi">3</span>

    <span class="n">model</span> <span class="o">=</span>
      <span class="s2">"355M"</span>
      <span class="o">|&gt;</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="n">file_path</span> <span class="o">=</span>
      <span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"llm_scratch_instruction_data"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"instruction-data.json"</span><span class="p">)</span>

    <span class="n">data</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">download_and_load_instructions_file</span><span class="p">(</span>
        <span class="n">file_path</span><span class="p">,</span>
        <span class="nv">@instruction_data_url</span>
      <span class="p">)</span>

    <span class="n">train_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.85</span><span class="p">)</span>
    <span class="n">test_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">val_portion</span> <span class="o">=</span> <span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">-</span> <span class="n">train_portion</span> <span class="o">-</span> <span class="n">test_portion</span>
    <span class="n">train_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">train_portion</span><span class="p">)</span>
    <span class="n">val_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span> <span class="o">+</span> <span class="n">test_portion</span><span class="p">,</span> <span class="n">val_portion</span><span class="p">)</span>

    <span class="n">train_dataset</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">train_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>
    <span class="n">val_dataset</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">val_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

    <span class="ss">:rand</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="ss">:exsss</span><span class="p">,</span> <span class="p">{</span><span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">})</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">train_dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">collate_fn:</span> <span class="o">&amp;</span><span class="n">binary_instruction_collate</span><span class="o">/</span><span class="mi">1</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span>
      <span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">val_dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">collate_fn:</span> <span class="o">&amp;</span><span class="n">binary_instruction_collate</span><span class="o">/</span><span class="mi">1</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span>
      <span class="p">)</span>

    <span class="n">train_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">5</span><span class="p">)</span>
    <span class="n">val_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">val_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">5</span><span class="p">)</span>

    <span class="n">input_text</span> <span class="o">=</span>
      <span class="n">val_data</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">()</span>

    <span class="n">token_ids</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">input_text</span>
        <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">device</span><span class="p">),</span>
        <span class="mi">35</span><span class="p">,</span>
        <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="mf">0.0</span><span class="p">,</span>
        <span class="no">nil</span><span class="p">,</span>
        <span class="mi">50_256</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">generated_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">response_text</span> <span class="o">=</span>
      <span class="n">generated_text</span> <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">input_text</span><span class="p">)</span><span class="o">..-</span><span class="mi">1</span><span class="o">//</span><span class="mi">1</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">trim</span><span class="p">()</span>

    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span> <span class="o">==</span> <span class="mi">1024</span>
    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">==</span> <span class="mi">1024</span>
    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">==</span> <span class="mi">24</span>
    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">==</span> <span class="mi">16</span>
    <span class="n">assert_in_delta</span> <span class="n">train_loss</span><span class="p">,</span> <span class="mf">3.7422078609466554</span><span class="p">,</span> <span class="mf">1.0e-5</span>
    <span class="n">assert_in_delta</span> <span class="n">val_loss</span><span class="p">,</span> <span class="mf">3.7619348049163817</span><span class="p">,</span> <span class="mf">1.0e-5</span>
    <span class="n">assert</span> <span class="n">input_text</span> <span class="o">==</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">(</span><span class="n">val_data</span><span class="p">))</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">starts_with?</span><span class="p">(</span><span class="n">generated_text</span><span class="p">,</span> <span class="n">input_text</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">response_text</span> <span class="o">==</span>
             <span class="s2">"### Response:</span><span class="se">\n\n</span><span class="s2">The chef cooks the meal every day.</span><span class="se">\n\n</span><span class="s2">### Instruction:</span><span class="se">\n\n</span><span class="s2">Convert the active sentence to passive: 'The chef cooks the"</span>

    <span class="n">num_epochs</span> <span class="o">=</span> <span class="mi">2</span>
    <span class="n">optimizer</span> <span class="o">=</span> <span class="no">Training</span><span class="o">.</span><span class="n">adamw</span><span class="p">(</span><span class="mf">0.00005</span><span class="p">,</span> <span class="ss">weight_decay:</span> <span class="mf">0.1</span><span class="p">)</span>

    <span class="p">{</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">train_losses</span><span class="p">,</span> <span class="n">val_losses</span><span class="p">,</span> <span class="n">tokens_seen</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Training</span><span class="o">.</span><span class="n">train_model_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">train_loader</span><span class="p">,</span>
        <span class="n">val_loader</span><span class="p">,</span>
        <span class="n">optimizer</span><span class="p">,</span>
        <span class="n">device</span><span class="p">,</span>
        <span class="n">num_epochs</span><span class="p">,</span>
        <span class="mi">5</span><span class="p">,</span>
        <span class="mi">5</span><span class="p">,</span>
        <span class="n">input_text</span><span class="p">,</span>
        <span class="n">tokenizer</span><span class="p">,</span>
        <span class="ss">generate_samples:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">return_optimizer:</span> <span class="no">true</span>
      <span class="p">)</span>

    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="s2">"ch7_instruction_finetuned_gpt2_355m_model_and_optimizer.nx"</span>
    <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">save_training_state!</span><span class="p">(</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="n">metrics_path</span> <span class="o">=</span> <span class="s2">"ch7_instruction_finetuning_metrics.json"</span>

    <span class="n">write_instruction_training_metrics!</span><span class="p">(</span>
      <span class="n">metrics_path</span><span class="p">,</span>
      <span class="n">num_epochs</span><span class="p">,</span>
      <span class="n">train_losses</span><span class="p">,</span>
      <span class="n">val_losses</span><span class="p">,</span>
      <span class="n">tokens_seen</span>
    <span class="p">)</span>

    <span class="n">assert</span> <span class="n">trained_model</span><span class="o">.</span><span class="n">__struct__</span> <span class="o">==</span> <span class="n">model</span><span class="o">.</span><span class="n">__struct__</span>
    <span class="n">assert</span> <span class="p">%</span><span class="no">Training</span><span class="o">.</span><span class="no">AdamW</span><span class="p">{}</span> <span class="o">=</span> <span class="n">trained_optimizer</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">)</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_losses</span><span class="p">)</span> <span class="o">==</span> <span class="mi">47</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">val_losses</span><span class="p">)</span> <span class="o">==</span> <span class="mi">47</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">)</span> <span class="o">==</span> <span class="mi">47</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">train_losses</span> <span class="o">++</span> <span class="n">val_losses</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_float</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_integer</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="n">tokens_seen</span> <span class="o">==</span> <span class="no">Enum</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">List</span><span class="o">.</span><span class="n">last</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">)</span> <span class="o">&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">binary_instruction_collate</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>

    <span class="k">try</span> <span class="k">do</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>
      <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">custom_collate</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>
    <span class="k">after</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">write_instruction_training_metrics!</span><span class="p">(</span>
         <span class="n">path</span><span class="p">,</span>
         <span class="n">num_epochs</span><span class="p">,</span>
         <span class="n">train_losses</span><span class="p">,</span>
         <span class="n">val_losses</span><span class="p">,</span>
         <span class="n">tokens_seen</span>
       <span class="p">)</span> <span class="k">do</span>
    <span class="n">metrics</span> <span class="o">=</span> <span class="p">%{</span>
      <span class="ss">num_epochs:</span> <span class="n">num_epochs</span><span class="p">,</span>
      <span class="ss">losses:</span> <span class="p">%{</span>
        <span class="ss">tokens_seen:</span> <span class="n">tokens_seen</span><span class="p">,</span>
        <span class="ss">train_values:</span> <span class="n">train_losses</span><span class="p">,</span>
        <span class="ss">val_values:</span> <span class="n">val_losses</span>
      <span class="p">}</span>
    <span class="p">}</span>

    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">encoded_metrics</span><span class="p">}</span> <span class="o">=</span> <span class="no">Jason</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">metrics</span><span class="p">,</span> <span class="ss">pretty:</span> <span class="no">true</span><span class="p">)</span>
    <span class="no">File</span><span class="o">.</span><span class="n">write!</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="n">encoded_metrics</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>
<p>Here is the size of the saved instruction fine-tuned model:</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  ls -alh *.nx
-rw-r--r--  1 karlosmid  staff   623M  5 lip  13:01 ch5_2_gpt_124m_emlx.nx
-rw-r--r--  1 karlosmid  staff   623M 10 lip  15:01 ch5_2_gpt_124m.nx
-rw-r--r--  1 karlosmid  staff   577M 10 lip  15:04 ch6_spam_classifier_model_and_optimizer.nx
-rw-r--r--  1 karlosmid  staff   4,6G 11 lip  13:37 ch7_instruction_finetuned_gpt2_355m_model_and_optimizer.nx
</code></pre></div></div>

<p>And here is the training output across epoch steps:</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>Ep 1 (Step 000000): Train loss 2.510, Val loss 2.540
Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Convert the active sentence to passive: 'The chef cooks the meal every day.'

### Response:
The chef cooks the meal every day.

### Instruction:
Write a response that appropriately completes the request.

### Response:
The chef cooks the meal every day.

### Instruction:
Write
Ep 1 (Step 000005): Train loss 1.093, Val loss 1.079
Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Convert the active sentence to passive: 'The chef cooks the meal every day.'

### Response:
The chef cooks the meal every day.&lt;|endoftext|&gt;The following is a list of items that make good gifts for Tia, who likes Modern, white items and dislikes Basic items.

Best Gifts for Tia:
Ep 1 (Step 000010): Train loss 0.960, Val loss 0.958
Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Convert the active sentence to passive: 'The chef cooks the meal every day.'

### Response:
The chef cooks the meal every day.&lt;|endoftext|&gt;The following is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
What is the difference between a 'good' and
Ep 1 (Step 000015): Train loss 0.913, Val loss 0.894
Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Convert the active sentence to passive: 'The chef cooks the meal every day.'

### Response:
The chef cooks the meal every day.&lt;|endoftext|&gt;The following is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Generate a sentence using the word 'excellent
Ep 1 (Step 000020): Train loss 0.878, Val loss 0.850
Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Convert the active sentence to passive: 'The chef cooks the meal every day.'

### Response:
The chef cooks the meal every day.&lt;|endoftext|&gt;The following is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
What is the opposite of 'happy'?

......

Ep 2 (Step 000220): Train loss 0.354, Val loss 0.645
Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Convert the active sentence to passive: 'The chef cooks the meal every day.'

### Response:
The meal is prepared by the chef every day.&lt;|endoftext|&gt;The following is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Translate 'I love you' into
Ep 2 (Step 000225): Train loss 0.343, Val loss 0.644
Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Convert the active sentence to passive: 'The chef cooks the meal every day.'

### Response:
The meal is prepared every day by the chef.&lt;|endoftext|&gt;The following is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
What is the capital of the United Kingdom
Ep 2 (Step 000230): Train loss 0.339, Val loss 0.634
Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Convert the active sentence to passive: 'The chef cooks the meal every day.'

### Response:
The meal is prepared every day by the chef.&lt;|endoftext|&gt;The following is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
What is the opposite of 'cold'?
.
Finished in 687.2 seconds (0.00s async, 687.2s sync)
</code></pre></div></div>

<p>And plot of training and validation losses over two epochs:</p>

<p><img src="/assets/2026/06/fine-tuning-training-validation-loss.png" alt="Training and validation loss over two epochs" /></p>

<p><em>Training and validation loss over two epochs.</em></p>

<h2 id="77-extracting-and-saving-responses">7.7 Extracting and saving responses</h2>

<p>It is time to evaluate how good our instruction fine-tuned LLM model is. We will first extract model responses for our test input set. The first approach is to print out the first three test data entries:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"7.7 evaluates fine-tuned instruction model on first test samples"</span> <span class="k">do</span>
    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="s2">"ch7_instruction_finetuned_gpt2_355m_model_and_optimizer.nx"</span>

    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="n">device</span> <span class="o">=</span> <span class="n">use_accelerated_backend</span><span class="p">()</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>
    <span class="ss">:rand</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="ss">:exsss</span><span class="p">,</span> <span class="p">{</span><span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">})</span>

    <span class="p">%{</span><span class="ss">model_state_dict:</span> <span class="n">model</span><span class="p">}</span> <span class="o">=</span> <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">load_training_state!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

    <span class="n">file_path</span> <span class="o">=</span>
      <span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"llm_scratch_instruction_data"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"instruction-data.json"</span><span class="p">)</span>

    <span class="n">data</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">download_and_load_instructions_file</span><span class="p">(</span>
        <span class="n">file_path</span><span class="p">,</span>
        <span class="nv">@instruction_data_url</span>
      <span class="p">)</span>

    <span class="n">train_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.85</span><span class="p">)</span>
    <span class="n">test_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">test_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span><span class="p">,</span> <span class="n">test_portion</span><span class="p">)</span>

    <span class="n">expected_samples</span> <span class="o">=</span> <span class="p">[</span>
      <span class="p">%{</span>
        <span class="ss">input_text:</span>
          <span class="s2">"Below is an instruction that describes a task. "</span> <span class="o">&lt;&gt;</span>
            <span class="s2">"Write a response that appropriately completes the request."</span> <span class="o">&lt;&gt;</span>
            <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Instruction:</span><span class="se">\n</span><span class="s2">Rewrite the sentence using a simile."</span> <span class="o">&lt;&gt;</span>
            <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Input:</span><span class="se">\n</span><span class="s2">The car is very fast."</span><span class="p">,</span>
        <span class="ss">correct_response:</span> <span class="s2">"The car is as fast as lightning."</span><span class="p">,</span>
        <span class="ss">model_response:</span> <span class="s2">"The car is as fast as a cheetah."</span>
      <span class="p">},</span>
      <span class="p">%{</span>
        <span class="ss">input_text:</span>
          <span class="s2">"Below is an instruction that describes a task. "</span> <span class="o">&lt;&gt;</span>
            <span class="s2">"Write a response that appropriately completes the request."</span> <span class="o">&lt;&gt;</span>
            <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Instruction:</span><span class="se">\n</span><span class="s2">What type of cloud is typically associated with thunderstorms?"</span><span class="p">,</span>
        <span class="ss">correct_response:</span>
          <span class="s2">"The type of cloud typically associated with thunderstorms is cumulonimbus."</span><span class="p">,</span>
        <span class="ss">model_response:</span>
          <span class="s2">"A thunderstorm is a type of cloud that typically forms when thunderstorms produce a dense, convective layer of air that is at least 10 miles thick."</span>
      <span class="p">},</span>
      <span class="p">%{</span>
        <span class="ss">input_text:</span>
          <span class="s2">"Below is an instruction that describes a task. "</span> <span class="o">&lt;&gt;</span>
            <span class="s2">"Write a response that appropriately completes the request."</span> <span class="o">&lt;&gt;</span>
            <span class="s2">"</span><span class="se">\n\n</span><span class="s2">### Instruction:</span><span class="se">\n</span><span class="s2">Name the author of 'Pride and Prejudice'."</span><span class="p">,</span>
        <span class="ss">correct_response:</span> <span class="s2">"Jane Austen."</span><span class="p">,</span>
        <span class="ss">model_response:</span> <span class="s2">"The author of 'Pride and Prejudice' is Jane Austen."</span>
      <span class="p">}</span>
    <span class="p">]</span>

    <span class="n">test_data</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">(</span><span class="n">expected_samples</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">each</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">entry</span><span class="p">,</span> <span class="n">expected</span><span class="p">}</span> <span class="o">-&gt;</span>
      <span class="n">input_text</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="n">entry</span><span class="p">)</span>

      <span class="n">token_ids</span> <span class="o">=</span>
        <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
          <span class="n">model</span><span class="p">,</span>
          <span class="n">input_text</span>
          <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>
          <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">device</span><span class="p">),</span>
          <span class="mi">256</span><span class="p">,</span>
          <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
          <span class="mf">0.0</span><span class="p">,</span>
          <span class="no">nil</span><span class="p">,</span>
          <span class="mi">50_256</span>
        <span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

      <span class="n">generated_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

      <span class="n">response_text</span> <span class="o">=</span>
        <span class="n">generated_text</span>
        <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">input_text</span><span class="p">)</span><span class="o">..-</span><span class="mi">1</span><span class="o">//</span><span class="mi">1</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">"### Response:"</span><span class="p">,</span> <span class="s2">""</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">trim</span><span class="p">()</span>

      <span class="n">assert</span> <span class="n">input_text</span> <span class="o">==</span> <span class="n">expected</span><span class="o">.</span><span class="n">input_text</span>
      <span class="n">assert</span> <span class="n">entry</span><span class="p">[</span><span class="s2">"output"</span><span class="p">]</span> <span class="o">==</span> <span class="n">expected</span><span class="o">.</span><span class="n">correct_response</span>
      <span class="n">assert</span> <span class="n">response_text</span> <span class="o">==</span> <span class="n">expected</span><span class="o">.</span><span class="n">model_response</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>By our own judgment, we can see that two responses are correct, while the cloud answer is wrong. We have three evaluation options:</p>

<ul>
  <li>MMLU uses short-answer and multiple-choice benchmarks</li>
  <li>LMSYS uses human preference comparisons with other LLMs</li>
  <li>automated conversational benchmarks, where we use another LLM, e.g. GPT-5.</li>
</ul>

<p>Since this chapter focuses on chat conversation, we will not use the MMLU option. We will use the third option: another LLM for evaluating our fine-tuned instruction model. For that, we need a small helper module that generates responses from a defined set and saves them into a file:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
  <span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">InstructionsEvaluation</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Helpers for evaluating an instruction fine-tuned model on instruction records.
  """</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span><span class="no">FineTuneDataLoader</span><span class="p">,</span> <span class="no">TextGeneration</span><span class="p">,</span> <span class="no">TextUtils</span><span class="p">}</span>

  <span class="nv">@default_output_path</span> <span class="s2">"instruction-data-with-response.json"</span>
  <span class="nv">@default_max_new_tokens</span> <span class="mi">256</span>
  <span class="nv">@default_eos_id</span> <span class="mi">50_256</span>

  <span class="nv">@type</span> <span class="n">instruction_record</span> <span class="p">::</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">instruction_record</span><span class="p">()</span>
  <span class="nv">@type</span> <span class="n">device</span> <span class="p">::</span> <span class="no">nil</span> <span class="o">|</span> <span class="ss">:default</span> <span class="o">|</span> <span class="n">atom</span><span class="p">()</span> <span class="o">|</span> <span class="n">tuple</span><span class="p">()</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Generates model responses for instruction records and writes them as JSON.

  Each input record is preserved and extended with a `"model_response"` key.
  The output mirrors the chapter 7 Python loop that writes
  `instruction-data-with-response.json`.

  ## Options

    * `:output_path` - JSON destination. Defaults to
      `"instruction-data-with-response.json"`.
    * `:max_new_tokens` - maximum generated tokens per record. Defaults to
      `256`.
    * `:context_size` - context length passed to `TextGeneration.generate/7`.
      Defaults to `model.cfg.context_length`.
    * `:eos_id` - end-of-text token id. Defaults to `50256`.

  Returns the enriched records after writing the file.

  ## Example

      test_data = [
        %{
          "instruction" =&gt; "Say hi.",
          "input" =&gt; "",
          "output" =&gt; "Hi."
        }
      ]

      LlmScratch.InstructionsEvaluation.write_responses!(
        test_data,
        model,
        "code-davinci-002",
        device,
        output_path: "instruction-data-with-response.json"
      )

  Example output:

      [
        %{
          "input" =&gt; "",
          "instruction" =&gt; "Say hi.",
          "model_response" =&gt; "Hi.",
          "output" =&gt; "Hi."
        }
      ]
  """</span>
  <span class="nv">@spec</span> <span class="n">write_responses!</span><span class="p">([</span><span class="n">instruction_record</span><span class="p">()],</span> <span class="n">struct</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">device</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span>
          <span class="p">[</span><span class="n">instruction_record</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">write_responses!</span><span class="p">(</span><span class="n">test_data</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span>
      <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">test_data</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">output_path</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:output_path</span><span class="p">,</span> <span class="nv">@default_output_path</span><span class="p">)</span>

    <span class="n">enriched_data</span> <span class="o">=</span> <span class="n">generate_responses</span><span class="p">(</span><span class="n">test_data</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="n">output_path</span>
    <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">dirname</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">mkdir_p!</span><span class="p">()</span>

    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">encoded</span><span class="p">}</span> <span class="o">=</span> <span class="no">Jason</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">enriched_data</span><span class="p">,</span> <span class="ss">pretty:</span> <span class="no">true</span><span class="p">)</span>
    <span class="no">File</span><span class="o">.</span><span class="n">write!</span><span class="p">(</span><span class="n">output_path</span><span class="p">,</span> <span class="n">encoded</span><span class="p">)</span>

    <span class="n">enriched_data</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Generates model responses for instruction records.

  Returns the original records with an added `"model_response"` field.

  ## Example

      test_data = [
        %{
          "instruction" =&gt; "Name one color.",
          "input" =&gt; "",
          "output" =&gt; "Blue."
        }
      ]

      LlmScratch.InstructionsEvaluation.generate_responses(
        test_data,
        model,
        "code-davinci-002",
        device
      )

  Example output:

      [
        %{
          "input" =&gt; "",
          "instruction" =&gt; "Name one color.",
          "model_response" =&gt; "Blue.",
          "output" =&gt; "Blue."
        }
      ]
  """</span>
  <span class="nv">@spec</span> <span class="n">generate_responses</span><span class="p">([</span><span class="n">instruction_record</span><span class="p">()],</span> <span class="n">struct</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">device</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span>
          <span class="p">[</span><span class="n">instruction_record</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">generate_responses</span><span class="p">(</span><span class="n">test_data</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span>
      <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">test_data</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">max_new_tokens</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:max_new_tokens</span><span class="p">,</span> <span class="nv">@default_max_new_tokens</span><span class="p">)</span>
    <span class="n">context_size</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get_lazy</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:context_size</span><span class="p">,</span> <span class="k">fn</span> <span class="o">-&gt;</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span> <span class="k">end</span><span class="p">)</span>
    <span class="n">eos_id</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:eos_id</span><span class="p">,</span> <span class="nv">@default_eos_id</span><span class="p">)</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">with_index</span><span class="p">(</span><span class="n">test_data</span><span class="p">,</span> <span class="k">fn</span> <span class="n">entry</span><span class="p">,</span> <span class="n">index</span> <span class="o">-&gt;</span>
      <span class="k">if</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:progress</span><span class="p">,</span> <span class="no">false</span><span class="p">)</span> <span class="k">do</span>
        <span class="no">IO</span><span class="o">.</span><span class="n">puts</span><span class="p">(</span><span class="s2">"Generating response </span><span class="si">#{</span><span class="n">index</span> <span class="o">+</span> <span class="mi">1</span><span class="si">}</span><span class="s2">/</span><span class="si">#{</span><span class="n">length</span><span class="p">(</span><span class="n">test_data</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
      <span class="k">end</span>

      <span class="no">Map</span><span class="o">.</span><span class="n">put</span><span class="p">(</span>
        <span class="n">entry</span><span class="p">,</span>
        <span class="s2">"model_response"</span><span class="p">,</span>
        <span class="n">generate_response</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">entry</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">max_new_tokens</span><span class="p">,</span> <span class="n">context_size</span><span class="p">,</span> <span class="n">eos_id</span><span class="p">)</span>
      <span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">generate_response</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">entry</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">max_new_tokens</span><span class="p">,</span> <span class="n">context_size</span><span class="p">,</span> <span class="n">eos_id</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">input_text</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="n">entry</span><span class="p">)</span>

    <span class="n">token_ids</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">input_text</span>
        <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">device</span><span class="p">),</span>
        <span class="n">max_new_tokens</span><span class="p">,</span>
        <span class="n">context_size</span><span class="p">,</span>
        <span class="mf">0.0</span><span class="p">,</span>
        <span class="no">nil</span><span class="p">,</span>
        <span class="n">eos_id</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">token_ids</span>
    <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">input_text</span><span class="p">)</span><span class="o">..-</span><span class="mi">1</span><span class="o">//</span><span class="mi">1</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">"### Response:"</span><span class="p">,</span> <span class="s2">""</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">trim</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">tensor</span>
  <span class="k">defp</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="ss">:default</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">tensor</span>
  <span class="k">defp</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>
<span class="k">end</span>
</code></pre></div></div>

<h2 id="78-evaluating-the-fine-tuned-llm">7.8 Evaluating the fine-tuned LLM</h2>

<p>We will evaluate our model using <a href="https://ollama.com/">Ollama</a>, so you first need to install it and start it before running the tests. Here is the Ollama helper that will be used in the tests:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">OllamaUtils</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Utilities for working with a local Ollama runtime.
  """</span>

  <span class="nv">@llama_server_process</span> <span class="s2">"ollama"</span>
  <span class="nv">@default_model</span> <span class="s2">"llama3"</span>
  <span class="nv">@default_chat_url</span> <span class="s2">"http://localhost:11434/api/chat"</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns whether Ollama's `llama-server` process is running.

  The check first uses `pgrep -x llama-server`, then falls back to scanning
  process command names with `ps` when `pgrep` is unavailable.
  """</span>
  <span class="nv">@spec</span> <span class="n">ollama_running?</span><span class="p">()</span> <span class="p">::</span> <span class="n">boolean</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">ollama_running?</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">System</span><span class="o">.</span><span class="n">find_executable</span><span class="p">(</span><span class="s2">"pgrep"</span><span class="p">)</span> <span class="k">do</span>
      <span class="no">nil</span> <span class="o">-&gt;</span> <span class="n">ps_process_running?</span><span class="p">(</span><span class="nv">@llama_server_process</span><span class="p">)</span>
      <span class="n">_pgrep</span> <span class="o">-&gt;</span> <span class="n">pgrep_process_running?</span><span class="p">(</span><span class="nv">@llama_server_process</span><span class="p">)</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Sends a prompt to Ollama's local chat API and returns the generated text.

  Ollama's standard local chat endpoint is
  `http://localhost:11434/api/chat`. The response is streamed as
  newline-delimited JSON; this function concatenates each
  `"message"."content"` chunk into one string.
  """</span>
  <span class="nv">@spec</span> <span class="n">query_model</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">query_model</span><span class="p">(</span><span class="n">prompt</span><span class="p">,</span> <span class="n">model</span> <span class="p">\\</span> <span class="nv">@default_model</span><span class="p">,</span> <span class="n">url</span> <span class="p">\\</span> <span class="nv">@default_chat_url</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">prompt</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">url</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">request_body</span> <span class="o">=</span> <span class="p">%{</span>
      <span class="ss">model:</span> <span class="n">model</span><span class="p">,</span>
      <span class="ss">messages:</span> <span class="p">[</span>
        <span class="p">%{</span><span class="ss">role:</span> <span class="s2">"user"</span><span class="p">,</span> <span class="ss">content:</span> <span class="n">prompt</span><span class="p">}</span>
      <span class="p">],</span>
      <span class="ss">options:</span> <span class="p">%{</span>
        <span class="ss">seed:</span> <span class="mi">123</span><span class="p">,</span>
        <span class="ss">temperature:</span> <span class="mi">0</span><span class="p">,</span>
        <span class="ss">num_ctx:</span> <span class="mi">2048</span>
      <span class="p">}</span>
    <span class="p">}</span>

    <span class="n">response</span> <span class="o">=</span> <span class="no">Req</span><span class="o">.</span><span class="n">post!</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="ss">json:</span> <span class="n">request_body</span><span class="p">,</span> <span class="ss">receive_timeout:</span> <span class="ss">:infinity</span><span class="p">)</span>

    <span class="k">case</span> <span class="n">response</span><span class="o">.</span><span class="n">status</span> <span class="k">do</span>
      <span class="mi">200</span> <span class="o">-&gt;</span>
        <span class="n">parse_chat_response</span><span class="p">(</span><span class="n">response</span><span class="o">.</span><span class="n">body</span><span class="p">)</span>

      <span class="n">status</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="s2">"Ollama chat request failed with status </span><span class="si">#{</span><span class="n">status</span><span class="si">}</span><span class="s2">: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">response</span><span class="o">.</span><span class="n">body</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">pgrep_process_running?</span><span class="p">(</span><span class="n">process_name</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">System</span><span class="o">.</span><span class="n">cmd</span><span class="p">(</span><span class="s2">"pgrep"</span><span class="p">,</span> <span class="p">[</span><span class="s2">"-x"</span><span class="p">,</span> <span class="n">process_name</span><span class="p">],</span> <span class="ss">stderr_to_stdout:</span> <span class="no">true</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">_output</span><span class="p">,</span> <span class="mi">0</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="no">true</span>
      <span class="p">{</span><span class="n">_output</span><span class="p">,</span> <span class="n">_status</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="no">false</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">ps_process_running?</span><span class="p">(</span><span class="n">process_name</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">System</span><span class="o">.</span><span class="n">cmd</span><span class="p">(</span><span class="s2">"ps"</span><span class="p">,</span> <span class="p">[</span><span class="s2">"-axo"</span><span class="p">,</span> <span class="s2">"comm"</span><span class="p">],</span> <span class="ss">stderr_to_stdout:</span> <span class="no">true</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">output</span><span class="p">,</span> <span class="mi">0</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="n">output</span>
        <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="p">,</span> <span class="ss">trim:</span> <span class="no">true</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">any?</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="no">Path</span><span class="o">.</span><span class="n">basename</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">)</span> <span class="o">==</span> <span class="n">process_name</span><span class="p">))</span>

      <span class="p">{</span><span class="n">_output</span><span class="p">,</span> <span class="n">_status</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="no">false</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_chat_response</span><span class="p">(</span><span class="n">body</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">body</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">body</span>
    <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="p">,</span> <span class="ss">trim:</span> <span class="no">true</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="n">decode_chat_chunk!</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">join</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">decode_chat_chunk!</span><span class="p">(</span><span class="n">line</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">line</span>
    <span class="o">|&gt;</span> <span class="no">Jason</span><span class="o">.</span><span class="n">decode!</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">get_in</span><span class="p">([</span><span class="s2">"message"</span><span class="p">,</span> <span class="s2">"content"</span><span class="p">])</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.||</span><span class="p">(</span><span class="s2">""</span><span class="p">)</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And here is a test that shows how we use Llama 3 through Ollama to evaluate our model for instruction tasks:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
  <span class="nv">@tag</span> <span class="ss">:ollama</span>
  <span class="n">test</span> <span class="s2">"7.8 queries local Ollama llama3 model for response scores"</span> <span class="k">do</span>
    <span class="n">assert</span> <span class="no">OllamaUtils</span><span class="o">.</span><span class="n">ollama_running?</span><span class="p">()</span>

    <span class="n">result</span> <span class="o">=</span> <span class="no">OllamaUtils</span><span class="o">.</span><span class="n">query_model</span><span class="p">(</span><span class="s2">"What do Llamas eat?"</span><span class="p">,</span> <span class="s2">"llama3"</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">contains?</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="s2">"Llamas are herbivores"</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">contains?</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="s2">"Minerals"</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">contains?</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="s2">"Grasses"</span><span class="p">)</span>

    <span class="n">response_data_path</span> <span class="o">=</span> <span class="s2">"instruction-data-with-response.json"</span>

    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">response_data_path</span><span class="p">)</span>

    <span class="n">test_data</span> <span class="o">=</span>
      <span class="n">response_data_path</span>
      <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Jason</span><span class="o">.</span><span class="n">decode!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>

    <span class="n">expected_samples</span> <span class="o">=</span> <span class="p">[</span>
      <span class="p">%{</span>
        <span class="ss">output:</span> <span class="s2">"The car is as fast as lightning."</span><span class="p">,</span>
        <span class="ss">model_response:</span> <span class="s2">"The car is as fast as a bullet."</span><span class="p">,</span>
        <span class="ss">score:</span> <span class="mi">85</span>
      <span class="p">},</span>
      <span class="p">%{</span>
        <span class="ss">output:</span> <span class="s2">"The type of cloud typically associated with thunderstorms is cumulonimbus."</span><span class="p">,</span>
        <span class="ss">model_response:</span> <span class="s2">"The type of cloud associated with thunderstorms is a cumulus cloud."</span><span class="p">,</span>
        <span class="ss">score:</span> <span class="mi">40</span>
      <span class="p">},</span>
      <span class="p">%{</span>
        <span class="ss">output:</span> <span class="s2">"Jane Austen."</span><span class="p">,</span>
        <span class="ss">model_response:</span> <span class="s2">"The author of 'Pride and Prejudice' is Jane Austen."</span><span class="p">,</span>
        <span class="ss">score:</span> <span class="mi">95</span>
      <span class="p">}</span>
    <span class="p">]</span>

    <span class="n">test_data</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">(</span><span class="n">expected_samples</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">each</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">entry</span><span class="p">,</span> <span class="n">expected</span><span class="p">}</span> <span class="o">-&gt;</span>
      <span class="n">prompt</span> <span class="o">=</span>
        <span class="s2">"Given the input `</span><span class="si">#{</span><span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="n">entry</span><span class="p">)</span><span class="si">}</span><span class="s2">` "</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"and correct output `</span><span class="si">#{</span><span class="n">entry</span><span class="p">[</span><span class="s2">"output"</span><span class="p">]</span><span class="si">}</span><span class="s2">`, "</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"score the model response `</span><span class="si">#{</span><span class="n">entry</span><span class="p">[</span><span class="s2">"model_response"</span><span class="p">]</span><span class="si">}</span><span class="s2">`"</span> <span class="o">&lt;&gt;</span>
          <span class="s2">" on a scale from 0 to 100, where 100 is the best score. "</span>

      <span class="n">score_response</span> <span class="o">=</span> <span class="no">OllamaUtils</span><span class="o">.</span><span class="n">query_model</span><span class="p">(</span><span class="n">prompt</span><span class="p">,</span> <span class="s2">"llama3"</span><span class="p">)</span>

      <span class="n">report</span> <span class="o">=</span>
        <span class="s2">"</span><span class="se">\n</span><span class="s2">Dataset response:</span><span class="se">\n</span><span class="s2">"</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"&gt;&gt; </span><span class="si">#{</span><span class="n">entry</span><span class="p">[</span><span class="s2">"output"</span><span class="p">]</span><span class="si">}</span><span class="se">\n</span><span class="s2">"</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"</span><span class="se">\n</span><span class="s2">Model response:</span><span class="se">\n</span><span class="s2">"</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"&gt;&gt; </span><span class="si">#{</span><span class="n">entry</span><span class="p">[</span><span class="s2">"model_response"</span><span class="p">]</span><span class="si">}</span><span class="se">\n</span><span class="s2">"</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"</span><span class="se">\n</span><span class="s2">Score:</span><span class="se">\n</span><span class="s2">"</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"&gt;&gt; </span><span class="si">#{</span><span class="n">score_response</span><span class="si">}</span><span class="se">\n</span><span class="s2">"</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"</span><span class="se">\n</span><span class="s2">-------------------------</span><span class="se">\n</span><span class="s2">"</span>

      <span class="n">assert</span> <span class="n">entry</span><span class="p">[</span><span class="s2">"output"</span><span class="p">]</span> <span class="o">==</span> <span class="n">expected</span><span class="o">.</span><span class="n">output</span>
      <span class="n">assert</span> <span class="n">entry</span><span class="p">[</span><span class="s2">"model_response"</span><span class="p">]</span> <span class="o">==</span> <span class="n">expected</span><span class="o">.</span><span class="n">model_response</span>
      <span class="n">assert</span> <span class="n">response_score</span><span class="p">(</span><span class="n">score_response</span><span class="p">)</span> <span class="o">==</span> <span class="n">expected</span><span class="o">.</span><span class="n">score</span>

      <span class="n">assert</span> <span class="n">report</span> <span class="o">=~</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">Dataset response:</span><span class="se">\n</span><span class="s2">&gt;&gt; </span><span class="si">#{</span><span class="n">expected</span><span class="o">.</span><span class="n">output</span><span class="si">}</span><span class="se">\n</span><span class="s2">"</span>
      <span class="n">assert</span> <span class="n">report</span> <span class="o">=~</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">Model response:</span><span class="se">\n</span><span class="s2">&gt;&gt; </span><span class="si">#{</span><span class="n">expected</span><span class="o">.</span><span class="n">model_response</span><span class="si">}</span><span class="se">\n</span><span class="s2">"</span>
      <span class="n">assert</span> <span class="n">report</span> <span class="o">=~</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">Score:</span><span class="se">\n</span><span class="s2">&gt;&gt; </span><span class="si">#{</span><span class="n">score_response</span><span class="si">}</span><span class="se">\n</span><span class="s2">"</span>
      <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">ends_with?</span><span class="p">(</span><span class="n">report</span><span class="p">,</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">-------------------------</span><span class="se">\n</span><span class="s2">"</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
</code></pre></div></div>

<p>We added a new function in the <code class="language-plaintext highlighter-rouge">InstructionsEvaluation</code> module that evaluates our model’s instruction outputs with a score in the range <code class="language-plaintext highlighter-rouge">0..100</code>:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@doc</span> <span class="sd">"""
  Scores generated model responses with a local Ollama model.

  `json_key` identifies the generated response field in each instruction
  record, for example `"model_response"`. Entries whose score cannot be
  parsed as an integer are skipped, matching the chapter 7 evaluation loop.
  """</span>
  <span class="nv">@spec</span> <span class="n">generate_model_scores</span><span class="p">([</span><span class="n">instruction_record</span><span class="p">()],</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="p">[</span><span class="n">integer</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">generate_model_scores</span><span class="p">(</span><span class="n">json_data</span><span class="p">,</span> <span class="n">json_key</span><span class="p">,</span> <span class="n">model</span> <span class="p">\\</span> <span class="s2">"llama3"</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">json_data</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">json_key</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">json_data</span><span class="p">,</span> <span class="p">[],</span> <span class="k">fn</span> <span class="n">entry</span><span class="p">,</span> <span class="n">scores</span> <span class="o">-&gt;</span>
      <span class="n">prompt</span> <span class="o">=</span>
        <span class="s2">"Given the input `</span><span class="si">#{</span><span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_input</span><span class="p">(</span><span class="n">entry</span><span class="p">)</span><span class="si">}</span><span class="s2">` "</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"and correct output `</span><span class="si">#{</span><span class="n">entry</span><span class="p">[</span><span class="s2">"output"</span><span class="p">]</span><span class="si">}</span><span class="s2">`, "</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"score the model response `</span><span class="si">#{</span><span class="n">entry</span><span class="p">[</span><span class="n">json_key</span><span class="p">]</span><span class="si">}</span><span class="s2">`"</span> <span class="o">&lt;&gt;</span>
          <span class="s2">" on a scale from 0 to 100, where 100 is the best score. "</span> <span class="o">&lt;&gt;</span>
          <span class="s2">"Respond with the integer number only."</span>

      <span class="n">score</span> <span class="o">=</span> <span class="no">OllamaUtils</span><span class="o">.</span><span class="n">query_model</span><span class="p">(</span><span class="n">prompt</span><span class="p">,</span> <span class="n">model</span><span class="p">)</span>

      <span class="k">try</span> <span class="k">do</span>
        <span class="p">[</span><span class="no">String</span><span class="o">.</span><span class="n">to_integer</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">trim</span><span class="p">(</span><span class="n">score</span><span class="p">))</span> <span class="o">|</span> <span class="n">scores</span><span class="p">]</span>
      <span class="k">rescue</span>
        <span class="no">ArgumentError</span> <span class="o">-&gt;</span>
          <span class="no">IO</span><span class="o">.</span><span class="n">puts</span><span class="p">(</span><span class="s2">"Could not convert score: </span><span class="si">#{</span><span class="n">score</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
          <span class="n">scores</span>
      <span class="k">end</span>
    <span class="k">end</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">()</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>And here is a test that uses this function. The average score for our instruction fine-tuned model is 50:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@tag</span> <span class="ss">:ollama</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">900_000</span>
  <span class="n">test</span> <span class="s2">"7.8 scores all saved instruction responses with Ollama"</span> <span class="k">do</span>
    <span class="n">assert</span> <span class="no">OllamaUtils</span><span class="o">.</span><span class="n">ollama_running?</span><span class="p">()</span>

    <span class="n">response_data_path</span> <span class="o">=</span> <span class="s2">"instruction-data-with-response.json"</span>

    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">response_data_path</span><span class="p">)</span>

    <span class="n">test_data</span> <span class="o">=</span>
      <span class="n">response_data_path</span>
      <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Jason</span><span class="o">.</span><span class="n">decode!</span><span class="p">()</span>

    <span class="n">scores</span> <span class="o">=</span> <span class="no">InstructionsEvaluation</span><span class="o">.</span><span class="n">generate_model_scores</span><span class="p">(</span><span class="n">test_data</span><span class="p">,</span> <span class="s2">"model_response"</span><span class="p">)</span>
    <span class="n">average_score</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">scores</span><span class="p">)</span> <span class="o">/</span> <span class="n">length</span><span class="p">(</span><span class="n">scores</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">scores</span><span class="p">)</span> <span class="o">==</span> <span class="mi">110</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">scores</span><span class="p">)</span> <span class="o">==</span> <span class="n">length</span><span class="p">(</span><span class="n">test_data</span><span class="p">)</span>
    <span class="n">assert_in_delta</span> <span class="n">average_score</span><span class="p">,</span> <span class="mf">50.32</span><span class="p">,</span> <span class="mf">0.4</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="exercises">Exercises</h2>

<h3 id="71-changing-prompt-styles">7.1 Changing prompt styles</h3>

<p>After fine-tuning the model with the Alpaca prompt style, try the Phi-3 prompt style shown in figure 7.4 and observe whether it affects the response quality of the model. We first need to add support for phi3 prompt format:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@doc</span> <span class="sd">"""
  Formats an instruction dataset entry using the requested prompt style.

  `:alpaca` delegates to `format_input/1`. `:phi3` mirrors the Phi-3 chat
  prompt style shown in chapter 7:

      &lt;|user|&gt;
      ...
      &lt;|end|&gt;
      &lt;|assistant|&gt;

  The returned string contains only the model input, not the expected response.
  """</span>
  <span class="nv">@spec</span> <span class="n">format_text</span><span class="p">(</span><span class="n">instruction_record</span><span class="p">(),</span> <span class="ss">:alpaca</span> <span class="o">|</span> <span class="ss">:phi3</span><span class="p">)</span> <span class="p">::</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">format_text</span><span class="p">(</span><span class="n">entry</span><span class="p">,</span> <span class="n">style</span> <span class="p">\\</span> <span class="ss">:alpaca</span><span class="p">)</span>

  <span class="k">def</span> <span class="n">format_text</span><span class="p">(</span><span class="n">entry</span><span class="p">,</span> <span class="ss">:alpaca</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">format_input</span><span class="p">(</span><span class="n">entry</span><span class="p">)</span>

  <span class="k">def</span> <span class="n">format_text</span><span class="p">(%{</span><span class="s2">"instruction"</span> <span class="o">=&gt;</span> <span class="n">instruction</span><span class="p">,</span> <span class="s2">"input"</span> <span class="o">=&gt;</span> <span class="n">input</span><span class="p">},</span> <span class="ss">:phi3</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">instruction</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">input</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">user_text</span> <span class="o">=</span>
      <span class="k">if</span> <span class="n">input</span> <span class="o">==</span> <span class="s2">""</span> <span class="k">do</span>
        <span class="n">instruction</span>
      <span class="k">else</span>
        <span class="n">instruction</span> <span class="o">&lt;&gt;</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">"</span> <span class="o">&lt;&gt;</span> <span class="n">input</span>
      <span class="k">end</span>

    <span class="s2">"&lt;|user|&gt;</span><span class="se">\n</span><span class="s2">"</span> <span class="o">&lt;&gt;</span> <span class="n">user_text</span> <span class="o">&lt;&gt;</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">&lt;|end|&gt;</span><span class="se">\n</span><span class="s2">&lt;|assistant|&gt;</span><span class="se">\n</span><span class="s2">"</span>
  <span class="k">end</span>
</code></pre></div></div>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="gh">diff --git a/lib/llm_scratch/instruction_dataset.ex b/lib/llm_scratch/instruction_dataset.ex
index a1658d8..c229439 100644
</span><span class="gd">--- a/lib/llm_scratch/instruction_dataset.ex
</span><span class="gi">+++ b/lib/llm_scratch/instruction_dataset.ex
</span><span class="p">@@ -3,9 +3,10 @@</span> defmodule LlmScratch.InstructionDataset do
   Pre-tokenized instruction fine-tuning dataset.
 
   This mirrors the PyTorch `InstructionDataset` example from chapter 7. A
<span class="gd">-  dataset is built from decoded instruction maps, formats each map as an
-  Alpaca-style prompt plus response, and pre-tokenizes the full text during
-  construction.
</span><span class="gi">+  dataset is built from decoded instruction maps, formats each map as a prompt
+  plus response, and pre-tokenizes the full text during construction. Alpaca
+  prompt formatting is used by default; Phi-3-style chat formatting can be
+  selected with `prompt_style: :phi3`.
</span> 
   ## Examples
 
<span class="p">@@ -26,27 +27,39 @@</span> defmodule LlmScratch.InstructionDataset do
        25, 198, 40313, 13]
   """
 
<span class="gd">-  @enforce_keys [:data, :encoded_texts]
-  defstruct [:data, :encoded_texts]
</span><span class="gi">+  @enforce_keys [:data, :encoded_texts, :prompt_style]
+  defstruct [:data, :encoded_texts, :prompt_style]
</span> 
   @end_of_text "&lt;|endoftext|&gt;"
   @pad_token_id 50_256
 
   @type instruction_record :: LlmScratch.FineTuneDataLoader.instruction_record()
   @type tokenizer :: String.t()
<span class="gi">+  @type prompt_style :: :alpaca | :phi3
</span>   @type t :: %__MODULE__{
           data: [instruction_record()],
<span class="gd">-          encoded_texts: [[integer()]]
</span><span class="gi">+          encoded_texts: [[integer()]],
+          prompt_style: prompt_style()
</span>         }
 
   @doc """
   Creates an instruction dataset from decoded instruction records.
 
<span class="gd">-  Each record is formatted with `LlmScratch.FineTuneDataLoader.format_input/1`,
-  then the target response is appended as:
</span><span class="gi">+  By default, each record is formatted with
+  `LlmScratch.FineTuneDataLoader.format_input/1`, then the target response is
+  appended as:
</span> 
       \\n\\n### Response:\\n...
 
<span class="gi">+  With `prompt_style: :phi3`, each record is formatted as:
+
+      &lt;|user|&gt;
+      ...
+      &lt;|end|&gt;
+      &lt;|assistant|&gt;
+      ...
+      &lt;|end|&gt;
+
</span>   The resulting full text is encoded immediately and stored in
   `dataset.encoded_texts`.
 
<span class="p">@@ -61,6 +74,8 @@</span> defmodule LlmScratch.InstructionDataset do
 
     * `:allowed_special` - special tokens allowed when `tokenizer` is a
       Tiktoken model name. Defaults to `["&lt;|endoftext|&gt;"]`.
<span class="gi">+    * `:prompt_style` - prompt format for full training examples. Supported
+      values are `:alpaca` and `:phi3`. Defaults to `:alpaca`.
</span> 
   ## Output
 
<span class="p">@@ -97,6 +112,8 @@</span> defmodule LlmScratch.InstructionDataset do
   """
   @spec new([instruction_record()], tokenizer(), keyword()) :: t()
   def new(data, tokenizer, opts \\ []) when is_list(data) do
<span class="gi">+    prompt_style = Keyword.get(opts, :prompt_style, :alpaca)
+
</span>     # Pre-tokenize each complete training example once so repeated dataset
     # access does not rebuild the prompt or call the tokenizer again.
     encoded_texts =
<span class="p">@@ -104,7 +121,7 @@</span> defmodule LlmScratch.InstructionDataset do
         entry
         # The model trains on the instruction/input prompt followed by the
         # expected response target, matching the chapter 7 PyTorch dataset.
<span class="gd">-        |&gt; full_text()
</span><span class="gi">+        |&gt; full_text(prompt_style)
</span>         |&gt; encode_text!(tokenizer, opts)
       end)
 
<span class="p">@@ -112,7 +129,8 @@</span> defmodule LlmScratch.InstructionDataset do
     # from `encoded_texts`, which mirrors the PyTorch dataset fields.
     %__MODULE__{
       data: data,
<span class="gd">-      encoded_texts: encoded_texts
</span><span class="gi">+      encoded_texts: encoded_texts,
+      prompt_style: prompt_style
</span>     }
   end
 
<span class="p">@@ -269,13 +287,24 @@</span> defmodule LlmScratch.InstructionDataset do
     }
   end
 
<span class="gd">-  defp full_text(%{"output" =&gt; output} = entry) when is_binary(output) do
</span><span class="gi">+  defp full_text(%{"output" =&gt; output} = entry, :alpaca) when is_binary(output) do
</span>     instruction_plus_input = LlmScratch.FineTuneDataLoader.format_input(entry)
     response_text = "\n\n### Response:\n#{output}"
 
     instruction_plus_input &lt;&gt; response_text
   end
 
<span class="gi">+  defp full_text(%{"output" =&gt; output} = entry, :phi3) when is_binary(output) do
+    entry
+    |&gt; LlmScratch.FineTuneDataLoader.format_text(:phi3)
+    |&gt; Kernel.&lt;&gt;(output)
+    |&gt; Kernel.&lt;&gt;("\n&lt;|end|&gt;")
+  end
+
+  defp full_text(_entry, prompt_style) do
+    raise ArgumentError, "unsupported prompt style: #{inspect(prompt_style)}"
+  end
+
</span>   defp encode_text!(text, tokenizer, opts) when is_binary(tokenizer) do
     allowed_special = Keyword.get(opts, :allowed_special, [@end_of_text])
     {:ok, token_ids} = Tiktoken.encode(tokenizer, text, allowed_special)
</code></pre></div></div>

<p>Here is a test for instruction fine-tuning using the Phi-3 format:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">:train_long</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">3_600_000</span>
  <span class="n">test</span> <span class="s2">"exercise 7.1 trains instruction model with Phi-3 prompt style"</span> <span class="k">do</span>
    <span class="n">device</span> <span class="o">=</span> <span class="n">use_accelerated_backend</span><span class="p">()</span>

    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>
    <span class="n">batch_size</span> <span class="o">=</span> <span class="mi">8</span>
    <span class="n">num_workers</span> <span class="o">=</span> <span class="mi">3</span>

    <span class="n">model</span> <span class="o">=</span>
      <span class="s2">"355M"</span>
      <span class="o">|&gt;</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="n">file_path</span> <span class="o">=</span>
      <span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"llm_scratch_instruction_data"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"instruction-data.json"</span><span class="p">)</span>

    <span class="n">data</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">download_and_load_instructions_file</span><span class="p">(</span>
        <span class="n">file_path</span><span class="p">,</span>
        <span class="nv">@instruction_data_url</span>
      <span class="p">)</span>

    <span class="n">train_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.85</span><span class="p">)</span>
    <span class="n">test_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">val_portion</span> <span class="o">=</span> <span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">-</span> <span class="n">train_portion</span> <span class="o">-</span> <span class="n">test_portion</span>
    <span class="n">train_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">train_portion</span><span class="p">)</span>
    <span class="n">val_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span> <span class="o">+</span> <span class="n">test_portion</span><span class="p">,</span> <span class="n">val_portion</span><span class="p">)</span>

    <span class="n">train_dataset</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">train_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">prompt_style:</span> <span class="ss">:phi3</span><span class="p">)</span>
    <span class="n">val_dataset</span> <span class="o">=</span> <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">val_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">prompt_style:</span> <span class="ss">:phi3</span><span class="p">)</span>

    <span class="ss">:rand</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="ss">:exsss</span><span class="p">,</span> <span class="p">{</span><span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">})</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">train_dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">collate_fn:</span> <span class="o">&amp;</span><span class="n">binary_instruction_collate</span><span class="o">/</span><span class="mi">1</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span>
      <span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">val_dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">collate_fn:</span> <span class="o">&amp;</span><span class="n">binary_instruction_collate</span><span class="o">/</span><span class="mi">1</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span>
      <span class="p">)</span>

    <span class="n">train_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">5</span><span class="p">)</span>
    <span class="n">val_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">val_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">5</span><span class="p">)</span>

    <span class="n">input_text</span> <span class="o">=</span>
      <span class="n">val_data</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">format_text</span><span class="p">(</span><span class="ss">:phi3</span><span class="p">)</span>

    <span class="n">token_ids</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">input_text</span>
        <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">device</span><span class="p">),</span>
        <span class="mi">35</span><span class="p">,</span>
        <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="mf">0.0</span><span class="p">,</span>
        <span class="no">nil</span><span class="p">,</span>
        <span class="mi">50_256</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">generated_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">num_epochs</span> <span class="o">=</span> <span class="mi">2</span>
    <span class="n">optimizer</span> <span class="o">=</span> <span class="no">Training</span><span class="o">.</span><span class="n">adamw</span><span class="p">(</span><span class="mf">0.00005</span><span class="p">,</span> <span class="ss">weight_decay:</span> <span class="mf">0.1</span><span class="p">)</span>

    <span class="p">{</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">train_losses</span><span class="p">,</span> <span class="n">val_losses</span><span class="p">,</span> <span class="n">tokens_seen</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Training</span><span class="o">.</span><span class="n">train_model_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">train_loader</span><span class="p">,</span>
        <span class="n">val_loader</span><span class="p">,</span>
        <span class="n">optimizer</span><span class="p">,</span>
        <span class="n">device</span><span class="p">,</span>
        <span class="n">num_epochs</span><span class="p">,</span>
        <span class="mi">5</span><span class="p">,</span>
        <span class="mi">5</span><span class="p">,</span>
        <span class="n">input_text</span><span class="p">,</span>
        <span class="n">tokenizer</span><span class="p">,</span>
        <span class="ss">generate_samples:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">return_optimizer:</span> <span class="no">true</span>
      <span class="p">)</span>

    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="s2">"ch7_instruction_finetuned_gpt2_355m_phi3_model_and_optimizer.nx"</span>
    <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">save_training_state!</span><span class="p">(</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="n">metrics_path</span> <span class="o">=</span> <span class="s2">"ch7_instruction_finetuning_phi3_metrics.json"</span>

    <span class="n">write_instruction_training_metrics!</span><span class="p">(</span>
      <span class="n">metrics_path</span><span class="p">,</span>
      <span class="n">num_epochs</span><span class="p">,</span>
      <span class="n">train_losses</span><span class="p">,</span>
      <span class="n">val_losses</span><span class="p">,</span>
      <span class="n">tokens_seen</span>
    <span class="p">)</span>

    <span class="n">assert</span> <span class="n">train_dataset</span><span class="o">.</span><span class="n">prompt_style</span> <span class="o">==</span> <span class="ss">:phi3</span>
    <span class="n">assert</span> <span class="n">val_dataset</span><span class="o">.</span><span class="n">prompt_style</span> <span class="o">==</span> <span class="ss">:phi3</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">starts_with?</span><span class="p">(</span><span class="n">input_text</span><span class="p">,</span> <span class="s2">"&lt;|user|&gt;</span><span class="se">\n</span><span class="s2">"</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">ends_with?</span><span class="p">(</span><span class="n">input_text</span><span class="p">,</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">&lt;|end|&gt;</span><span class="se">\n</span><span class="s2">&lt;|assistant|&gt;</span><span class="se">\n</span><span class="s2">"</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">starts_with?</span><span class="p">(</span><span class="n">generated_text</span><span class="p">,</span> <span class="n">input_text</span><span class="p">)</span>
    <span class="n">assert</span> <span class="n">is_float</span><span class="p">(</span><span class="n">train_loss</span><span class="p">)</span>
    <span class="n">assert</span> <span class="n">is_float</span><span class="p">(</span><span class="n">val_loss</span><span class="p">)</span>
    <span class="n">assert</span> <span class="n">trained_model</span><span class="o">.</span><span class="n">__struct__</span> <span class="o">==</span> <span class="n">model</span><span class="o">.</span><span class="n">__struct__</span>
    <span class="n">assert</span> <span class="p">%</span><span class="no">Training</span><span class="o">.</span><span class="no">AdamW</span><span class="p">{}</span> <span class="o">=</span> <span class="n">trained_optimizer</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">)</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_losses</span><span class="p">)</span> <span class="o">==</span> <span class="mi">47</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">val_losses</span><span class="p">)</span> <span class="o">==</span> <span class="mi">47</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">)</span> <span class="o">==</span> <span class="mi">47</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">train_losses</span> <span class="o">++</span> <span class="n">val_losses</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_float</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_integer</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="n">tokens_seen</span> <span class="o">==</span> <span class="no">Enum</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">List</span><span class="o">.</span><span class="n">last</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">)</span> <span class="o">&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>And here is the training output:</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  Ep 1 (Step 000000): Train loss 2.150, Val loss 2.147
&lt;|user|&gt;
Convert the active sentence to passive: 'The chef cooks the meal every day.'
&lt;|end|&gt;
&lt;|assistant|&gt;
&lt;|user|&gt;
&lt;|end|&gt;
&lt;|assistant|&gt;
&lt;|end|&gt;
&lt;|assistant|&gt;
&lt;|end|&gt;
&lt;|assistant|&gt;
&lt;|end|&gt;
Ep 1 (Step 000005): Train loss 1.200, Val loss 1.174
&lt;|user|&gt;
Convert the active sentence to passive: 'The chef cooks the meal every day.'
&lt;|end|&gt;
&lt;|assistant|&gt;
The chef cooks the meal every day.
&lt;|end|&gt;&lt;|endoftext|&gt;The following is a list of all the Pokémon in the Pokémon Trading Card Game.
For a list of all Pokémon in the Pokémon Trading Card Game, see Pokémon.
If
Ep 1 (Step 000010): Train loss 1.061, Val loss 1.044
&lt;|user|&gt;
Convert the active sentence to passive: 'The chef cooks the meal every day.'
&lt;|end|&gt;
&lt;|assistant|&gt;
The chef cooks the meal every day.
&lt;|end|&gt;&lt;|endoftext|&gt;The following is a guest post by Dr. David J. Karp, a professor of psychiatry at the University of California, San Francisco.
The term 'psychosis'
Ep 1 (Step 000015): Train loss 1.005, Val loss 0.975
&lt;|user|&gt;
Convert the active sentence to passive: 'The chef cooks the meal every day.'
&lt;|end|&gt;
&lt;|assistant|&gt;
The chef cooks the meal every day.
&lt;|end|&gt;&lt;|endoftext|&gt;The following is a guest post by Dr. David B. Smith.
The following is a guest post by Dr. David B. Smith.
The following is a guest

...

Ep 2 (Step 000215): Train loss 0.411, Val loss 0.716
&lt;|user|&gt;
Convert the active sentence to passive: 'The chef cooks the meal every day.'
&lt;|end|&gt;
&lt;|assistant|&gt;
The meal is prepared by the chef every day.
&lt;|end|&gt;&lt;|endoftext|&gt;The following is a list of all the animals that are native to the United States.
1. Squirrel
2. Eagle
3. Bat
4. Eagle
Ep 2 (Step 000220): Train loss 0.385, Val loss 0.716
&lt;|user|&gt;
Convert the active sentence to passive: 'The chef cooks the meal every day.'
&lt;|end|&gt;
&lt;|assistant|&gt;
The meal is prepared every day by the chef.
&lt;|end|&gt;&lt;|endoftext|&gt;The following is a list of all the animals that are native to the United States.
1. Squirrel
2. Eagle
3. Bat
4. Eagle
Ep 2 (Step 000225): Train loss 0.372, Val loss 0.714
&lt;|user|&gt;
Convert the active sentence to passive: 'The chef cooks the meal every day.'
&lt;|end|&gt;
&lt;|assistant|&gt;
The meal is prepared every day by the chef.
&lt;|end|&gt;&lt;|endoftext|&gt;The following is a list of all the animals that are native to the United States.
1. Squirrel
2. Eagle
3. Bat
4. Cro
Ep 2 (Step 000230): Train loss 0.367, Val loss 0.701
&lt;|user|&gt;
Convert the active sentence to passive: 'The chef cooks the meal every day.'
&lt;|end|&gt;
&lt;|assistant|&gt;
The meal is cooked every day by the chef.
&lt;|end|&gt;&lt;|endoftext|&gt;The following is a list of all the animals that are classified as birds.
Birds
Cetacea
Cetacea is a type of bird.

.
Finished in 669.8 seconds (0.00s async, 669.8s sync)
</code></pre></div></div>

<p>This test calculates the Llama 3 score through Ollama. The score is worse than with the Alpaca prompt format:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">:train</span>
  <span class="nv">@tag</span> <span class="ss">:ollama</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">1_800_000</span>
  <span class="n">test</span> <span class="s2">"exercise 7.1 scores Phi-3 prompt style instruction model with Ollama"</span> <span class="k">do</span>
    <span class="n">assert</span> <span class="no">OllamaUtils</span><span class="o">.</span><span class="n">ollama_running?</span><span class="p">()</span>

    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="s2">"ch7_instruction_finetuned_gpt2_355m_phi3_model_and_optimizer.nx"</span>

    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="n">device</span> <span class="o">=</span> <span class="n">use_accelerated_backend</span><span class="p">()</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>

    <span class="p">%{</span><span class="ss">model_state_dict:</span> <span class="n">model</span><span class="p">}</span> <span class="o">=</span> <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">load_training_state!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

    <span class="n">file_path</span> <span class="o">=</span>
      <span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"llm_scratch_instruction_data"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"instruction-data.json"</span><span class="p">)</span>

    <span class="n">data</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">download_and_load_instructions_file</span><span class="p">(</span>
        <span class="n">file_path</span><span class="p">,</span>
        <span class="nv">@instruction_data_url</span>
      <span class="p">)</span>

    <span class="n">train_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.85</span><span class="p">)</span>
    <span class="n">test_portion</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">test_data</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">train_portion</span><span class="p">,</span> <span class="n">test_portion</span><span class="p">)</span>

    <span class="n">output_path</span> <span class="o">=</span> <span class="s2">"instruction-data-with-response-phi3.json"</span>

    <span class="n">enriched_data</span> <span class="o">=</span>
      <span class="n">test_data</span>
      <span class="o">|&gt;</span> <span class="no">InstructionsEvaluation</span><span class="o">.</span><span class="n">write_responses!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span>
        <span class="ss">output_path:</span> <span class="n">output_path</span><span class="p">,</span>
        <span class="ss">prompt_style:</span> <span class="ss">:phi3</span>
      <span class="p">)</span>

    <span class="n">scores</span> <span class="o">=</span> <span class="no">InstructionsEvaluation</span><span class="o">.</span><span class="n">generate_model_scores</span><span class="p">(</span><span class="n">enriched_data</span><span class="p">,</span> <span class="s2">"model_response"</span><span class="p">)</span>
    <span class="n">average_score</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">scores</span><span class="p">)</span> <span class="o">/</span> <span class="n">length</span><span class="p">(</span><span class="n">scores</span><span class="p">)</span>
    <span class="n">metrics_path</span> <span class="o">=</span> <span class="s2">"ch7_instruction_finetuning_phi3_ollama_scores.json"</span>

    <span class="n">write_ollama_score_metrics!</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">,</span> <span class="n">scores</span><span class="p">,</span> <span class="n">average_score</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">output_path</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">)</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">enriched_data</span><span class="p">)</span> <span class="o">==</span> <span class="mi">110</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">scores</span><span class="p">)</span> <span class="o">==</span> <span class="mi">110</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">scores</span><span class="p">)</span> <span class="o">==</span> <span class="n">length</span><span class="p">(</span><span class="n">enriched_data</span><span class="p">)</span>
    <span class="n">assert_in_delta</span> <span class="n">average_score</span><span class="p">,</span> <span class="mf">46.67</span><span class="p">,</span> <span class="mf">0.4</span>
  <span class="k">end</span>
</code></pre></div></div>

<h3 id="exercise-72-instruction-and-input-masking">Exercise 7.2 Instruction and input masking</h3>

<p>After completing the chapter and fine-tuning the model with <code class="language-plaintext highlighter-rouge">InstructionDataset</code>, replace the instruction and input tokens with the <code class="language-plaintext highlighter-rouge">-100</code> mask to use the instruction masking method illustrated in figure 7.13. Then evaluate whether this has a positive effect on model performance. First, we need to support masking input prompts in targets:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="gh">diff --git a/lib/llm_scratch/instruction_dataset.ex b/lib/llm_scratch/instruction_dataset.ex
index c229439..342a03b 100644
</span><span class="gd">--- a/lib/llm_scratch/instruction_dataset.ex
</span><span class="gi">+++ b/lib/llm_scratch/instruction_dataset.ex
</span><span class="p">@@ -27,19 +27,22 @@</span> defmodule LlmScratch.InstructionDataset do
        25, 198, 40313, 13]
   """
 
<span class="gd">-  @enforce_keys [:data, :encoded_texts, :prompt_style]
-  defstruct [:data, :encoded_texts, :prompt_style]
</span><span class="gi">+  @enforce_keys [:data, :encoded_texts, :prompt_style, :mask_out_instructions_in_target]
+  defstruct [:data, :encoded_texts, :prompt_style, :mask_out_instructions_in_target]
</span> 
   @end_of_text "&lt;|endoftext|&gt;"
   @pad_token_id 50_256
 
   @type instruction_record :: LlmScratch.FineTuneDataLoader.instruction_record()
<span class="gi">+  @type encoded_example ::
+          [integer()] | %{token_ids: [integer()], prompt_length: non_neg_integer()}
</span>   @type tokenizer :: String.t()
   @type prompt_style :: :alpaca | :phi3
   @type t :: %__MODULE__{
           data: [instruction_record()],
<span class="gd">-          encoded_texts: [[integer()]],
-          prompt_style: prompt_style()
</span><span class="gi">+          encoded_texts: [encoded_example()],
+          prompt_style: prompt_style(),
+          mask_out_instructions_in_target: boolean()
</span>         }
 
   @doc """
<span class="p">@@ -76,12 +79,18 @@</span> defmodule LlmScratch.InstructionDataset do
       Tiktoken model name. Defaults to `["&lt;|endoftext|&gt;"]`.
     * `:prompt_style` - prompt format for full training examples. Supported
       values are `:alpaca` and `:phi3`. Defaults to `:alpaca`.
<span class="gi">+    * `:mask_out_instructions_in_target` - when `true`, instruction/input
+      prompt targets are masked with `-100` so the loss is computed only after
+      the prompt. Padding targets are still masked independently by
+      `custom_collate/4`. Defaults to `false`.
</span> 
   ## Output
 
   Returns a `%LlmScratch.InstructionDataset{}`. The `:data` field contains the
<span class="gd">-  original records, and `:encoded_texts` contains one pre-tokenized token-id
-  list per record.
</span><span class="gi">+  original records, and `:encoded_texts` contains one pre-tokenized example per
+  record. With `mask_out_instructions_in_target: false`, each example is a
+  token-id list. With `mask_out_instructions_in_target: true`, each example
+  also carries prompt-length metadata used by `custom_collate/4`.
</span> 
   ## Examples
 
<span class="p">@@ -113,16 +122,13 @@</span> defmodule LlmScratch.InstructionDataset do
   @spec new([instruction_record()], tokenizer(), keyword()) :: t()
   def new(data, tokenizer, opts \\ []) when is_list(data) do
     prompt_style = Keyword.get(opts, :prompt_style, :alpaca)
<span class="gi">+    mask_out_instructions_in_target = Keyword.get(opts, :mask_out_instructions_in_target, false)
</span> 
     # Pre-tokenize each complete training example once so repeated dataset
     # access does not rebuild the prompt or call the tokenizer again.
     encoded_texts =
       Enum.map(data, fn entry -&gt;
<span class="gd">-        entry
-        # The model trains on the instruction/input prompt followed by the
-        # expected response target, matching the chapter 7 PyTorch dataset.
-        |&gt; full_text(prompt_style)
-        |&gt; encode_text!(tokenizer, opts)
</span><span class="gi">+        encode_example(entry, tokenizer, opts, prompt_style, mask_out_instructions_in_target)
</span>       end)
 
     # Keep the decoded records for inspection while serving encoded examples
<span class="p">@@ -130,7 +136,8 @@</span> defmodule LlmScratch.InstructionDataset do
     %__MODULE__{
       data: data,
       encoded_texts: encoded_texts,
<span class="gd">-      prompt_style: prompt_style
</span><span class="gi">+      prompt_style: prompt_style,
+      mask_out_instructions_in_target: mask_out_instructions_in_target
</span>     }
   end
 
<span class="p">@@ -226,7 +233,7 @@</span> defmodule LlmScratch.InstructionDataset do
       &gt;
   """
   @spec custom_collate(
<span class="gd">-          tuple() | [[integer()]],
</span><span class="gi">+          tuple() | [encoded_example()],
</span>           integer(),
           integer(),
           nil | pos_integer()
<span class="p">@@ -248,15 +255,17 @@</span> defmodule LlmScratch.InstructionDataset do
 
   def custom_collate(batch, pad_token_id, ignore_index, allowed_max_length)
       when is_list(batch) do
<span class="gi">+    examples = Enum.map(batch, &amp;normalize_example/1)
+
</span>     # Find the longest sequence length after the one extra pad token that the
     # Python collate function appends to every item.
     batch_max_length =
<span class="gd">-      batch
-      |&gt; Enum.map(&amp;(Kernel.length(&amp;1) + 1))
</span><span class="gi">+      examples
+      |&gt; Enum.map(&amp;(Kernel.length(&amp;1.token_ids) + 1))
</span>       |&gt; Enum.max(fn -&gt; 0 end)
 
     {inputs, targets} =
<span class="gd">-      Enum.map(batch, fn item -&gt;
</span><span class="gi">+      Enum.map(examples, fn %{token_ids: item, prompt_length: prompt_length} -&gt;
</span>         # Elixir data is immutable, so this builds the equivalent of
         # `new_item = item.copy(); new_item += [pad_token_id]`.
         new_item = item ++ [pad_token_id]
<span class="p">@@ -272,6 +281,7 @@</span> defmodule LlmScratch.InstructionDataset do
           padded
           |&gt; Enum.slice(1, batch_max_length - 1)
           |&gt; mask_extra_padding_targets(pad_token_id, ignore_index)
<span class="gi">+          |&gt; mask_prompt_targets(prompt_length, ignore_index)
</span> 
         # Optionally cap both rows to the model context length.
         {
<span class="p">@@ -287,8 +297,51 @@</span> defmodule LlmScratch.InstructionDataset do
     }
   end
 
<span class="gi">+  defp encode_example(entry, tokenizer, opts, prompt_style, false) do
+    entry
+    # The model trains on the instruction/input prompt followed by the
+    # expected response target, matching the chapter 7 PyTorch dataset.
+    |&gt; full_text(prompt_style)
+    |&gt; encode_text!(tokenizer, opts)
+  end
+
+  defp encode_example(entry, tokenizer, opts, prompt_style, true) do
+    token_ids =
+      entry
+      |&gt; full_text(prompt_style)
+      |&gt; encode_text!(tokenizer, opts)
+
+    prompt_length =
+      entry
+      |&gt; prompt_text(prompt_style)
+      |&gt; encode_text!(tokenizer, opts)
+      |&gt; Kernel.length()
+
+    %{token_ids: token_ids, prompt_length: prompt_length}
+  end
+
+  defp encode_example(
+         _entry,
+         _tokenizer,
+         _opts,
+         _prompt_style,
+         mask_out_instructions_in_target
+       ) do
+    raise ArgumentError,
+          "expected :mask_out_instructions_in_target to be a boolean, got: #{inspect(mask_out_instructions_in_target)}"
+  end
+
+  defp normalize_example(token_ids) when is_list(token_ids) do
+    %{token_ids: token_ids, prompt_length: nil}
+  end
+
+  defp normalize_example(%{token_ids: token_ids, prompt_length: prompt_length})
+       when is_list(token_ids) and is_integer(prompt_length) and prompt_length &gt;= 0 do
+    %{token_ids: token_ids, prompt_length: prompt_length}
+  end
+
</span>   defp full_text(%{"output" =&gt; output} = entry, :alpaca) when is_binary(output) do
<span class="gd">-    instruction_plus_input = LlmScratch.FineTuneDataLoader.format_input(entry)
</span><span class="gi">+    instruction_plus_input = prompt_text(entry, :alpaca)
</span>     response_text = "\n\n### Response:\n#{output}"
 
     instruction_plus_input &lt;&gt; response_text
<span class="p">@@ -305,6 +358,10 @@</span> defmodule LlmScratch.InstructionDataset do
     raise ArgumentError, "unsupported prompt style: #{inspect(prompt_style)}"
   end
 
<span class="gi">+  defp prompt_text(entry, prompt_style) do
+    LlmScratch.FineTuneDataLoader.format_text(entry, prompt_style)
+  end
+
</span>   defp encode_text!(text, tokenizer, opts) when is_binary(tokenizer) do
     allowed_special = Keyword.get(opts, :allowed_special, [@end_of_text])
     {:ok, token_ids} = Tiktoken.encode(tokenizer, text, allowed_special)
<span class="p">@@ -331,6 +388,19 @@</span> defmodule LlmScratch.InstructionDataset do
     masked_targets
   end
 
<span class="gi">+  defp mask_prompt_targets(targets, nil, _ignore_index), do: targets
+
+  defp mask_prompt_targets(targets, prompt_length, ignore_index) do
+    mask_count = max(prompt_length - 1, 0)
+
+    targets
+    |&gt; Enum.with_index()
+    |&gt; Enum.map(fn
+      {_target, index} when index &lt; mask_count -&gt; ignore_index
+      {target, _index} -&gt; target
+    end)
+  end
+
</span>   defp maybe_truncate(token_ids, nil), do: token_ids
   defp maybe_truncate(token_ids, max_length), do: Enum.take(token_ids, max_length)
 end
</code></pre></div></div>

<p>Here is the only change that we need to make in the test:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">train_dataset</span> <span class="o">=</span>
      <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">train_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">mask_out_instructions_in_target:</span> <span class="no">true</span><span class="p">)</span>

    <span class="n">val_dataset</span> <span class="o">=</span>
      <span class="no">InstructionDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">val_data</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">mask_out_instructions_in_target:</span> <span class="no">true</span><span class="p">)</span>
</code></pre></div></div>

<p>Score is slightly worse: 49.78</p>

<h3 id="exercise-73-fine-tuning-on-the-original-alpaca-dataset">Exercise 7.3 Fine-tuning on the original Alpaca dataset</h3>

<p>The Alpaca dataset, by researchers at Stanford, is one of the earliest and most popular openly shared instruction datasets, consisting of 52,002 entries. As an alternative to the <code class="language-plaintext highlighter-rouge">instruction-data.json</code> file we use here, consider fine-tuning an LLM on this dataset. The dataset is available at https://mng.bz/NBnE. This dataset contains 52,002 entries, which is approximately 50 times more than the dataset we used here, and most entries are longer. Thus, I highly recommend using a GPU for the training, which will accelerate the fine-tuning process. If you encounter out-of-memory errors, consider reducing the <code class="language-plaintext highlighter-rouge">batch_size</code> from 8 to 4, 2, or even 1. Lowering the <code class="language-plaintext highlighter-rouge">allowed_max_length</code> from 1,024 to 512 or 256 can also help manage memory problems.</p>

<p>Per epoch, we have 5,525 steps, which is 20x more than Sebastian’s training set. With one epoch, training took 24,000 seconds and the average Llama 3 score through Ollama was 43.15. With two epochs, training took 41,060.8 seconds and improved the average score to 50.54.</p>

<h3 id="exercise-74-parameter-efficient-fine-tuning-with-lora">Exercise 7.4 Parameter-efficient fine-tuning with LoRA</h3>

<p>To instruction fine-tune an LLM more efficiently, modify the code in this chapter to use the low-rank adaptation method (LoRA) from appendix E. Compare the training run time and model performance before and after the modification.</p>

<p>With LoRA, training on Sebastian’s instruction set took 648 seconds, compared with 687 seconds without LoRA, with a slightly worse accuracy score of 49.9, compared with 50.32 without LoRA. Here are the commits that implemented the LoRA algorithm: <a href="https://github.com/karlosmid/llm-from-scratch/commit/2c43a61b05c74d612302c89e03f01e79637ec347">feat: 🎸 lora support</a>, <a href="https://github.com/karlosmid/llm-from-scratch/commit/852f5f51c60e007b9663a09752a0f640ab4363f4">feat: 🎸 lora support in model inspect</a>, and <a href="https://github.com/karlosmid/llm-from-scratch/commit/a231974b42c9c7a786d6a029cf6bc1e2f2ed97d8">test: 💍 lora test</a>.</p>]]></content><author><name></name></author><category term="llm-from-scratch" /><category term="gpt" /><category term="tutorials" /><category term="llm" /><category term="gpt" /><category term="workbook" /><summary type="html"><![CDATA[Fine-tuning a pretrained GPT-2 model in Elixir so it can follow instructions, generate task-specific responses, and be evaluated with Ollama.]]></summary></entry><entry><title type="html">Recap of fine-tuning for classification</title><link href="https://karlosmid.com/2026/06/recap-on-fine-tuning-for-classification/" rel="alternate" type="text/html" title="Recap of fine-tuning for classification" /><published>2026-06-12T09:00:00+00:00</published><updated>2026-06-12T09:00:00+00:00</updated><id>https://karlosmid.com/2026/06/recap-on-fine-tuning-for-classification</id><content type="html" xml:base="https://karlosmid.com/2026/06/recap-on-fine-tuning-for-classification/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>I am recapping the blog post <a href="/fine-tuning-for-classification/">Fine-tuning for classification</a> by reviewing the chapter 6 workbook questions from Build an LLM from Scratch and reading Giles’ blog post on chapter 6.</p>

<h2 id="giles-blog-posts">Giles blog posts</h2>

<p>Whoa, only one blog post from Giles on the classification topic: <a href="https://www.gilesthomas.com/2025/10/llm-from-scratch-23-fine-tuning-classification">part23</a>.</p>

<p>In part23, Giles nicely notes all the a-ha moments from chapter 6. As usual, he triggers new a-ha moments for me by explaining those concepts in his own words. For example, we use only the last token logits to calculate the cross-entropy loss with the target class. Why? Because that last token actually represents the whole message, and we want classification for the whole message! Regarding the datasets, we extract the same number of ham and spam messages because we want to avoid training our model to favor ham or spam. He also notices that only the training data drops the last batch when it does not have the same amount of data as the other batches. For the test and validation sets, this is not the case. The reason is very simple: that last partial batch could be more significant in training than other batches with the same data length because it contains less data. Giles also noticed that the dropout of the loaded GPT-2 model is 0.0, which is logical because we do not want to drop data in the frozen layers. Giles had a problem with classification training when we only trained the output head, last transformer block, and final normalization layer. So he trained the whole model.</p>

<h2 id="workbook-questions">Workbook questions</h2>

<h3 id="quick-questions">Quick questions</h3>

<p>1 The primary purpose of <strong>__</strong>_ is to train a model to recognize and predict specific class labels.</p>

<p>A classification-finetuning<br />
B regression<br />
C clustering</p>

<p>My answer: A</p>

<p>Correct answer: +</p>

<p>2 What are the two main categories of fine tuning for language models?</p>

<p>A Instruction-finetuning and classification-finetuning<br />
B Generative and discriminative models<br />
C Pretraining and fine-tuning<br />
D Supervised and unsupervised learning</p>

<p>My answer: A</p>

<p>Correct answer: +</p>

<p>3 What is the purpose of padding text messages in the spam dataset?</p>

<p>A To ensure that all text messages have the same length for batch processing.<br />
B To reduce the computational cost of processing the text messages.<br />
C To remove irrelevant information from the text message.<br />
D To improve the model’s ability to undersatnd the context of the text message.</p>

<p>My answer: A +</p>

<p>Correct answer:</p>

<p>4 What is the role of padding token in the SpamDataset module?</p>

<p>A To mark the beginning and the end of a text message<br />
B To indicate the start of a new sentence.<br />
C To represent unknown or out-of-vocabulary words.<br />
D To fill in the gaps in shorter text messages to match the length of the longest message.</p>

<p>My answer: D</p>

<p>Correct answer: +</p>

<p>5 Why is the last output token of the GPT model chosen for classification-finetuning?</p>

<p>A Because the last token is the most likely to contain the class label information.<br />
B Because the last token is the easiest to process for the model.<br />
C Becaue the last token is the most important token in the text message.<br />
D Because the last token accumulates information from all preceding tokens due to casual attention mask.</p>

<p>My answer: D</p>

<p>Correct answer: +</p>

<p>6 What is the purpose of the cross entropy loss function in the spam classification task?</p>

<p>A To determine the number of epochs required for training.<br />
B To calculate the accuracy of the model’s predictions.<br />
C To measure the difference between the model’s predicted probabilities and the actual class labels.<br />
D To indentify the most important features in the text messages.</p>

<p>My answer: B</p>

<p>Correct answer: C</p>

<p>7 What is the significance of the model’s accuracy on the test set?</p>

<p>A It indicates how well the model generalizes to new, unseen data.<br />
B It determines the number of epochs required for training.<br />
C It reflects the model’s performance on the training data.<br />
D It indicates the model’s ability to learn from the training data.</p>

<p>My answer: A</p>

<p>Correct answer:</p>

<h3 id="61-different-categories-of-fine-tuning">6.1 Different categories of fine-tuning</h3>

<p>1 What are the two most common ways to fine-tune language models?</p>

<p>My answer: instructions and classification fine-tuning.</p>

<p>Correct answer: <em>The two most common ways to fine-tune language models are instruction fine-tuning and classification fine-tuning. Instruction fine-tuning focuses on training the model to understand and execute tasks based on natural language prompts, while classification fine-tuning trains the model to recognize specific class labels.</em></p>

<p>Workbook answer also described those two ways.</p>

<p>2 Describe the purpose of instruction fine-tuning and provide an example.</p>

<p>My answer: Purpose of instruction fine-tuning is to train llm how to do specific tasks. This is supervised training, contrary to llm pretraining phase. As a training data, we give to llm a set of input text with corresponding actions. Example is set of instructions how to buy online an airline ticket.</p>

<p>Correct answer: <em>Instruction fine-tuning aims to improve a model’s ability to understand and execute tasks based on natural language instructions. For example, training a model to translate English sentences into German using specific instructions would be an instance of instruction fine-tuning.</em></p>

<p>I would say that my answer is correct.</p>

<p>3 Explain the concept of classification fine-tuning and give an example.</p>

<p>My answer: In classification fine-tuning we are training llm how to classify input text into several classification classes. This is also supervised training. Example is a set of sms messages where each message is classified as spam/ham class.</p>

<p>Correct answer: <em>Classification fine-tuning involves training a model to recognize specific class labels. For instance, training a model to identify whether an email is spam or not spam is an example of classification fine-tuning. This approach is also used in tasks like image classification and sentiment analysis.</em></p>

<p>My answer is correct.</p>

<p>4 What is the key limitation of a classification fine-tuned model?</p>

<p>My answer: ?</p>

<p>Correct answer: <em>A classification fine-tuned model is restricted to predicting classes it has encountered during training. It can only classify data into the predefined categories it was trained on, limiting its ability to handle tasks outside its training scope.</em></p>

<p>5 Compare the flexibility of instruction fine-tuned models and classification fine-tuned models.</p>

<p>My answer: ?</p>

<p>Correct answer: <em>Instruction fine-tuned models are more flexible and can handle a broader range of tasks based on user instructions. Classification fine-tuned models are highly specialized and excel at categorizing data into predefined classes, but they lack the flexibility of instruction fine-tuned models.</em></p>

<p>6 When is instruction fine-tuning the preferred approach?</p>

<p>My answer: ?</p>

<p>Correct answer: <em>Instruction fine-tuning is best suited for models that need to handle a variety of tasks based on complex user instructions, improving flexibility and interaction quality. It is ideal for applications requiring adaptability and the ability to respond to diverse user requests.</em></p>

<p>I did not know how to answer three questions.</p>

<h3 id="62-preparing-the-dataset">6.2 Preparing the dataset</h3>

<p>1 What is the purpose of the dataset used in this section and what type of data does it contain?</p>

<p>My answer: Purpose of this dataset is to do classification training of existing GPT model so model is able to classify input thext into two classes, ham and spam. It contains sms messages with human categorizations into ham and spam classes.</p>

<p>Correct answer: <em>The dataset used is a collection of text messages classified as either ‘spam’ or ‘non-spam’ (also known as ‘ham’). This dataset is used to demonstrate the process of fine-tuning a large language model for classification tasks</em></p>

<p>2 Why is the dataset undersampled to include an equal number of spam and non-spam messages?</p>

<p>My answer: The reason is that we do not want to train the model to be favour of any of the class groups. With having 50/50 representation of spam/ham messages, our model would be trained with highest classification quality.</p>

<p>Corredt answer: <em>The original dataset has a significant imbalance between ‘spam’ and ‘non- spam’ messages. Undersampling creates a balanced dataset, which is beneficial for training a classification model as it prevents the model from being biased towards the majority class</em></p>

<p>I did not answer that dataset is imbalanced, and in the case of using original dataset, model would be trained to favour class that has majority in dataset.</p>

<p>3 How are the string class labels ham and spam converted into integer class labels?</p>

<p>My answer: ham is 0 and spam is 1</p>

<p>Correct answer: <em>The ‘string’ class labels are converted into integer class labels (0 and 1) using a mapping dictionary. This process is similar to converting text into token IDs, but instead of using the GPT vocabulary, it uses only two token IDs.</em></p>

<p>4 Describe the purpose of the <code class="language-plaintext highlighter-rouge">random_split</code> function and how it divides the dataset.</p>

<p>My answer: Purpose of this function is to randomly shuffle input data and split it to training, validation and testing sets.</p>

<p>Correct answer: <em>The random_split function divides the dataset into three parts: training, validation, and testing. The training set is used to train the model, the validation set is used to adjust hyperparameters and prevent overfitting, and the testing set is used to evaluate the model’s performance on unseen data.</em></p>

<p>Train and validation datasets are used in training function. Validation is used to avoid training overfeeting during the training process. Training dataset is used to evaluate already trained model.</p>

<p>5 Pieces of the code have been removed from two places in this listing. Which of these terms have been removed and where should they go?</p>

<p>A “spam B “ham” C num_spam D ham_records</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">def</span> <span class="n">create_balanced_dataset</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">random_state</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:random_state</span><span class="p">,</span> <span class="mi">123</span><span class="p">)</span>

    <span class="n">spam_records</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">filter</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span><span class="o">.</span><span class="n">label</span> <span class="o">==</span> <span class="s2">"spam"</span><span class="p">))</span>
    <span class="n">ham_records</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">filter</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span><span class="o">.</span><span class="n">label</span> <span class="o">==</span> <span class="n">__1__</span><span class="p">))</span>
    <span class="n">num_spam</span> <span class="o">=</span> <span class="n">length</span><span class="p">(</span><span class="n">spam_records</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">length</span><span class="p">(</span><span class="n">ham_records</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">__2__</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
            <span class="s2">"expected at least </span><span class="si">#{</span><span class="n">num_spam</span><span class="si">}</span><span class="s2"> ham records, got </span><span class="si">#{</span><span class="n">length</span><span class="p">(</span><span class="n">ham_records</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>

    <span class="n">ham_subset</span> <span class="o">=</span>
      <span class="n">ham_records</span>
      <span class="o">|&gt;</span> <span class="n">pandas_sample</span><span class="p">(</span><span class="n">num_spam</span><span class="p">,</span> <span class="n">random_state</span><span class="p">)</span>

    <span class="n">ham_subset</span> <span class="o">++</span> <span class="n">spam_records</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; B, 2 =&gt; num_spam</p>

<p>Correct answer: +</p>

<p>6 What is the significance of saving the dataset as CSV files?</p>

<p>My answer: ?</p>

<p>Correct answer: <em>Saving the dataset as CSV files allows for easy reuse of the data in future steps. This ensures that the prepared dataset can be readily accessed for further analysis and model training.</em></p>

<p>I was expecting something more complicated, like more efficient training.</p>

<p>7 Match the terms on the left to their description on the right:</p>

<table>
  <thead>
    <tr>
      <th>Term</th>
      <th> </th>
      <th>Description</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>Instruction Fine-Tuning</td>
      <td> </td>
      <td>The process of training a language model on a set of tasks using specific instructions to improve its ability to understand and execute tasks described in natural language prompts.</td>
    </tr>
    <tr>
      <td>Classification Fine-Tuning</td>
      <td> </td>
      <td>A model that can perform well across a variety of tasks.</td>
    </tr>
    <tr>
      <td>Generalist Model</td>
      <td> </td>
      <td>A model that is highly trained to perform a specific task.</td>
    </tr>
    <tr>
      <td>Specialized Model</td>
      <td> </td>
      <td>A method of training a model to recognize a specific set of class labels, such as ‘spam’ and ‘not spam’.</td>
    </tr>
  </tbody>
</table>

<p>My answer: 1 =&gt; 1, 2 =&gt; 4, 3 =&gt; 2, 4 =&gt; 3</p>

<p>Correct answer:  +</p>

<h3 id="63-creating-data-loaders">6.3 Creating data loaders</h3>

<p>1 What are the two primary options for batching text messages of varying lenghts?</p>

<p>My answer: We have two options, to pad all the messages to the longest message, or to trim all messages to the shortes messages. In former case we need more resources, but in latter case we loss training data and we have model that has lower classification accuracy.</p>

<p>Correct answer: <em>The two options are truncating all messages to the length of the shortest message or padding all messages to the length of the longest message.  Truncation is computationally cheaper but can lead to information loss, while padding preserves the entire content of all messages.</em></p>

<p>2 Pieces of the code have been removed from four places in this listing. Which of these terms have been removed and where sould they go?</p>

<p>A train_frac B validation_frac C train_end D suffled_records</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">def</span> <span class="n">random_split</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">train_frac</span><span class="p">,</span> <span class="n">validation_frac</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">validate_split_fractions!</span><span class="p">(</span><span class="n">train_frac</span><span class="p">,</span> <span class="n">validation_frac</span><span class="p">)</span>

    <span class="n">random_state</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:random_state</span><span class="p">,</span> <span class="mi">123</span><span class="p">)</span>
    <span class="n">shuffled_records</span> <span class="o">=</span> <span class="n">pandas_sample</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">length</span><span class="p">(</span><span class="n">records</span><span class="p">),</span> <span class="n">random_state</span><span class="p">)</span>
    <span class="n">train_end</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">)</span> <span class="o">*</span> <span class="n">__1__</span><span class="p">)</span>
    <span class="n">validation_end</span> <span class="o">=</span> <span class="n">train_end</span> <span class="o">+</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">__2__</span><span class="p">)</span> <span class="o">*</span> <span class="n">validation_frac</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">__3__</span><span class="p">),</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">,</span> <span class="n">train_end</span><span class="p">,</span> <span class="n">validation_end</span> <span class="o">-</span> <span class="n">train_end</span><span class="p">),</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">,</span> <span class="n">validation_end</span><span class="p">,</span> <span class="n">length</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">)</span> <span class="o">-</span> <span class="n">validation_end</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; A, 2 =&gt; D, 3 =&gt; C</p>

<p>Correct answer: +</p>

<p>3 What is the purpose of padding token and how is it used in the SpamDataset class?</p>

<p>My answer: Purpose of padding token is to make that all messages in input set has same length.</p>

<p>Correct answer: <em>The padding token is used to ensure that all text messages in a batch have the same length. In the SpamDataset class, shorter messages are padded with the padding token ID (50256) to match the length of the longest message</em></p>

<p>I did not answer that padding token is endoftext token.</p>

<p>4 What is the role of SpamDataset in the data loading process?</p>

<p>My answer: Role of this module is to load input data from csv file and to parse it into structured data. We have :data that contains pair of message text and class integer value. :encoded_text contains corresponding message tokens that are padded/truncated.</p>

<p>Correct answer: <em>The SpamDataset class handles several key tasks: loading data from CSV files, tokenizing text messages using the GPT-2 tokenizer, and padding or truncating sequences to a uniform length. It also provides methods to access individual data samples and the overall dataset length.</em></p>

<p>5 Label this diagramby assigning the label numbers to the matching description.</p>

<p><img src="/assets/2026/06/spam-dataset-padding-diagram.png" alt="Spam dataset padding diagram" /></p>

<p>My answer: 1 =&gt; tokenize text, 2 =&gt; token IDs, 3 =&gt; pad to longest sequence, 4 =&gt; padded token IDs, 5 =&gt; longest message has no padding</p>

<p>Correct answer: +</p>

<p>6 How are the validation and test sets padded and truncated in relation to the training set?</p>

<p>My answer: We only pad trining set and do not pad/truncate validation and test sets.</p>

<p>Correct answer: <em>The validation and test sets are padded to match the length of the longest training sequence. Any samples exceeding this length are truncated. This ensures consistency in input lengths across all datasets.</em></p>

<p>We padd all three sets where max length is the longest training message.</p>

<p>7 Describe the structure of a single training batch in terms of input and target tensors.</p>

<p>My answer: target tensor follows input tensor with difference that target tensor is slided related to input tensor.</p>

<p>Correct answer: <em>A single training batch consists of eight text messages represented as token IDs, each with 120 tokens. The corresponding class labels for each message are stored in a separate tensor. This structure allows for efficient processing of multiple training examples simultaneously.</em></p>

<p>I forgot to mention that targets are class values.</p>

<p>8 What is the purpose of the DataLoader module and how is it used to create training, validation and test loaders?</p>

<p>My answer: It loads input data from file to memory structure of batches of input, target pairs.</p>

<p>Correct answer: <em>The DataLoader class is used to create data loaders that efficiently load and process data in batches. It takes a dataset as input and allows for customization of parameters like batch size, shuffling, and number of workers. Separate data loaders are created for training, validation, and testing, each with specific configurations.</em></p>

<p>Dataloaders are about loading input data into batches.</p>

<p>### 6.4 Initializing a model with pretrained weights</p>

<p>1 What is the purpose of initializing a pretrained model before fine-tuning for clasification?</p>

<p>My answer: purpose is to have model that is capable in processing free input text and ready to be trained for classification classes. Doing that, we ship expensive model pretraining phase.</p>

<p>Correct answer: <em>Initializing a pretrained model prepares it for classification fine-tuning by loading the weights learned during pretraining on unlabeled data. This allows the model to leverage its existing knowledge and accelerate the learning process for the specific classification task</em></p>

<p>Main goal is that classification training is faster.</p>

<p>2 Describe the process of loading pretrained weights into the GPT model.</p>

<p>My answer: After we load the model, we freeze model layers that we do not need to train for classification. We are training for classification output head layer, last transformer layer, and final linear layer.</p>

<p>Correct answer: <em>The process involves using the download_and_load_gpt2 function to retrieve the pretrained weights based on the chosen model size. These weights are then loaded into the GPTModel using the load_weights_into_gpt function, ensuring the model is ready for fine-tuning.</em></p>

<p>I missed here what was actually asked.</p>

<p>3 How is the model’s ability to generate coherent text used to verify that the pretrained weights have been loaded correctly?</p>

<p>My answer: ?</p>

<p>Correct answer: <em>By providing a prompt and generating text using the generate_text_simple function, we can assess whether the model produces coherent and meaningful output. If the generated text is sensible, it indicates that the pretrained weights have been loaded successfully.</em></p>

<p>I missed here that question is about verifying the model load process.</p>

<p>4 What is the purpose of prompting the model with a spam message before fine-tuning?</p>

<p>My answer: To verify that model can not classify input to ham/spam classes.</p>

<p>Correct answer: <em>Prompting the model with a spam message allows us to evaluate its initial ability to classify spam. This provides a baseline understanding of the model’s performance before fine-tuning and helps identify areas for improvement.</em></p>

<h3 id="65-adding-a-classification-head">6.5 Adding a classification head</h3>

<p>1 Why is it necessary to modify the output layer of a pretrained LLM for classification fine-tuning?</p>

<p>My answer: We need to modify output layer because output layer of a pretrained LLM predicts next token from the whole tokenizer vocabulary set, while for classification task, we only need set of class values.</p>

<p>Correct answer: <em>The original output layer of a pretrained LLM is designed for language generation, mapping hidden representations to a large vocabulary of tokens. For classification, we need a smaller output layer that maps to the specific classes we want to predict.</em></p>

<p>2 Explain the rationale behind using a number of output nodes equal to the number of classes in a classification task.</p>

<p>My answer: ?</p>

<p>Correct answer: <em>Using a separate output node for each class allows for a more general approach to classification, as it avoids the need to modify the loss function for binary tasks. This approach is easily adaptable to multi-class problems.</em></p>

<p>A-ha moment is that we do not need to modify loss function.</p>

<p>3 Why is it often sufficient to fine-tune only the last layers of a pretrained LLM for a new task?</p>

<p>My answer: Because last layers contain token logits for next token of last token of input that we want to classify.</p>

<p>Correct answer: <em>The lower layers of a pretrained LLM typically capture general language structures and semantics, while the upper layers learn task-specific features. Fine-tuning only the last layers allows for efficient adaptation to new tasks without disrupting the learned general language knowledge.</em></p>

<p>A-ha moment is that separation on lower/upper layers. Upper layers are related to knoweladge related to tasks/classification.</p>

<p>4 Describe the process of freezing and unfrezzing layers in a pretrained LLM for fine-tuning.</p>

<p>My answer: When we freez a model layer, that means that we skip that layer in training loop so we do not update layer weights. Unfreezing the layer, layer becomes part of training loop and its weights get updated.</p>

<p>Correct answer: <em>Freezing layers prevents their weights from being updated during training. This is done by setting the requires_grad attribute of their parameters to False. Unfreezing layers allows their weights to be updated by setting requires_grad to True.</em></p>

<p>5 Why is the last token in a sequence considered the most informative for classification tasks using a casual attention mask?</p>

<p>My answer: Because last token contains all attentions weights with all previous token that consists the text that we want to classify.</p>

<p>Correct answer: <em>The causal attention mask restricts each token’s attention to itself and preceding tokens. As a result, the last token accumulates information from all previous tokens, making it the most comprehensive representation of the input sequence.</em></p>

<p>I did not mention that root cause is usage of casual mask.</p>

<p>6 Explain the significance of focusing on the last output token for classification fine-tuning in a GPT-like model.</p>

<p>My answer: Because next token of last token is classification class value.</p>

<p>Correct answer: <em>Since the last token in a sequence has access to information from all preceding tokens due to the causal attention mask, it provides the most relevant context for classification. Therefore, fine-tuning based on the last token’s output allows for more accurate predictions.</em></p>

<p>A-ha moment is that last token provides the most relevant context.</p>

<h3 id="66-calculating-the-classification-loss-and-accuracy">6.6 Calculating the classification loss and accuracy</h3>

<p>1 Explain how the model’s output is converted into a class label prediction for spam classification.</p>

<p>My answer: We are replacing last model output head so this last layer has last dimension is not equal to size of tokenizer vocabulary, but size is equal to number of elements of class classification. We first initialize this output head as linear layer with size of {768, n} where n is 2, as we have spam/ham classes.</p>

<p>Correct answer: <em>The model’s output for the last token is a 2-dimensional tensor representing the probability scores for each class (spam or not spam).  The class label is determined by finding the index of the highest probability score using the argmax function.</em></p>

<p>I did not mention in answer that we are using <code class="language-plaintext highlighter-rouge">argmax</code> to determine the index of highest class probability.</p>

<p>2 What is the purpose of the <code class="language-plaintext highlighter-rouge">softmax</code> function in this context, and why is it optional?</p>

<p>My answer: Purpose of <code class="language-plaintext highlighter-rouge">softmax</code> function is to transfrom raw logits values to probability values (range of 0 to 1). It is optional as sum of all logits values is 1.</p>

<p>Correct answer: <em>The softmax function converts the model’s output into probabilities that sum to 1. It is optional because the largest output values directly correspond to the highest probability scores, so argmax can be applied directly to the output tensor.</em></p>

<p>I missread optional for optimal :(</p>

<p>3 Describe the function of the <code class="language-plaintext highlighter-rouge">calc_accuracy_loader</code> function and how it is used to calculate the classification accuracy.</p>

<p>My answer: It calculates classification accuraccy over batches from data loader. For each batch counts the number of valid predictions by comparing last logits row (last token) and targets. In the end we sum good predictions over batches and divide that number by sum batch examples over all batches.</p>

<p>Correct answer: <em>The calc_accuracy_loader function iterates through a data loader, applies the argmax prediction to each input, and calculates the proportion of correct predictions. It returns the classification accuracy as a percentage.</em></p>

<p>4 Why is cross-entropy loss used as a proxy for maximizing classification accuracy?</p>

<p>My answer: ?</p>

<p>Correct answer: <em>Classification accuracy is not a differentiable function, making it unsuitable for optimization. Cross-entropy loss is a differentiable function that can be used to optimize the model’s parameters to indirectly maximize classification accuracy</em></p>

<p>5 Explain the key difference between the <code class="language-plaintext highlighter-rouge">calc_loss_batch</code> function used for classification and the one used for language modeling.</p>

<p>My answer: The key difference is that for classification training we are taking into account only logits row that belongs to the last logit. In language model training we are using all logits row (all tokens) as we want to capture attention between all tokens.</p>

<p>Correct answer: <em>The classification calc_loss_batch function focuses on optimizing only the last token’s output, while the language modeling version optimizes the output of all tokens. This difference reflects the focus on predicting the class label for the last token in classification</em></p>

<p>A-ha moment is what difference between last token and all tokens reflects.</p>

<h3 id="67-fine-tuning-the-model-on-supervised-data">6.7 Fine-tuning the model on supervised data</h3>

<p>1 What is the primary difference between the training loop used for fine-tuning and the one used for pretraining?</p>

<p>My answer: The primary difference is that we do not train for classification all model layers, but just output head layer, last transformer layer, and final linearization layer.</p>

<p>Correct answer: <em>The primary difference is that during fine-tuning, we calculate the classification accuracy instead of generating sample text to evaluate the model’s performance</em></p>

<p>Wrong answer.</p>

<p>2 Describe the key modifications made to the <code class="language-plaintext highlighter-rouge">train_classifier_simple</code> function compared to the <code class="language-plaintext highlighter-rouge">train_model_simple</code> function used for pretraining.</p>

<p>My answer: First difference is that for calculating training accuracy we are using just last token from logits. Second difference is that we only train in model there layers.</p>

<p>Correct answer:a_The train_classifier_simple function tracks the number of training examples seen instead of tokens and calculates accuracy after each epoch. It also omits the step of printing a sample text._</p>

<p>My answer is wrong.</p>

<p>3 What is the purpose of the <code class="language-plaintext highlighter-rouge">evaluate_model</code> function in the context of fine-tuning?</p>

<p>My answer: It calculates training and validation losess so we could check do we have model classification overfeeting.</p>

<p>Correct answer: <em>The evaluate_model function calculates the loss for both the training and validation sets, providing insights into the model’s performance on both seen and unseen data</em></p>

<p>4 Explain the significance of the loss curves plotted here:</p>

<p><img src="/assets/2026/06/classification-loss-plot.png" alt="Classification training and validation loss plot" /></p>

<p>My answer: This graph shows training and validation losses over seen examples and number of epochs. We can see that both loss functions drop after epoch 2 to almost zero value, showing that our model is trained for classification very quickly with very good accuraccy.</p>

<p>Correct answer: <em>The loss curves in Figure 6.16 show a sharp decline in both training and validation loss during the initial epochs, indicating effective learning. The close proximity of the curves suggests that the model is generalizing well to unseen data and not overfitting</em></p>

<p>I missed to answer that when training and validtion curves are close, model is not overfitting.</p>

<p>5 What factors influence the choice of the number of epochs during fine-tuning?</p>

<p>My answer: Main factor is at which epoch number loss function becomes near to zero values.</p>

<p>Correct answer: <em>The number of epochs depends on the dataset’s complexity and the task’s difficulty. Overfitting may necessitate reducing the number of epochs, while insufficient training might require increasing them</em></p>

<p>A-ha moment is that lower number of epochs could help in overfitting problem.</p>

<p>6 Pieces of the code have been removed from three places in this listing. Which of these terms have been removed and where should they go?</p>

<p>A adam B adamW C train_loader D val_loader E test_loader</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="n">optimizer</span> <span class="o">=</span> <span class="no">Training</span><span class="o">.</span><span class="n">__1__</span><span class="p">(</span><span class="mf">5.0e-5</span><span class="p">,</span> <span class="ss">weight_decay:</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">num_epochs</span> <span class="o">=</span> <span class="mi">5</span>

    <span class="p">{</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">train_losses</span><span class="p">,</span> <span class="n">val_losses</span><span class="p">,</span> <span class="n">train_accs</span><span class="p">,</span> <span class="n">val_accs</span><span class="p">,</span>
     <span class="n">examples_seen</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Training</span><span class="o">.</span><span class="n">train_classifier_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">__2__</span><span class="p">,</span>
        <span class="n">__3__</span><span class="p">,</span>
        <span class="n">optimizer</span><span class="p">,</span>
        <span class="n">device</span><span class="p">,</span>
        <span class="n">num_epochs</span><span class="p">,</span>
        <span class="mi">50</span><span class="p">,</span>
        <span class="mi">5</span><span class="p">,</span>
        <span class="ss">return_optimizer:</span> <span class="no">true</span>
      <span class="p">)</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; A, 2 =&gt; train_loader, 3 =&gt; val_loader</p>

<p>Correct answer: +</p>

<p>7 How does <code class="language-plaintext highlighter-rouge">eval_iter</code> parameter affect the accuracy estimations during training?</p>

<p>My answer: Not aplicable as we do not have this parameter in Elixir code.</p>

<p>Correct answer: <em>The eval_iter parameter determines the number of batches used to calculate training and validation accuracy. A smaller eval_iter value leads to faster training but less accurate performance estimations.</em></p>

<p>8 Put these steps to prepare and fine-tune the GPT-2 Model into the correct order:</p>

<p>A Load pretrained GPT-2 model weights.<br />
  B Train the model using the training data loader for a specified number of epochs.<br />
  C Freeze all model parameters except the output layer and the last transformer block.<br />
  D Initialize an AdamW optimizer.<br />
  E Define functions to calculate loss and accuracy for a batch and a data loader.<br />
  F Set the <code class="language-plaintext highlighter-rouge">requires_grad</code> attribute to <code class="language-plaintext highlighter-rouge">True</code> for the new output layer and the last transformer block.<br />
  G Replace the original output layer with a new linear layer mapping to two classes (spam vs not spam)<br />
  H Evaluate the model’s performance on the validation set after each epoch.</p>

<p>My answer: 1 =&gt; A, 2 =&gt; C, 3 =&gt; G, 4 =&gt; F, 5 =&gt; E, 6 =&gt; D, 7 =&gt; B, 8 =&gt; H</p>

<p>Correct answer: +</p>

<h3 id="68-using-the-llm-as-a-spam-classifier">6.8 Using the LLM as a spam classifier</h3>

<p>1 Describe the process of using a fine-tuned LLM for spam classification.</p>

<p>My answer: Input text that we want to classify with trained model we first encode with tokenizer to token IDS. Then we trim or padd this input text to the length of the longest training input. That text we eval using model to get the classfication class value.</p>

<p>Correct answer: <em>The process involves using a pre-trained LLM, fine-tuned for classification, to predict the class label (spam or not spam) for a given text message.  This is achieved by converting the text into token IDs, feeding it to the model, and obtaining the predicted class label based on the model’s output.</em></p>

<p>2 Explain the role of the <code class="language-plaintext highlighter-rouge">classify_review</code> function in spam classification.</p>

<p>My answer: This function classifies input text to be spam or not spam.</p>

<p>Correct answer: <em>The classify_review function takes a text message as input, preprocesses it, converts it into token IDs, feeds it to the fine-tuned model, and predicts the class label (spam or not spam) based on the model’s output. It then returns the corresponding class name.</em></p>

<p>3 How is the classification accuracy of a spam classification model evaluated?</p>

<p>My answer: It is evaluated during the training when we compare predicted batch logits with target logits.</p>

<p>Correct answer: <em>The classification accuracy is evaluated by calculating the fraction or percentage of correct predictions made by the model on a test dataset.  This metric indicates how well the model is able to correctly classify spam and non-spam messages.</em></p>

<p>4 In this code listing, order the stages to match the numbers in the listing:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">def</span> <span class="n">classify_review</span><span class="p">(</span>
        <span class="n">text</span><span class="p">,</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">tokenizer</span><span class="p">,</span>
        <span class="n">device</span> <span class="p">\\</span> <span class="ss">:default</span><span class="p">,</span>
        <span class="n">max_length</span> <span class="p">\\</span> <span class="no">nil</span><span class="p">,</span>
        <span class="n">pad_token_id</span> <span class="p">\\</span> <span class="mi">50_256</span>
      <span class="p">)</span> <span class="k">do</span>
    <span class="n">__1__</span> <span class="n">input_ids</span> <span class="o">=</span> <span class="n">encode_text</span><span class="p">(</span><span class="n">text</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>
    <span class="n">supported_context_length</span> <span class="o">=</span> <span class="n">context_length</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
    <span class="n">target_length</span> <span class="o">=</span> <span class="n">target_input_length</span><span class="p">(</span><span class="n">input_ids</span><span class="p">,</span> <span class="n">max_length</span><span class="p">,</span> <span class="n">supported_context_length</span><span class="p">)</span>

    <span class="c1"># Classifier inference uses the same fixed-length preparation as training:</span>
    <span class="c1"># truncate long messages first, then pad short messages so the model sees</span>
    <span class="c1"># the expected context width.</span>
    <span class="n">input_ids</span> <span class="o">=</span>
      <span class="n">input_ids</span>
     <span class="n">__2__</span> <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">target_length</span><span class="p">)</span>
      <span class="n">__3__</span><span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">++</span> <span class="no">List</span><span class="o">.</span><span class="n">duplicate</span><span class="p">(</span><span class="n">pad_token_id</span><span class="p">,</span> <span class="n">target_length</span> <span class="o">-</span> <span class="n">length</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">))))</span>

    <span class="n">__4__</span> <span class="n">input_tensor</span> <span class="o">=</span>
      <span class="n">input_ids</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">new_axis</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="n">maybe_transfer_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

    <span class="c1"># Calling forward/2 is inference-only in this codebase. Dropout is only used</span>
    <span class="c1"># by train/3, and gradients are only created inside explicit value_and_grad</span>
    <span class="c1"># calls, so this is the Nx counterpart to model.eval() and no_grad().</span>
    <span class="n">predicted_label</span> <span class="o">=</span>
      <span class="n">model</span>
     <span class="n">__5__</span> <span class="o">|&gt;</span> <span class="n">forward!</span><span class="p">(</span><span class="n">input_tensor</span><span class="p">)</span>
    <span class="n">__6__</span>  <span class="o">|&gt;</span> <span class="n">last_position_logits</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">()</span>

   <span class="n">__7__</span> <span class="k">if</span> <span class="n">predicted_label</span> <span class="o">==</span> <span class="mi">1</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="s2">"spam"</span><span class="p">,</span> <span class="k">else</span><span class="p">:</span> <span class="s2">"not spam"</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Models inference without gradient tracking<br />
  Add batch dimension<br />
  Truncate sequences if they are too long<br />
  Prepare inputs to the model<br />
  Pad sequences to the longest sequence<br />
  Logits of the last output token<br />
  Return the classified result</p>

<p>My answer: Models inference without gradient tracking =&gt; 5, Add batch dimension =&gt; 4, Truncate sequence if they are too long =&gt; 2, Prepare inputs to the model =&gt; 1, Pad sequences to the longes sequence =&gt; 3, Logits of the last output token =&gt; 6, Return the classified result =&gt; 7</p>

<p>Correct answer: +</p>

<p>5 What is the purpose of saving a fine-tuned spam classification model?</p>

<p>My answer: By saving the model, we could load it latter so we do not need to retrain it again.</p>

<p>Correct answer: <em>Saving the model allows for its reuse later without the need for retraining. This is useful for deploying the model for real-time spam detection or for further experimentation and analysis.</em></p>

<p>6 How does classification fine-tuning process differ from the pretraining process of an LLM?</p>

<p>My answer: We are trining classification fine-tuning process using just last token of input example. We also train only on output head layer, last transformer layer and last linearization layer. Output head layer is mapped on only two values, as our classification has two classes. We padd input text to the length of longest input message.</p>

<p>Correct answer: <em>While both processes involve converting text into token IDs and using a cross-entropy loss function, classification fine-tuning focuses on training the model to output a correct class label, whereas pretraining aims to predict the next token in a sequence. Classification fine-tuning also involves replacing the output layer of the LLM with a smaller classification layer</em></p>]]></content><author><name></name></author><category term="llm-from-scratch" /><category term="gpt" /><category term="tutorials" /><category term="llm" /><category term="gpt" /><category term="workbook" /><summary type="html"><![CDATA[Workbook notes on fine-tuning a pretrained GPT-2 model for SMS spam classification in Elixir.]]></summary></entry><entry><title type="html">Fine-tuning for classification</title><link href="https://karlosmid.com/2026/06/fine-tuning-for-classification/" rel="alternate" type="text/html" title="Fine-tuning for classification" /><published>2026-06-05T09:00:00+00:00</published><updated>2026-06-05T09:00:00+00:00</updated><id>https://karlosmid.com/2026/06/fine-tuning-for-classification</id><content type="html" xml:base="https://karlosmid.com/2026/06/fine-tuning-for-classification/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>This post fine-tunes a pretrained GPT-2 model for SMS spam classification in Elixir. It prepares a balanced ham/spam dataset, builds padded data loaders, replaces the language-model output head with a classification head, freezes most model layers, trains the classifier, plots loss and accuracy, adds Apple GPU support with EMLX, and finally uses the saved model as a spam filter.</p>

<p>It follows <a href="_posts/2026-05-14-pretraining-on-unlabeled-data-in-elixir.md">pretraining on unlabeled data in Elixir</a> and uses the source code from my <a href="https://github.com/karlosmid/llm-from-scratch">GitHub repository</a>.</p>

<h2 id="61-different-categories-of-fine-tuning">6.1 Different categories of fine-tuning</h2>

<p>We have a pretrained foundation model. We also downloaded OpenAI’s largest GPT-2 model, with 1 billion parameters. This means we are ready to fine-tune the model for our needs. We can do instruction fine-tuning or classification fine-tuning. Instruction fine-tuning teaches the model to perform tasks from prompts, for example language translation. In classification fine-tuning, we classify input data: is a text message spam or not spam, or in the famous image-classification example, is an image a hot dog or not a hot dog? Classification fine-tuning is usually easier than instruction fine-tuning because it targets a narrower task and needs a smaller dataset.</p>

<h2 id="62-preparing-the-data-set">6.2 Preparing the data set</h2>

<p>We are going to download a set of text messages classified as <code class="language-plaintext highlighter-rouge">ham</code> and <code class="language-plaintext highlighter-rouge">spam</code>. The dataset contains more <code class="language-plaintext highlighter-rouge">ham</code> messages than <code class="language-plaintext highlighter-rouge">spam</code> messages, so we balance the two classes by sampling an equal number of examples from each class. Then we randomly split the data into three sets: training, validation, and test.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">FineTuneDataLoader</span> <span class="k">do</span>
  <span class="kn">import</span> <span class="no">Bitwise</span>

  <span class="nv">@moduledoc</span> <span class="sd">"""
  Helpers for loading small supervised fine-tuning datasets.

  The SMS spam dataset is a two-column TSV file with labels (`ham` or `spam`)
  and text. Records can be loaded as maps or converted to Nx tensor pairs for
  use with `LlmScratch.DataLoader`.
  """</span>

  <span class="nv">@spam_url</span> <span class="s2">"https://archive.ics.uci.edu/static/public/228/sms+spam+collection.zip"</span>
  <span class="nv">@zip_path</span> <span class="s2">"sms_spam_collection.zip"</span>
  <span class="nv">@extracted_path</span> <span class="s2">"sms_spam_collection"</span>
  <span class="nv">@data_filename</span> <span class="s2">"SMSSpamCollection.tsv"</span>
  <span class="nv">@raw_filename</span> <span class="s2">"SMSSpamCollection"</span>
  <span class="nv">@label_ids</span> <span class="p">%{</span><span class="s2">"ham"</span> <span class="o">=&gt;</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">"spam"</span> <span class="o">=&gt;</span> <span class="mi">1</span><span class="p">}</span>
  <span class="nv">@pad_token_id</span> <span class="mi">50_256</span>
  <span class="nv">@mt19937_n</span> <span class="mi">624</span>
  <span class="nv">@mt19937_m</span> <span class="mi">397</span>
  <span class="nv">@mt19937_matrix_a</span> <span class="mh">0x9908B0DF</span>
  <span class="nv">@mt19937_upper_mask</span> <span class="mh">0x80000000</span>
  <span class="nv">@mt19937_lower_mask</span> <span class="mh">0x7FFFFFFF</span>
  <span class="nv">@uint32_mask</span> <span class="mh">0xFFFFFFFF</span>

  <span class="nv">@type</span> <span class="n">spam_record</span> <span class="p">::</span> <span class="p">%{</span>
          <span class="ss">label:</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">label_id:</span> <span class="mi">0</span> <span class="o">|</span> <span class="mi">1</span><span class="p">,</span>
          <span class="ss">text:</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Downloads and extracts the UCI SMS spam dataset if the TSV is not present.

  Accepts a keyword list of path and URL options.

  ## Input Parameters

    * `opts` - optional keyword list controlling where the dataset is downloaded,
      extracted, and saved.

  ## Options

    * `:url` - dataset zip URL. Defaults to the UCI SMS spam collection URL.
    * `:zip_path` - local zip cache path. Defaults to
      `"sms_spam_collection.zip"`.
    * `:extracted_path` - directory where the zip is extracted. Defaults to
      `"sms_spam_collection"`.
    * `:data_file_path` - final TSV path. Defaults to
      `"sms_spam_collection/SMSSpamCollection.tsv"`.

  ## Output

  Returns the final TSV path as a string. If `:data_file_path` already exists,
  no download or extraction is performed. If the zip exists but the TSV does
  not, the existing zip is reused.
  """</span>
  <span class="nv">@spec</span> <span class="n">download_and_unzip_spam_data</span><span class="p">(</span><span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">download_and_unzip_spam_data</span><span class="p">(</span><span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">url</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:url</span><span class="p">,</span> <span class="nv">@spam_url</span><span class="p">)</span>
    <span class="n">zip_path</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:zip_path</span><span class="p">,</span> <span class="nv">@zip_path</span><span class="p">)</span>
    <span class="n">extracted_path</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:extracted_path</span><span class="p">,</span> <span class="nv">@extracted_path</span><span class="p">)</span>
    <span class="n">data_file_path</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:data_file_path</span><span class="p">,</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">extracted_path</span><span class="p">,</span> <span class="nv">@data_filename</span><span class="p">))</span>

    <span class="k">if</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">data_file_path</span><span class="p">)</span> <span class="k">do</span>
      <span class="n">data_file_path</span>
    <span class="k">else</span>
      <span class="no">File</span><span class="o">.</span><span class="n">mkdir_p!</span><span class="p">(</span><span class="no">Path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">zip_path</span><span class="p">))</span>
      <span class="no">File</span><span class="o">.</span><span class="n">mkdir_p!</span><span class="p">(</span><span class="n">extracted_path</span><span class="p">)</span>

      <span class="k">unless</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">zip_path</span><span class="p">)</span> <span class="k">do</span>
        <span class="n">download_file!</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="n">zip_path</span><span class="p">)</span>
      <span class="k">end</span>

      <span class="n">unzip_file!</span><span class="p">(</span><span class="n">zip_path</span><span class="p">,</span> <span class="n">extracted_path</span><span class="p">)</span>
      <span class="n">rename_raw_dataset!</span><span class="p">(</span><span class="n">extracted_path</span><span class="p">,</span> <span class="n">data_file_path</span><span class="p">)</span>

      <span class="n">data_file_path</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Loads the SMS spam TSV into a list of maps.

  ## Input Parameters

    * `data_file_path` - path to a two-column TSV file. Each row must have a
      label (`ham` or `spam`) in the first column and the message text in the
      second column.

  ## Output

  Returns a list of `%{label: label, label_id: label_id, text: text}` maps.
  `label_id` is `0` for `ham` and `1` for `spam`.
  """</span>
  <span class="nv">@spec</span> <span class="n">load_spam_data</span><span class="p">(</span><span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="p">[</span><span class="n">spam_record</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">load_spam_data</span><span class="p">(</span><span class="n">data_file_path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">data_file_path</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">stream!</span><span class="p">([],</span> <span class="ss">:line</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Stream</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="no">String</span><span class="o">.</span><span class="n">trim_trailing</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Stream</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="no">String</span><span class="o">.</span><span class="n">trim_trailing</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="s2">"</span><span class="se">\r</span><span class="s2">"</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Stream</span><span class="o">.</span><span class="n">reject</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">==</span> <span class="s2">""</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="n">parse_spam_row!</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Counts records by label.

  ## Input Parameters

    * `records` - list of records returned by `load_spam_data/1`, or any list
      of maps containing a `:label` key.

  ## Output

  Returns a map where each key is a label value and each value is the number
  of records with that label.
  """</span>
  <span class="nv">@spec</span> <span class="n">label_counts</span><span class="p">([</span><span class="n">map</span><span class="p">()])</span> <span class="p">::</span> <span class="p">%{</span><span class="n">optional</span><span class="p">(</span><span class="n">term</span><span class="p">())</span> <span class="o">=&gt;</span> <span class="n">non_neg_integer</span><span class="p">()}</span>
  <span class="k">def</span> <span class="n">label_counts</span><span class="p">(</span><span class="n">records</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Enum</span><span class="o">.</span><span class="n">frequencies_by</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="o">&amp;</span> <span class="nv">&amp;1</span><span class="o">.</span><span class="n">label</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Creates a balanced SMS spam dataset by undersampling ham records.

  This mirrors the pandas example:

      num_spam = df[df["Label"] == "spam"].shape[0]
      ham_subset = df[df["Label"] == "ham"].sample(num_spam, random_state=123)
      balanced_df = pd.concat([ham_subset, df[df["Label"] == "spam"]])

  ## Input Parameters

    * `records` - list of records returned by `load_spam_data/1`.
    * `opts` - optional keyword list controlling the deterministic sample.

  ## Options

    * `:random_state` - integer seed used for sampling ham records. Defaults
      to `123`. Sampling mirrors pandas' NumPy-backed `sample(random_state: ...)`
      ordering.

  ## Output

  Returns a list containing a deterministic sample of ham records followed by
  all spam records. The returned dataset has the same number of ham and spam
  records.
  """</span>
  <span class="nv">@spec</span> <span class="n">create_balanced_dataset</span><span class="p">([</span><span class="n">spam_record</span><span class="p">()],</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="p">[</span><span class="n">spam_record</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">create_balanced_dataset</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">random_state</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:random_state</span><span class="p">,</span> <span class="mi">123</span><span class="p">)</span>

    <span class="n">spam_records</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">filter</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span><span class="o">.</span><span class="n">label</span> <span class="o">==</span> <span class="s2">"spam"</span><span class="p">))</span>
    <span class="n">ham_records</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">filter</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span><span class="o">.</span><span class="n">label</span> <span class="o">==</span> <span class="s2">"ham"</span><span class="p">))</span>
    <span class="n">num_spam</span> <span class="o">=</span> <span class="n">length</span><span class="p">(</span><span class="n">spam_records</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">length</span><span class="p">(</span><span class="n">ham_records</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">num_spam</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
            <span class="s2">"expected at least </span><span class="si">#{</span><span class="n">num_spam</span><span class="si">}</span><span class="s2"> ham records, got </span><span class="si">#{</span><span class="n">length</span><span class="p">(</span><span class="n">ham_records</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>

    <span class="n">ham_subset</span> <span class="o">=</span>
      <span class="n">ham_records</span>
      <span class="o">|&gt;</span> <span class="n">pandas_sample</span><span class="p">(</span><span class="n">num_spam</span><span class="p">,</span> <span class="n">random_state</span><span class="p">)</span>

    <span class="n">ham_subset</span> <span class="o">++</span> <span class="n">spam_records</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Converts string class labels to integer class labels.

  This mirrors the pandas example:

      balanced_df["Label"] = balanced_df["Label"].map({"ham": 0, "spam": 1})

  ## Input Parameters

    * `records` - list of maps with `:label` set to `"ham"` or `"spam"`.

  ## Output

  Returns a list of records with `:label` replaced by `0` for ham and `1` for
  spam. Other fields, including `:text` and `:label_id`, are preserved.
  """</span>
  <span class="nv">@spec</span> <span class="n">map_labels_to_ids</span><span class="p">([</span><span class="n">spam_record</span><span class="p">()])</span> <span class="p">::</span> <span class="p">[</span><span class="n">map</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">map_labels_to_ids</span><span class="p">(</span><span class="n">records</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="k">fn</span> <span class="p">%{</span><span class="ss">label:</span> <span class="n">label</span><span class="p">}</span> <span class="o">=</span> <span class="n">record</span> <span class="o">-&gt;</span>
      <span class="p">%{</span><span class="n">record</span> <span class="o">|</span> <span class="ss">label:</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="nv">@label_ids</span><span class="p">,</span> <span class="n">label</span><span class="p">)}</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Randomly shuffles records and splits them into train, validation, and test sets.

  This mirrors the pandas example:

      df = df.sample(frac=1, random_state=123).reset_index(drop=True)
      train_end = int(len(df) * train_frac)
      validation_end = train_end + int(len(df) * validation_frac)
      train_df = df[:train_end]
      validation_df = df[train_end:validation_end]
      test_df = df[validation_end:]

  The test split is implied as the remaining records after the train and
  validation splits.

  ## Input Parameters

    * `records` - list of records to shuffle and split.
    * `train_frac` - fraction of records assigned to the train split.
    * `validation_frac` - fraction of records assigned to the validation split.
    * `opts` - optional keyword list controlling deterministic shuffling.

  ## Options

    * `:random_state` - integer seed used for shuffling. Defaults to `123`.
      Shuffling mirrors pandas' NumPy-backed `sample(frac=1, random_state=...)`
      ordering.

  ## Output

  Returns `{train_records, validation_records, test_records}`. The train and
  validation split sizes are truncated with `trunc/1`, matching Python's
  `int(...)`; the test split receives the remainder.
  """</span>
  <span class="nv">@spec</span> <span class="n">random_split</span><span class="p">([</span><span class="n">map</span><span class="p">()],</span> <span class="n">number</span><span class="p">(),</span> <span class="n">number</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="p">{[</span><span class="n">map</span><span class="p">()],</span> <span class="p">[</span><span class="n">map</span><span class="p">()],</span> <span class="p">[</span><span class="n">map</span><span class="p">()]}</span>
  <span class="k">def</span> <span class="n">random_split</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">train_frac</span><span class="p">,</span> <span class="n">validation_frac</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">validate_split_fractions!</span><span class="p">(</span><span class="n">train_frac</span><span class="p">,</span> <span class="n">validation_frac</span><span class="p">)</span>

    <span class="n">random_state</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:random_state</span><span class="p">,</span> <span class="mi">123</span><span class="p">)</span>
    <span class="n">shuffled_records</span> <span class="o">=</span> <span class="n">pandas_sample</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">length</span><span class="p">(</span><span class="n">records</span><span class="p">),</span> <span class="n">random_state</span><span class="p">)</span>
    <span class="n">train_end</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">)</span> <span class="o">*</span> <span class="n">train_frac</span><span class="p">)</span>
    <span class="n">validation_end</span> <span class="o">=</span> <span class="n">train_end</span> <span class="o">+</span> <span class="n">trunc</span><span class="p">(</span><span class="n">length</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">)</span> <span class="o">*</span> <span class="n">validation_frac</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">train_end</span><span class="p">),</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">,</span> <span class="n">train_end</span><span class="p">,</span> <span class="n">validation_end</span> <span class="o">-</span> <span class="n">train_end</span><span class="p">),</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">,</span> <span class="n">validation_end</span><span class="p">,</span> <span class="n">length</span><span class="p">(</span><span class="n">shuffled_records</span><span class="p">)</span> <span class="o">-</span> <span class="n">validation_end</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Builds a reusable dataloader for the SMS spam dataset.

  By default, records are tokenized into `{input_ids, label}` Nx tensor pairs.
  Pass `tokenize: false` to batch raw record maps.

  ## Input Parameters

    * `opts` - optional keyword list controlling dataset location,
      tokenization, and batching.

  ## Options

    * `:data_file_path` - existing TSV path to load. When omitted,
      `download_and_unzip_spam_data/1` is called with the same options.
    * `:url` - dataset zip URL used only when `:data_file_path` is omitted.
    * `:zip_path` - local zip cache path used only when downloading or
      extracting.
    * `:extracted_path` - extraction directory used only when downloading or
      extracting.
    * `:tokenize` - when `true`, returns batches of `{input_ids, label}` Nx
      tensor pairs. When `false`, batches raw record maps. Defaults to `true`.
    * `:tokenizer_model` - Tiktoken model name used when tokenizing. Defaults
      to `"code-davinci-002"`.
    * `:max_length` - fixed token length for each text tensor. Longer messages
      are truncated and shorter messages are padded. Defaults to `120`.
    * `:pad_token_id` - token id used for padding. Defaults to GPT-2's
      end-of-text token id, `50256`.
    * `:batch_size` - number of examples per batch. Defaults to `32`.
    * `:shuffle` - whether to shuffle records before batching. Defaults to
      `true`.
    * `:drop_last` - whether to drop an incomplete final batch. Defaults to
      `false`.
    * `:num_workers` - worker count passed through to `LlmScratch.DataLoader`.
      Defaults to `0`.

  ## Output

  Returns the map produced by `LlmScratch.DataLoader.new/2`, including
  `:stream`, `:batches`, `:length`, `:batch_size`, `:shuffle`, `:drop_last`,
  and `:num_workers`.
  """</span>
  <span class="nv">@spec</span> <span class="n">create_spam_dataloader</span><span class="p">(</span><span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">map</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">create_spam_dataloader</span><span class="p">(</span><span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">data_file_path</span> <span class="o">=</span>
      <span class="no">Keyword</span><span class="o">.</span><span class="n">get_lazy</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:data_file_path</span><span class="p">,</span> <span class="k">fn</span> <span class="o">-&gt;</span>
        <span class="n">download_and_unzip_spam_data</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="n">records</span> <span class="o">=</span> <span class="n">load_spam_data</span><span class="p">(</span><span class="n">data_file_path</span><span class="p">)</span>

    <span class="n">dataset</span> <span class="o">=</span>
      <span class="k">if</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:tokenize</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span> <span class="k">do</span>
        <span class="n">to_tensor_dataset</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
      <span class="k">else</span>
        <span class="n">records</span>
      <span class="k">end</span>

    <span class="no">LlmScratch</span><span class="o">.</span><span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span>
      <span class="ss">batch_size:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:batch_size</span><span class="p">,</span> <span class="mi">32</span><span class="p">),</span>
      <span class="ss">shuffle:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:shuffle</span><span class="p">,</span> <span class="no">true</span><span class="p">),</span>
      <span class="ss">drop_last:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:drop_last</span><span class="p">,</span> <span class="no">false</span><span class="p">),</span>
      <span class="ss">num_workers:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:num_workers</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
    <span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Converts records to `{input_ids, label}` tensor pairs.

  Texts are encoded with Tiktoken, truncated to `:max_length`, and padded with
  GPT-2's end-of-text token id.

  ## Input Parameters

    * `records` - list of `%{label_id: label_id, text: text}` maps.
    * `opts` - optional keyword list controlling tokenization.

  ## Options

    * `:tokenizer_model` - Tiktoken model name used for encoding text. Defaults
      to `"code-davinci-002"`.
    * `:max_length` - fixed number of token ids in each input tensor. Defaults
      to `120`.
    * `:pad_token_id` - token id appended until each input reaches
      `:max_length`. Defaults to `50256`.

  ## Output

  Returns a list of `{input_ids, label}` tuples. `input_ids` is an Nx tensor of
  shape `{max_length}` and signed 64-bit integer type. `label` is a scalar Nx
  tensor containing `0` for `ham` or `1` for `spam`.
  """</span>
  <span class="nv">@spec</span> <span class="n">to_tensor_dataset</span><span class="p">([</span><span class="n">spam_record</span><span class="p">()],</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="p">[{</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()}]</span>
  <span class="k">def</span> <span class="n">to_tensor_dataset</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:tokenizer_model</span><span class="p">,</span> <span class="s2">"code-davinci-002"</span><span class="p">)</span>
    <span class="n">max_length</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:max_length</span><span class="p">,</span> <span class="mi">120</span><span class="p">)</span>
    <span class="n">pad_token_id</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:pad_token_id</span><span class="p">,</span> <span class="nv">@pad_token_id</span><span class="p">)</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="k">fn</span> <span class="p">%{</span><span class="ss">label_id:</span> <span class="n">label_id</span><span class="p">,</span> <span class="ss">text:</span> <span class="n">text</span><span class="p">}</span> <span class="o">-&gt;</span>
      <span class="n">input_ids</span> <span class="o">=</span>
        <span class="n">text</span>
        <span class="o">|&gt;</span> <span class="n">encode_text!</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="n">pad_or_truncate</span><span class="p">(</span><span class="n">max_length</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>

      <span class="p">{</span><span class="n">input_ids</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">label_id</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})}</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_spam_row!</span><span class="p">(</span><span class="n">row</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">String</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">row</span><span class="p">,</span> <span class="s2">"</span><span class="se">\t</span><span class="s2">"</span><span class="p">,</span> <span class="ss">parts:</span> <span class="mi">2</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">[</span><span class="n">label</span><span class="p">,</span> <span class="n">text</span><span class="p">]</span> <span class="ow">when</span> <span class="n">is_map_key</span><span class="p">(</span><span class="nv">@label_ids</span><span class="p">,</span> <span class="n">label</span><span class="p">)</span> <span class="o">-&gt;</span>
        <span class="p">%{</span><span class="ss">label:</span> <span class="n">label</span><span class="p">,</span> <span class="ss">label_id:</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="nv">@label_ids</span><span class="p">,</span> <span class="n">label</span><span class="p">),</span> <span class="ss">text:</span> <span class="n">text</span><span class="p">}</span>

      <span class="p">[</span><span class="n">label</span><span class="p">,</span> <span class="n">_text</span><span class="p">]</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"unknown SMS spam label </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">label</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>

      <span class="n">_</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected a two-column TSV row, got </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">row</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">encode_text!</span><span class="p">(</span><span class="n">text</span><span class="p">,</span> <span class="n">model</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">token_ids</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">text</span><span class="p">,</span> <span class="p">[</span><span class="s2">"&lt;|endoftext|&gt;"</span><span class="p">])</span>
    <span class="n">token_ids</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">pad_or_truncate</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">max_length</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">token_ids</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">max_length</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="k">fn</span> <span class="n">ids</span> <span class="o">-&gt;</span> <span class="n">ids</span> <span class="o">++</span> <span class="no">List</span><span class="o">.</span><span class="n">duplicate</span><span class="p">(</span><span class="n">pad_token_id</span><span class="p">,</span> <span class="n">max_length</span> <span class="o">-</span> <span class="n">length</span><span class="p">(</span><span class="n">ids</span><span class="p">))</span> <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">pandas_sample</span><span class="p">(</span><span class="n">records</span><span class="p">,</span> <span class="n">count</span><span class="p">,</span> <span class="n">random_state</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">if</span> <span class="n">count</span> <span class="o">&gt;</span> <span class="n">length</span><span class="p">(</span><span class="n">records</span><span class="p">)</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"cannot take a sample larger than the population"</span>
    <span class="k">end</span>

    <span class="n">records_tuple</span> <span class="o">=</span> <span class="no">List</span><span class="o">.</span><span class="n">to_tuple</span><span class="p">(</span><span class="n">records</span><span class="p">)</span>

    <span class="p">{</span><span class="n">indexes</span><span class="p">,</span> <span class="n">_rng</span><span class="p">}</span> <span class="o">=</span>
      <span class="n">records</span>
      <span class="o">|&gt;</span> <span class="n">length</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="n">numpy_random_state_permutation</span><span class="p">(</span><span class="n">random_state</span><span class="p">)</span>

    <span class="n">indexes</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">count</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="n">elem</span><span class="p">(</span><span class="n">records_tuple</span><span class="p">,</span> <span class="nv">&amp;1</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">numpy_random_state_permutation</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{[],</span> <span class="n">mt19937_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">numpy_random_state_permutation</span><span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">rng</span> <span class="o">=</span> <span class="n">mt19937_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span>

    <span class="p">{</span><span class="n">permuted</span><span class="p">,</span> <span class="n">rng</span><span class="p">}</span> <span class="o">=</span>
      <span class="p">(</span><span class="n">size</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span><span class="o">..</span><span class="mi">1</span><span class="o">//-</span><span class="mi">1</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">({</span><span class="no">List</span><span class="o">.</span><span class="n">to_tuple</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">to_list</span><span class="p">(</span><span class="mi">0</span><span class="o">..</span><span class="p">(</span><span class="n">size</span> <span class="o">-</span> <span class="mi">1</span><span class="p">))),</span> <span class="n">rng</span><span class="p">},</span> <span class="k">fn</span> <span class="n">i</span><span class="p">,</span> <span class="p">{</span><span class="n">items</span><span class="p">,</span> <span class="n">rng</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="p">{</span><span class="n">j</span><span class="p">,</span> <span class="n">rng</span><span class="p">}</span> <span class="o">=</span> <span class="n">mt19937_random_interval</span><span class="p">(</span><span class="n">rng</span><span class="p">,</span> <span class="n">i</span><span class="p">)</span>
        <span class="p">{</span><span class="n">swap_tuple</span><span class="p">(</span><span class="n">items</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">),</span> <span class="n">rng</span><span class="p">}</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="p">{</span><span class="no">Tuple</span><span class="o">.</span><span class="n">to_list</span><span class="p">(</span><span class="n">permuted</span><span class="p">),</span> <span class="n">rng</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">mt19937_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">mt</span> <span class="o">=</span>
      <span class="mi">1</span><span class="o">..</span><span class="p">(</span><span class="nv">@mt19937_n</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">([</span><span class="n">seed</span> <span class="o">&amp;&amp;&amp;</span> <span class="nv">@uint32_mask</span><span class="p">],</span> <span class="k">fn</span> <span class="n">i</span><span class="p">,</span> <span class="n">acc</span> <span class="o">-&gt;</span>
        <span class="n">previous</span> <span class="o">=</span> <span class="n">hd</span><span class="p">(</span><span class="n">acc</span><span class="p">)</span>
        <span class="n">next</span> <span class="o">=</span> <span class="mi">1_812_433_253</span> <span class="o">*</span> <span class="n">bxor</span><span class="p">(</span><span class="n">previous</span><span class="p">,</span> <span class="n">previous</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">30</span><span class="p">)</span> <span class="o">+</span> <span class="n">i</span> <span class="o">&amp;&amp;&amp;</span> <span class="nv">@uint32_mask</span>
        <span class="p">[</span><span class="n">next</span> <span class="o">|</span> <span class="n">acc</span><span class="p">]</span>
      <span class="k">end</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">to_tuple</span><span class="p">()</span>

    <span class="p">%{</span><span class="ss">mt:</span> <span class="n">mt</span><span class="p">,</span> <span class="ss">index:</span> <span class="nv">@mt19937_n</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">mt19937_random_interval</span><span class="p">(</span><span class="n">rng</span><span class="p">,</span> <span class="n">max</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">mask</span> <span class="o">=</span>
      <span class="n">max</span>
      <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">|||</span> <span class="nv">&amp;1</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">1</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">|||</span> <span class="nv">&amp;1</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">2</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">|||</span> <span class="nv">&amp;1</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">4</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">|||</span> <span class="nv">&amp;1</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">8</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">|||</span> <span class="nv">&amp;1</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">16</span><span class="p">))</span>

    <span class="n">mt19937_random_interval</span><span class="p">(</span><span class="n">rng</span><span class="p">,</span> <span class="n">max</span><span class="p">,</span> <span class="n">mask</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">mt19937_random_interval</span><span class="p">(</span><span class="n">rng</span><span class="p">,</span> <span class="n">max</span><span class="p">,</span> <span class="n">mask</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">value</span><span class="p">,</span> <span class="n">rng</span><span class="p">}</span> <span class="o">=</span> <span class="n">mt19937_extract_number</span><span class="p">(</span><span class="n">rng</span><span class="p">)</span>
    <span class="n">value</span> <span class="o">=</span> <span class="n">value</span> <span class="o">&amp;&amp;&amp;</span> <span class="n">mask</span>

    <span class="k">if</span> <span class="n">value</span> <span class="o">&lt;=</span> <span class="n">max</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">value</span><span class="p">,</span> <span class="n">rng</span><span class="p">}</span>
    <span class="k">else</span>
      <span class="n">mt19937_random_interval</span><span class="p">(</span><span class="n">rng</span><span class="p">,</span> <span class="n">max</span><span class="p">,</span> <span class="n">mask</span><span class="p">)</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">mt19937_extract_number</span><span class="p">(%{</span><span class="ss">index:</span> <span class="nv">@mt19937_n</span><span class="p">}</span> <span class="o">=</span> <span class="n">rng</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">rng</span>
    <span class="o">|&gt;</span> <span class="n">mt19937_twist</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">mt19937_extract_number</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">mt19937_extract_number</span><span class="p">(%{</span><span class="ss">mt:</span> <span class="n">mt</span><span class="p">,</span> <span class="ss">index:</span> <span class="n">index</span><span class="p">}</span> <span class="o">=</span> <span class="n">rng</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">y</span> <span class="o">=</span> <span class="n">elem</span><span class="p">(</span><span class="n">mt</span><span class="p">,</span> <span class="n">index</span><span class="p">)</span>
    <span class="n">y</span> <span class="o">=</span> <span class="n">bxor</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="n">y</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">11</span><span class="p">)</span>
    <span class="n">y</span> <span class="o">=</span> <span class="n">bxor</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="n">y</span> <span class="o">&lt;&lt;&lt;</span> <span class="mi">7</span> <span class="o">&amp;&amp;&amp;</span> <span class="mh">0x9D2C5680</span><span class="p">)</span>
    <span class="n">y</span> <span class="o">=</span> <span class="n">bxor</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="n">y</span> <span class="o">&lt;&lt;&lt;</span> <span class="mi">15</span> <span class="o">&amp;&amp;&amp;</span> <span class="mh">0xEFC60000</span><span class="p">)</span>
    <span class="n">y</span> <span class="o">=</span> <span class="n">bxor</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="n">y</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">18</span><span class="p">)</span>

    <span class="p">{</span><span class="n">y</span> <span class="o">&amp;&amp;&amp;</span> <span class="nv">@uint32_mask</span><span class="p">,</span> <span class="p">%{</span><span class="n">rng</span> <span class="o">|</span> <span class="ss">index:</span> <span class="n">index</span> <span class="o">+</span> <span class="mi">1</span><span class="p">}}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">mt19937_twist</span><span class="p">(%{</span><span class="ss">mt:</span> <span class="n">mt</span><span class="p">}</span> <span class="o">=</span> <span class="n">rng</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">mt</span> <span class="o">=</span>
      <span class="mi">0</span><span class="o">..</span><span class="p">(</span><span class="nv">@mt19937_n</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">mt</span><span class="p">,</span> <span class="k">fn</span> <span class="n">i</span><span class="p">,</span> <span class="n">acc</span> <span class="o">-&gt;</span>
        <span class="n">x</span> <span class="o">=</span>
          <span class="p">(</span><span class="n">elem</span><span class="p">(</span><span class="n">acc</span><span class="p">,</span> <span class="n">i</span><span class="p">)</span> <span class="o">&amp;&amp;&amp;</span> <span class="nv">@mt19937_upper_mask</span><span class="p">)</span> <span class="o">+</span>
            <span class="p">(</span><span class="n">elem</span><span class="p">(</span><span class="n">acc</span><span class="p">,</span> <span class="n">rem</span><span class="p">(</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span> <span class="nv">@mt19937_n</span><span class="p">))</span> <span class="o">&amp;&amp;&amp;</span> <span class="nv">@mt19937_lower_mask</span><span class="p">)</span>

        <span class="n">x_a</span> <span class="o">=</span> <span class="n">x</span> <span class="o">&gt;&gt;&gt;</span> <span class="mi">1</span>
        <span class="n">x_a</span> <span class="o">=</span> <span class="k">if</span> <span class="n">rem</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="n">bxor</span><span class="p">(</span><span class="n">x_a</span><span class="p">,</span> <span class="nv">@mt19937_matrix_a</span><span class="p">),</span> <span class="k">else</span><span class="p">:</span> <span class="n">x_a</span>

        <span class="n">value</span> <span class="o">=</span> <span class="n">bxor</span><span class="p">(</span><span class="n">elem</span><span class="p">(</span><span class="n">acc</span><span class="p">,</span> <span class="n">rem</span><span class="p">(</span><span class="n">i</span> <span class="o">+</span> <span class="nv">@mt19937_m</span><span class="p">,</span> <span class="nv">@mt19937_n</span><span class="p">)),</span> <span class="n">x_a</span><span class="p">)</span> <span class="o">&amp;&amp;&amp;</span> <span class="nv">@uint32_mask</span>
        <span class="n">put_elem</span><span class="p">(</span><span class="n">acc</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="p">%{</span><span class="n">rng</span> <span class="o">|</span> <span class="ss">mt:</span> <span class="n">mt</span><span class="p">,</span> <span class="ss">index:</span> <span class="mi">0</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">swap_tuple</span><span class="p">(</span><span class="n">tuple</span><span class="p">,</span> <span class="n">index</span><span class="p">,</span> <span class="n">index</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">tuple</span>

  <span class="k">defp</span> <span class="n">swap_tuple</span><span class="p">(</span><span class="n">tuple</span><span class="p">,</span> <span class="n">left</span><span class="p">,</span> <span class="n">right</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">left_value</span> <span class="o">=</span> <span class="n">elem</span><span class="p">(</span><span class="n">tuple</span><span class="p">,</span> <span class="n">left</span><span class="p">)</span>
    <span class="n">right_value</span> <span class="o">=</span> <span class="n">elem</span><span class="p">(</span><span class="n">tuple</span><span class="p">,</span> <span class="n">right</span><span class="p">)</span>

    <span class="n">tuple</span>
    <span class="o">|&gt;</span> <span class="n">put_elem</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">right_value</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">put_elem</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">left_value</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_split_fractions!</span><span class="p">(</span><span class="n">train_frac</span><span class="p">,</span> <span class="n">validation_frac</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">cond</span> <span class="k">do</span>
      <span class="n">train_frac</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">validation_frac</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"split fractions must be greater than or equal to 0"</span>

      <span class="n">train_frac</span> <span class="o">+</span> <span class="n">validation_frac</span> <span class="o">&gt;</span> <span class="mi">1</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"train and validation fractions must sum to 1 or less"</span>

      <span class="no">true</span> <span class="o">-&gt;</span>
        <span class="ss">:ok</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">download_file!</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="n">destination</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">tmp_destination</span> <span class="o">=</span> <span class="n">destination</span> <span class="o">&lt;&gt;</span> <span class="s2">".download"</span>
    <span class="no">File</span><span class="o">.</span><span class="n">rm</span><span class="p">(</span><span class="n">tmp_destination</span><span class="p">)</span>

    <span class="n">response</span> <span class="o">=</span>
      <span class="no">Req</span><span class="o">.</span><span class="n">get!</span><span class="p">(</span>
        <span class="n">url</span><span class="p">,</span>
        <span class="ss">into:</span> <span class="no">File</span><span class="o">.</span><span class="n">stream!</span><span class="p">(</span><span class="n">tmp_destination</span><span class="p">),</span>
        <span class="ss">receive_timeout:</span> <span class="mi">120_000</span><span class="p">,</span>
        <span class="ss">retry:</span> <span class="no">false</span>
      <span class="p">)</span>

    <span class="k">if</span> <span class="n">response</span><span class="o">.</span><span class="n">status</span> <span class="ow">in</span> <span class="mi">200</span><span class="o">..</span><span class="mi">299</span> <span class="k">do</span>
      <span class="no">File</span><span class="o">.</span><span class="n">rename!</span><span class="p">(</span><span class="n">tmp_destination</span><span class="p">,</span> <span class="n">destination</span><span class="p">)</span>
    <span class="k">else</span>
      <span class="no">File</span><span class="o">.</span><span class="n">rm</span><span class="p">(</span><span class="n">tmp_destination</span><span class="p">)</span>
      <span class="k">raise</span> <span class="s2">"failed to download </span><span class="si">#{</span><span class="n">url</span><span class="si">}</span><span class="s2">: HTTP </span><span class="si">#{</span><span class="n">response</span><span class="o">.</span><span class="n">status</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">rescue</span>
    <span class="n">error</span> <span class="o">-&gt;</span>
      <span class="no">File</span><span class="o">.</span><span class="n">rm</span><span class="p">(</span><span class="n">destination</span> <span class="o">&lt;&gt;</span> <span class="s2">".download"</span><span class="p">)</span>
      <span class="n">reraise</span> <span class="n">error</span><span class="p">,</span> <span class="n">__STACKTRACE__</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">unzip_file!</span><span class="p">(</span><span class="n">zip_path</span><span class="p">,</span> <span class="n">extracted_path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">zip_path</span>
    <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">to_charlist</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="ss">:zip</span><span class="o">.</span><span class="n">unzip</span><span class="p">([{</span><span class="ss">:cwd</span><span class="p">,</span> <span class="no">String</span><span class="o">.</span><span class="n">to_charlist</span><span class="p">(</span><span class="n">extracted_path</span><span class="p">)}])</span>
    <span class="o">|&gt;</span> <span class="k">case</span> <span class="k">do</span>
      <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">_files</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="ss">:ok</span>
      <span class="p">{</span><span class="ss">:error</span><span class="p">,</span> <span class="n">reason</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="k">raise</span> <span class="s2">"failed to unzip </span><span class="si">#{</span><span class="n">zip_path</span><span class="si">}</span><span class="s2">: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">reason</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">rename_raw_dataset!</span><span class="p">(</span><span class="n">extracted_path</span><span class="p">,</span> <span class="n">data_file_path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">raw_file_path</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">extracted_path</span><span class="p">,</span> <span class="nv">@raw_filename</span><span class="p">)</span>

    <span class="k">if</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">raw_file_path</span><span class="p">)</span> <span class="k">do</span>
      <span class="no">File</span><span class="o">.</span><span class="n">rename!</span><span class="p">(</span><span class="n">raw_file_path</span><span class="p">,</span> <span class="n">data_file_path</span><span class="p">)</span>
    <span class="k">else</span>
      <span class="k">raise</span> <span class="s2">"expected extracted dataset file at </span><span class="si">#{</span><span class="n">raw_file_path</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>
<span class="k">end</span>

</code></pre></div></div>

<p>And here is a test that uses <code class="language-plaintext highlighter-rouge">FineTuneDataLoader</code>:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmFromScratch6Test</span> <span class="k">do</span>
  <span class="kn">use</span> <span class="no">ExUnit</span><span class="o">.</span><span class="no">Case</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">FineTuneDataLoader</span>

  <span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="n">test</span> <span class="s2">"6.1 loads SMS spam TSV rows and counts labels"</span> <span class="k">do</span>
    <span class="n">tmp_dir</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">(),</span> <span class="s2">"llm_scratch_sms_spam_collection"</span><span class="p">)</span>
    <span class="n">data_file_path</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="s2">"SMSSpamCollection.tsv"</span><span class="p">)</span>

    <span class="n">data_file_path</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">download_and_unzip_spam_data</span><span class="p">(</span>
        <span class="ss">zip_path:</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">tmp_dir</span><span class="p">,</span> <span class="s2">"sms_spam_collection.zip"</span><span class="p">),</span>
        <span class="ss">extracted_path:</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">tmp_dir</span><span class="p">,</span> <span class="s2">"sms_spam_collection"</span><span class="p">),</span>
        <span class="ss">data_file_path:</span> <span class="n">data_file_path</span>
      <span class="p">)</span>

    <span class="n">records</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">load_spam_data</span><span class="p">(</span><span class="n">data_file_path</span><span class="p">)</span>
    <span class="n">label_counts</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">label_counts</span><span class="p">(</span><span class="n">records</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">records</span><span class="p">)</span> <span class="o">==</span> <span class="mi">5_574</span>
    <span class="n">assert</span> <span class="n">label_counts</span><span class="p">[</span><span class="s2">"ham"</span><span class="p">]</span> <span class="o">==</span> <span class="mi">4_827</span>
    <span class="n">assert</span> <span class="n">label_counts</span><span class="p">[</span><span class="s2">"spam"</span><span class="p">]</span> <span class="o">==</span> <span class="mi">747</span>

    <span class="n">balanced_records</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">create_balanced_dataset</span><span class="p">(</span><span class="n">records</span><span class="p">)</span>
    <span class="n">balanced_label_counts</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">label_counts</span><span class="p">(</span><span class="n">balanced_records</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">balanced_records</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1_494</span>
    <span class="n">assert</span> <span class="n">balanced_label_counts</span><span class="p">[</span><span class="s2">"ham"</span><span class="p">]</span> <span class="o">==</span> <span class="mi">747</span>
    <span class="n">assert</span> <span class="n">balanced_label_counts</span><span class="p">[</span><span class="s2">"spam"</span><span class="p">]</span> <span class="o">==</span> <span class="mi">747</span>

    <span class="n">mapped_records</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">map_labels_to_ids</span><span class="p">(</span><span class="n">balanced_records</span><span class="p">)</span>
    <span class="n">mapped_label_counts</span> <span class="o">=</span> <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">label_counts</span><span class="p">(</span><span class="n">mapped_records</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">mapped_label_counts</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="mi">747</span>
    <span class="n">assert</span> <span class="n">mapped_label_counts</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">==</span> <span class="mi">747</span>

    <span class="p">{</span><span class="n">train_records</span><span class="p">,</span> <span class="n">validation_records</span><span class="p">,</span> <span class="n">test_records</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">FineTuneDataLoader</span><span class="o">.</span><span class="n">random_split</span><span class="p">(</span><span class="n">mapped_records</span><span class="p">,</span> <span class="mf">0.7</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">)</span>

    <span class="n">write_csv!</span><span class="p">(</span><span class="s2">"train.csv"</span><span class="p">,</span> <span class="n">train_records</span><span class="p">)</span>
    <span class="n">write_csv!</span><span class="p">(</span><span class="s2">"validation.csv"</span><span class="p">,</span> <span class="n">validation_records</span><span class="p">)</span>
    <span class="n">write_csv!</span><span class="p">(</span><span class="s2">"test.csv"</span><span class="p">,</span> <span class="n">test_records</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_records</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1_045</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">validation_records</span><span class="p">)</span> <span class="o">==</span> <span class="mi">149</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">test_records</span><span class="p">)</span> <span class="o">==</span> <span class="mi">300</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="s2">"train.csv"</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="s2">"validation.csv"</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="s2">"test.csv"</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">write_csv!</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="n">records</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">rows</span> <span class="o">=</span>
      <span class="n">records</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="n">record</span> <span class="o">-&gt;</span>
        <span class="p">[</span><span class="n">csv_field</span><span class="p">(</span><span class="n">record</span><span class="o">.</span><span class="n">label</span><span class="p">),</span> <span class="s2">","</span><span class="p">,</span> <span class="n">csv_field</span><span class="p">(</span><span class="n">record</span><span class="o">.</span><span class="n">text</span><span class="p">),</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="p">]</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="no">File</span><span class="o">.</span><span class="n">write!</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="p">[</span><span class="s2">"Label,Text</span><span class="se">\n</span><span class="s2">"</span><span class="p">,</span> <span class="n">rows</span><span class="p">])</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">csv_field</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">value</span>
    <span class="o">|&gt;</span> <span class="n">to_string</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">"</span><span class="se">\"</span><span class="s2">"</span><span class="p">,</span> <span class="s2">"</span><span class="se">\"\"</span><span class="s2">"</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="s2">"</span><span class="se">\"</span><span class="si">#{</span><span class="nv">&amp;1</span><span class="si">}</span><span class="se">\"</span><span class="s2">"</span><span class="p">)</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<h2 id="63-creating-data-loaders">6.3 Creating data loaders</h2>

<p>For pretraining our <code class="language-plaintext highlighter-rouge">GPTModel</code>, we prepared data with <code class="language-plaintext highlighter-rouge">LlmScratch.GptDatasetV1</code>. The input text is split into fixed-size chunks, and each chunk contains an input part and a target part. We use sliding-window and stride parameters to create the target part. After that, we organize chunks into batches. Here we have messages of different lengths, so we need a different technique. We could truncate all messages to the length of the shortest message, or pad all messages to the length of the longest message. Since truncating to the shortest message could lose important training information, we use padding.</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>Input text message              1. Tokenize texts                 2. Pad to longest sequence
------------------              -----------------                 --------------------------

This is the first text   ---&gt;   1212, 318, 262, 717,       ---&gt;   1212, 318, 262, 717,
message                         2420, 3275                        2420, 3275, 50256, 50256,
                                                                  50256, 50256, 50256

This is another text     ---&gt;   1212, 318, 1194,           ---&gt;   1212, 318, 1194, 2420,
message                         2420, 3275                        3275, 50256, 50256, 50256,
                                                                  50256, 50256, 50256

This is the third text   ---&gt;   1212, 318, 262, 2368,      ---&gt;   1212, 318, 262, 2368,
message, which is               2420, 3275, 11, 543,              2420, 3275, 11, 543,
very long                       318, 845, 890                     318, 845, 890

                                                                  ^ no padding is needed here
                                                                    because this message already
                                                                    has the longest token sequence

                                      Padding token ID: 50256
</code></pre></div></div>

<p>Shorter messages are padded with the token ID <code class="language-plaintext highlighter-rouge">50256</code> (&lt;|endoftext|&gt;) until every message in the batch has the same sequence length as the longest message.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">SpamDataset</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  CSV-backed spam classification dataset.

  This mirrors the PyTorch `SpamDataset` example from chapter 6. A dataset is
  built from a CSV file containing `Label` and `Text` columns, pre-tokenizes all
  texts during construction, truncates sequences that exceed the configured
  maximum length, and pads shorter sequences with a fixed token id.
  """</span>

  <span class="nv">@enforce_keys</span> <span class="p">[</span><span class="ss">:data</span><span class="p">,</span> <span class="ss">:encoded_texts</span><span class="p">,</span> <span class="ss">:max_length</span><span class="p">,</span> <span class="ss">:pad_token_id</span><span class="p">]</span>
  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:data</span><span class="p">,</span> <span class="ss">:encoded_texts</span><span class="p">,</span> <span class="ss">:max_length</span><span class="p">,</span> <span class="ss">:pad_token_id</span><span class="p">]</span>

  <span class="nv">@end_of_text</span> <span class="s2">"&lt;|endoftext|&gt;"</span>
  <span class="nv">@pad_token_id</span> <span class="mi">50_256</span>

  <span class="nv">@type</span> <span class="n">row</span> <span class="p">::</span> <span class="p">%{</span><span class="ss">label:</span> <span class="n">integer</span><span class="p">(),</span> <span class="ss">text:</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()}</span>
  <span class="nv">@type</span> <span class="n">tokenizer</span> <span class="p">::</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span> <span class="o">|</span> <span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="p">[</span><span class="n">integer</span><span class="p">()])</span>
  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">data:</span> <span class="p">[</span><span class="n">row</span><span class="p">()],</span>
          <span class="ss">encoded_texts:</span> <span class="p">[[</span><span class="n">integer</span><span class="p">()]],</span>
          <span class="ss">max_length:</span> <span class="n">non_neg_integer</span><span class="p">(),</span>
          <span class="ss">pad_token_id:</span> <span class="n">integer</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Creates a spam dataset from a CSV file.

  ## Input Parameters

    * `csv_file` - path to a CSV file with a header row containing `Label` and
      `Text` columns. Labels may be integer strings (`"0"`, `"1"`) or the SMS
      spam labels (`"ham"`, `"spam"`).
    * `tokenizer` - either a Tiktoken model name accepted by `Tiktoken.encode/3`
      or a one-argument function that receives text and returns a list of token
      ids.
    * `opts` - optional keyword list controlling sequence length and padding.

  ## Options

    * `:max_length` - fixed sequence length. When omitted, the longest
      pre-tokenized sequence length in the CSV is used.
    * `:pad_token_id` - token id appended to shorter sequences. Defaults to
      GPT-2's end-of-text token id, `50256`.
    * `:allowed_special` - special tokens allowed when `tokenizer` is a
      Tiktoken model name. Defaults to `["&lt;|endoftext|&gt;"]`.

  ## Output

  Returns a `%LlmScratch.SpamDataset{}`. The `:data` field contains parsed
  `%{label: integer, text: text}` rows, and `:encoded_texts` contains one
  padded/truncated token-id list per row.
  """</span>
  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">tokenizer</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">csv_file</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">csv_file</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">data</span> <span class="o">=</span> <span class="n">read_csv!</span><span class="p">(</span><span class="n">csv_file</span><span class="p">)</span>
    <span class="n">encoded_texts</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">encode_text!</span><span class="p">(</span><span class="nv">&amp;1</span><span class="o">.</span><span class="n">text</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">))</span>
    <span class="n">max_length</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:max_length</span><span class="p">)</span> <span class="o">||</span> <span class="n">longest_length</span><span class="p">(</span><span class="n">encoded_texts</span><span class="p">)</span>
    <span class="n">pad_token_id</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:pad_token_id</span><span class="p">,</span> <span class="nv">@pad_token_id</span><span class="p">)</span>

    <span class="n">encoded_texts</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">encoded_texts</span><span class="p">,</span> <span class="k">fn</span> <span class="n">encoded_text</span> <span class="o">-&gt;</span>
        <span class="n">encoded_text</span>
        <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">max_length</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="n">pad_to_length</span><span class="p">(</span><span class="n">max_length</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">data:</span> <span class="n">data</span><span class="p">,</span>
      <span class="ss">encoded_texts:</span> <span class="n">encoded_texts</span><span class="p">,</span>
      <span class="ss">max_length:</span> <span class="n">max_length</span><span class="p">,</span>
      <span class="ss">pad_token_id:</span> <span class="n">pad_token_id</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns one encoded text and label pair by index.

  ## Input Parameters

    * `dataset` - `%LlmScratch.SpamDataset{}` returned by `new/3`.
    * `index` - zero-based row index.

  ## Output

  Returns `{encoded, label}` where `encoded` is an Nx tensor of shape
  `{dataset.max_length}` and signed 64-bit integer type, and `label` is a scalar
  signed 64-bit integer Nx tensor.
  """</span>
  <span class="nv">@spec</span> <span class="n">get</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="n">non_neg_integer</span><span class="p">())</span> <span class="p">::</span> <span class="p">{</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()}</span>
  <span class="k">def</span> <span class="n">get</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">dataset</span><span class="p">,</span> <span class="n">index</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">index</span><span class="p">)</span> <span class="ow">and</span> <span class="n">index</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="n">encoded</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">,</span> <span class="n">index</span><span class="p">)</span>
    <span class="n">label</span> <span class="o">=</span> <span class="n">dataset</span><span class="o">.</span><span class="n">data</span> <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">index</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="ss">:label</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">encoded</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}),</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">label</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns the number of rows in the dataset.

  ## Input Parameters

    * `dataset` - `%LlmScratch.SpamDataset{}` returned by `new/3`.

  ## Output

  Returns the number of CSV data rows as a non-negative integer.
  """</span>
  <span class="nv">@spec</span> <span class="n">length</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">length</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">dataset</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">dataset</span><span class="o">.</span><span class="n">data</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns the longest stored encoded sequence length.

  ## Input Parameters

    * `dataset` - `%LlmScratch.SpamDataset{}` returned by `new/3`.

  ## Output

  Returns the maximum token count among `dataset.encoded_texts`. Because
  `new/3` stores padded/truncated sequences, this matches `dataset.max_length`
  for datasets created by this module.
  """</span>
  <span class="nv">@spec</span> <span class="n">longest_encoded_length</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">longest_encoded_length</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">dataset</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">longest_length</span><span class="p">(</span><span class="n">dataset</span><span class="o">.</span><span class="n">encoded_texts</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">encode_text!</span><span class="p">(</span><span class="n">text</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">_opts</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_function</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">tokenizer</span><span class="o">.</span><span class="p">(</span><span class="n">text</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">encode_text!</span><span class="p">(</span><span class="n">text</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">allowed_special</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:allowed_special</span><span class="p">,</span> <span class="p">[</span><span class="nv">@end_of_text</span><span class="p">])</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">token_ids</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="n">text</span><span class="p">,</span> <span class="n">allowed_special</span><span class="p">)</span>
    <span class="n">token_ids</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">longest_length</span><span class="p">(</span><span class="n">encoded_texts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">encoded_texts</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="mi">0</span> <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">pad_to_length</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">max_length</span><span class="p">,</span> <span class="n">pad_token_id</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">token_ids</span> <span class="o">++</span> <span class="no">List</span><span class="o">.</span><span class="n">duplicate</span><span class="p">(</span><span class="n">pad_token_id</span><span class="p">,</span> <span class="n">max_length</span> <span class="o">-</span> <span class="no">Kernel</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">token_ids</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">read_csv!</span><span class="p">(</span><span class="n">csv_file</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">csv_file</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">stream!</span><span class="p">([],</span> <span class="ss">:line</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Stream</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="no">String</span><span class="o">.</span><span class="n">trim_trailing</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Stream</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="no">String</span><span class="o">.</span><span class="n">trim_trailing</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="s2">"</span><span class="se">\r</span><span class="s2">"</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">to_list</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">parse_csv_rows!</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_csv_rows!</span><span class="p">([]),</span> <span class="k">do</span><span class="p">:</span> <span class="p">[]</span>

  <span class="k">defp</span> <span class="n">parse_csv_rows!</span><span class="p">([</span><span class="n">header</span> <span class="o">|</span> <span class="n">rows</span><span class="p">])</span> <span class="k">do</span>
    <span class="n">header_indexes</span> <span class="o">=</span>
      <span class="n">header</span>
      <span class="o">|&gt;</span> <span class="n">parse_csv_line!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">with_index</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Map</span><span class="o">.</span><span class="n">new</span><span class="p">()</span>

    <span class="n">label_index</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">header_indexes</span><span class="p">,</span> <span class="s2">"Label"</span><span class="p">)</span>
    <span class="n">text_index</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">header_indexes</span><span class="p">,</span> <span class="s2">"Text"</span><span class="p">)</span>

    <span class="n">rows</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reject</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">==</span> <span class="s2">""</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="n">row</span> <span class="o">-&gt;</span>
      <span class="n">columns</span> <span class="o">=</span> <span class="n">parse_csv_line!</span><span class="p">(</span><span class="n">row</span><span class="p">)</span>

      <span class="p">%{</span>
        <span class="ss">label:</span> <span class="n">columns</span> <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">label_index</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="n">parse_label!</span><span class="p">(),</span>
        <span class="ss">text:</span> <span class="no">Enum</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">columns</span><span class="p">,</span> <span class="n">text_index</span><span class="p">)</span>
      <span class="p">}</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_label!</span><span class="p">(</span><span class="s2">"ham"</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="mi">0</span>
  <span class="k">defp</span> <span class="n">parse_label!</span><span class="p">(</span><span class="s2">"spam"</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="mi">1</span>

  <span class="k">defp</span> <span class="n">parse_label!</span><span class="p">(</span><span class="n">label</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Integer</span><span class="o">.</span><span class="n">parse</span><span class="p">(</span><span class="n">label</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">label_id</span><span class="p">,</span> <span class="s2">""</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="n">label_id</span>

      <span class="n">_</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected integer, ham, or spam label, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">label</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_csv_line!</span><span class="p">(</span><span class="n">line</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">line</span>
    <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">to_charlist</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">parse_csv_chars</span><span class="p">([],</span> <span class="p">[],</span> <span class="no">false</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">to_string</span><span class="p">()))</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_csv_chars</span><span class="p">([],</span> <span class="n">field</span><span class="p">,</span> <span class="n">fields</span><span class="p">,</span> <span class="no">false</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="p">[</span><span class="n">field</span> <span class="o">|</span> <span class="n">fields</span><span class="p">]</span>

  <span class="k">defp</span> <span class="n">parse_csv_chars</span><span class="p">([],</span> <span class="n">_field</span><span class="p">,</span> <span class="n">_fields</span><span class="p">,</span> <span class="no">true</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="k">raise</span><span class="p">(</span><span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"unterminated CSV quote"</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">parse_csv_chars</span><span class="p">([</span><span class="sx">?"</span><span class="p">,</span> <span class="sx">?"</span> <span class="o">|</span> <span class="n">rest</span><span class="p">],</span> <span class="n">field</span><span class="p">,</span> <span class="n">fields</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">parse_csv_chars</span><span class="p">(</span><span class="n">rest</span><span class="p">,</span> <span class="p">[</span><span class="sx">?"</span> <span class="o">|</span> <span class="n">field</span><span class="p">],</span> <span class="n">fields</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_csv_chars</span><span class="p">([</span><span class="sx">?"</span> <span class="o">|</span> <span class="n">rest</span><span class="p">],</span> <span class="n">field</span><span class="p">,</span> <span class="n">fields</span><span class="p">,</span> <span class="n">in_quotes</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">parse_csv_chars</span><span class="p">(</span><span class="n">rest</span><span class="p">,</span> <span class="n">field</span><span class="p">,</span> <span class="n">fields</span><span class="p">,</span> <span class="ow">not</span> <span class="n">in_quotes</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_csv_chars</span><span class="p">([</span><span class="sx">?,</span> <span class="o">|</span> <span class="n">rest</span><span class="p">],</span> <span class="n">field</span><span class="p">,</span> <span class="n">fields</span><span class="p">,</span> <span class="no">false</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">parse_csv_chars</span><span class="p">(</span><span class="n">rest</span><span class="p">,</span> <span class="p">[],</span> <span class="p">[</span><span class="n">field</span> <span class="o">|</span> <span class="n">fields</span><span class="p">],</span> <span class="no">false</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">parse_csv_chars</span><span class="p">([</span><span class="n">char</span> <span class="o">|</span> <span class="n">rest</span><span class="p">],</span> <span class="n">field</span><span class="p">,</span> <span class="n">fields</span><span class="p">,</span> <span class="n">in_quotes</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">parse_csv_chars</span><span class="p">(</span><span class="n">rest</span><span class="p">,</span> <span class="p">[</span><span class="n">char</span> <span class="o">|</span> <span class="n">field</span><span class="p">],</span> <span class="n">fields</span><span class="p">,</span> <span class="n">in_quotes</span><span class="p">)</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And here is a test that uses <code class="language-plaintext highlighter-rouge">SpamDataset</code> with our existing <code class="language-plaintext highlighter-rouge">DataLoader</code>:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"6.3 creating data loaders"</span> <span class="k">do</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">encoded_tokens</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="s2">"&lt;|endoftext|&gt;"</span><span class="p">,</span> <span class="p">[</span><span class="s2">"&lt;|endoftext|&gt;"</span><span class="p">])</span>

    <span class="n">assert</span> <span class="n">encoded_tokens</span> <span class="o">==</span> <span class="p">[</span><span class="mi">50_256</span><span class="p">]</span>

    <span class="n">train_dataset</span> <span class="o">=</span> <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"train.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="no">nil</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span> <span class="o">==</span> <span class="mi">120</span>

    <span class="n">val_dataset</span> <span class="o">=</span>
      <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"validation.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span><span class="p">)</span>

    <span class="n">test_dataset</span> <span class="o">=</span>
      <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"test.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span><span class="p">)</span>

    <span class="n">num_workers</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="n">batch_size</span> <span class="o">=</span> <span class="mi">8</span>
    <span class="ss">:rand</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="ss">:exsss</span><span class="p">,</span> <span class="p">{</span><span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">})</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="n">train_dataset</span>
      <span class="o">|&gt;</span> <span class="n">dataset_samples</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">true</span>
      <span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="n">val_dataset</span>
      <span class="o">|&gt;</span> <span class="n">dataset_samples</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span>
      <span class="p">)</span>

    <span class="n">test_loader</span> <span class="o">=</span>
      <span class="n">test_dataset</span>
      <span class="o">|&gt;</span> <span class="n">dataset_samples</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span>
        <span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="n">num_workers</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span>
      <span class="p">)</span>

    <span class="p">{</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">}</span> <span class="o">=</span>
      <span class="n">train_loader</span><span class="o">.</span><span class="n">batches</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="n">collate_batch</span><span class="p">()</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">input_batch</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">8</span><span class="p">,</span> <span class="mi">120</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">target_batch</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">8</span><span class="p">}</span>

    <span class="n">assert</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">==</span> <span class="mi">130</span>
    <span class="n">assert</span> <span class="n">val_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">==</span> <span class="mi">19</span>
    <span class="n">assert</span> <span class="n">test_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">==</span> <span class="mi">38</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">dataset_samples</span><span class="p">(</span><span class="n">dataset</span><span class="p">)</span> <span class="k">do</span>
    <span class="mi">0</span><span class="o">..</span><span class="p">(</span><span class="no">SpamDataset</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">dataset</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="no">SpamDataset</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="nv">&amp;1</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">collate_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">labels</span><span class="p">}</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">unzip</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">inputs</span><span class="p">),</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">labels</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Each batch has eight messages (this is batch_size), and each message has 120 tokens (we padded them to max message length). This is input data. Each target has eight (again batch_size value) class for each of eight messages, with values 0 or 1 for ham or spam.</p>

<h2 id="64-initializing-a-model-with-pretrained-weights">6.4 Initializing a model with pretrained weights</h2>

<p>We first need to load OpenAI GPT2 small model:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">900_000</span>
  <span class="n">test</span> <span class="s2">"6.4 loads OpenAI GPT-2 and generates classification prompts"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">choose_model</span> <span class="o">=</span> <span class="s2">"gpt2-small (124M)"</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>
    <span class="n">config</span> <span class="o">=</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">openai_gpt2</span><span class="p">(</span><span class="n">choose_model</span><span class="p">)</span>

    <span class="n">model_size</span> <span class="o">=</span>
      <span class="n">choose_model</span>
      <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">" "</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">last</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">trim_leading</span><span class="p">(</span><span class="s2">"("</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">trim_trailing</span><span class="p">(</span><span class="s2">")"</span><span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="n">model_size</span><span class="p">,</span> <span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">config</span><span class="o">.</span><span class="n">vocab_size</span> <span class="o">==</span> <span class="mi">50_257</span>
    <span class="n">assert</span> <span class="n">config</span><span class="o">.</span><span class="n">context_length</span> <span class="o">==</span> <span class="mi">1024</span>
    <span class="n">assert</span> <span class="n">config</span><span class="o">.</span><span class="n">drop_rate</span> <span class="o">==</span> <span class="mf">0.0</span>
    <span class="n">assert</span> <span class="n">config</span><span class="o">.</span><span class="n">qkv_bias</span> <span class="o">==</span> <span class="no">true</span>
    <span class="n">assert</span> <span class="n">config</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">==</span> <span class="mi">768</span>
    <span class="n">assert</span> <span class="n">config</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">==</span> <span class="mi">12</span>
    <span class="n">assert</span> <span class="n">config</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">==</span> <span class="mi">12</span>

    <span class="n">decoded_text_1</span> <span class="o">=</span>
      <span class="n">model</span>
      <span class="o">|&gt;</span> <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate_text_simple</span><span class="p">(</span>
        <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="s2">"Every effort moves you"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">),</span>
        <span class="mi">15</span><span class="p">,</span>
        <span class="n">config</span><span class="o">.</span><span class="n">context_length</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">decoded_text_1</span> <span class="o">==</span>
             <span class="s2">"Every effort moves you forward.</span><span class="se">\n\n</span><span class="s2">The first step is to understand the importance of your work"</span>

    <span class="n">text_2</span> <span class="o">=</span>
      <span class="s2">"Is the following text 'spam'? Answer with 'yes' or 'no':"</span> <span class="o">&lt;&gt;</span>
        <span class="s2">" 'You are a winner you have been specially"</span> <span class="o">&lt;&gt;</span>
        <span class="s2">" selected to receive $1000 cash or a $2000 award.'"</span>

    <span class="n">decoded_text_2</span> <span class="o">=</span>
      <span class="n">model</span>
      <span class="o">|&gt;</span> <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate_text_simple</span><span class="p">(</span>
        <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">text_2</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">),</span>
        <span class="mi">23</span><span class="p">,</span>
        <span class="n">config</span><span class="o">.</span><span class="n">context_length</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">decoded_text_2</span> <span class="o">==</span>
             <span class="s2">"Is the following text 'spam'? Answer with 'yes' or 'no': 'You are a winner you have been specially selected to receive $1000 cash or a $2000 award.'</span><span class="se">\n\n</span><span class="s2">The following text 'spam'? Answer with 'yes' or 'no': 'You are a winner"</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>The model fails to classify one message because we have not trained it for this classification task yet.</p>

<h2 id="65-adding-a-classification-head">6.5 Adding a classification head</h2>

<p>To classify text with our pretrained <code class="language-plaintext highlighter-rouge">GPT2Model</code>, we need to replace its output layer. The output layer maps the hidden representation to a vocabulary of 50,257 tokens. In our ham/spam classifier, the output size is only two. This means we can train the model on the ham/spam classification dataset by updating only a small part of the model: the output head, the last transformer block, and the final normalization layer. Sebastian adds the last transformer layer and final normalization layer based on his practical experience training classification models. <code class="language-plaintext highlighter-rouge">PyTorch</code> has a nice feature for printing the model architecture. Here is an Elixir implementation of the <code class="language-plaintext highlighter-rouge">Inspect</code> protocol for the <code class="language-plaintext highlighter-rouge">GPTModel</code> module:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="k">defimpl</span> <span class="no">Inspect</span><span class="p">,</span> <span class="ss">for:</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">GPTModel</span> <span class="k">do</span>
  <span class="kn">import</span> <span class="no">Inspect</span><span class="o">.</span><span class="no">Algebra</span>

  <span class="k">def</span> <span class="n">inspect</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">_opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">model</span>
    <span class="o">|&gt;</span> <span class="n">inspect_model</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">string</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">inspect_model</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">lines</span> <span class="o">=</span> <span class="p">[</span>
      <span class="s2">"GPTModel("</span><span class="p">,</span>
      <span class="s2">"  (tok_emb): </span><span class="si">#{</span><span class="n">embedding</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">tok_emb</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"  (pos_emb): </span><span class="si">#{</span><span class="n">embedding</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"  (drop_emb): </span><span class="si">#{</span><span class="n">dropout</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">drop_emb</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"  (trf_blocks): Sequential("</span>
    <span class="p">]</span>

    <span class="n">block_lines</span> <span class="o">=</span>
      <span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">with_index</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">flat_map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">block</span><span class="p">,</span> <span class="n">index</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="n">transformer_block</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">index</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">join</span><span class="p">(</span>
      <span class="n">lines</span> <span class="o">++</span>
        <span class="n">block_lines</span> <span class="o">++</span>
        <span class="p">[</span><span class="s2">"  )"</span><span class="p">,</span> <span class="s2">"  (final_norm): LayerNorm()"</span><span class="p">,</span> <span class="s2">"  (out_head): </span><span class="si">#{</span><span class="n">out_head</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span> <span class="s2">")"</span><span class="p">],</span>
      <span class="s2">"</span><span class="se">\n</span><span class="s2">"</span>
    <span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">transformer_block</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">index</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">[</span>
      <span class="s2">"    (</span><span class="si">#{</span><span class="n">index</span><span class="si">}</span><span class="s2">): TransformerBlock("</span><span class="p">,</span>
      <span class="s2">"      (att): MultiHeadAttention("</span><span class="p">,</span>
      <span class="s2">"        (W_query): </span><span class="si">#{</span><span class="n">linear</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">w_q</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"        (W_key): </span><span class="si">#{</span><span class="n">linear</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">w_k</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"        (W_value): </span><span class="si">#{</span><span class="n">linear</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">w_v</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"        (out_proj): </span><span class="si">#{</span><span class="n">linear</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">out_proj</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"        (dropout): </span><span class="si">#{</span><span class="n">dropout</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">dropout</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"      )"</span><span class="p">,</span>
      <span class="s2">"      (ff): FeedForward("</span><span class="p">,</span>
      <span class="s2">"        (layers): Sequential("</span><span class="p">,</span>
      <span class="s2">"          (0): </span><span class="si">#{</span><span class="n">linear</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"          (1): GELU()"</span><span class="p">,</span>
      <span class="s2">"          (2): </span><span class="si">#{</span><span class="n">linear</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"        )"</span><span class="p">,</span>
      <span class="s2">"      )"</span><span class="p">,</span>
      <span class="s2">"      (norm1): LayerNorm()"</span><span class="p">,</span>
      <span class="s2">"      (norm2): LayerNorm()"</span><span class="p">,</span>
      <span class="s2">"      (drop_resid): </span><span class="si">#{</span><span class="n">dropout</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">drop_shortcut</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
      <span class="s2">"    )"</span>
    <span class="p">]</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">embedding</span><span class="p">(</span><span class="n">embedding</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">vocab_size</span><span class="p">,</span> <span class="n">embedding_dim</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">embedding</span><span class="o">.</span><span class="n">weight</span><span class="p">)</span>
    <span class="s2">"Embedding(</span><span class="si">#{</span><span class="n">vocab_size</span><span class="si">}</span><span class="s2">, </span><span class="si">#{</span><span class="n">embedding_dim</span><span class="si">}</span><span class="s2">)"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">linear</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}</span> <span class="o">=</span> <span class="n">layer</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">in_features</span><span class="p">,</span> <span class="n">out_features</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span>

    <span class="s2">"Linear(in_features=</span><span class="si">#{</span><span class="n">in_features</span><span class="si">}</span><span class="s2">, out_features=</span><span class="si">#{</span><span class="n">out_features</span><span class="si">}</span><span class="s2">, bias=</span><span class="si">#{</span><span class="no">Map</span><span class="o">.</span><span class="n">has_key?</span><span class="p">(</span><span class="n">layer</span><span class="p">,</span> <span class="ss">:bias</span><span class="p">)</span><span class="si">}</span><span class="s2">)"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">out_head</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}</span> <span class="o">=</span> <span class="n">layer</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">in_features</span><span class="p">,</span> <span class="n">out_features</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span>

    <span class="s2">"Linear(in_features=</span><span class="si">#{</span><span class="n">in_features</span><span class="si">}</span><span class="s2">, out_features=</span><span class="si">#{</span><span class="n">out_features</span><span class="si">}</span><span class="s2">, bias=</span><span class="si">#{</span><span class="no">Map</span><span class="o">.</span><span class="n">has_key?</span><span class="p">(</span><span class="n">layer</span><span class="p">,</span> <span class="ss">:bias</span><span class="p">)</span><span class="si">}</span><span class="s2">)"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">dropout</span><span class="p">(</span><span class="n">dropout</span><span class="p">)</span> <span class="k">do</span>
    <span class="s2">"Dropout(p=</span><span class="si">#{</span><span class="n">dropout</span> <span class="o">*</span> <span class="mf">1.0</span><span class="si">}</span><span class="s2">, inplace=false)"</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>PyTorch has a feature for freezing model blocks. When a block is frozen, it is skipped during training. To support this in Elixir, we need to change the <code class="language-plaintext highlighter-rouge">GPTModel</code> and <code class="language-plaintext highlighter-rouge">Training</code> modules. Here is the diff:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="gh">diff --git a/lib/llm_scratch/gpt_model.ex b/lib/llm_scratch/gpt_model.ex
index a6e5b25..1da835a 100644
</span><span class="gd">--- a/lib/llm_scratch/gpt_model.ex
</span><span class="gi">+++ b/lib/llm_scratch/gpt_model.ex
</span><span class="p">@@ -26,7 +26,18 @@</span> defmodule LlmScratch.GPTModel do
     TransformerBlock
   }
 
<span class="gd">-  defstruct [:cfg, :tok_emb, :pos_emb, :drop_emb, :trf_blocks, :final_norm, :out_head]
</span><span class="gi">+  @trainable_fields [:tok_emb, :pos_emb, :trf_blocks, :final_norm, :out_head]
+
+  defstruct [
+    :cfg,
+    :tok_emb,
+    :pos_emb,
+    :drop_emb,
+    :trf_blocks,
+    :final_norm,
+    :out_head,
+    trainable: :all
+  ]
</span> 
   @type linear_no_bias :: %{kernel: Nx.Tensor.t()}
 
<span class="p">@@ -37,7 +48,8 @@</span> defmodule LlmScratch.GPTModel do
           drop_emb: float(),
           trf_blocks: [TransformerBlock.t()],
           final_norm: DummyLayerNorm.t(),
<span class="gd">-          out_head: linear_no_bias()
</span><span class="gi">+          out_head: linear_no_bias(),
+          trainable: :all | [atom() | {:trf_block, non_neg_integer()}]
</span>         }
 
   @spec new(GPTConfig.t(), keyword()) :: t()
<span class="p">@@ -62,7 +74,113 @@</span> defmodule LlmScratch.GPTModel do
       drop_emb: GPTConfig.embedding_dropout(cfg),
       trf_blocks: transformer_blocks(cfg, seed + 2, norm_eps),
       final_norm: DummyLayerNorm.new(cfg.emb_dim, eps: norm_eps),
<span class="gd">-      out_head: init_out_head(cfg, seed + 2 + cfg.n_layers * 8)
</span><span class="gi">+      out_head: init_out_head(cfg, seed + 2 + cfg.n_layers * 8),
+      trainable: :all
+    }
+  end
+
+  @doc """
+  Freezes all model parameters for fine-tuning workflows.
+
+  Nx tensors do not have a mutable `requires_grad` flag like PyTorch tensors.
+  Instead, this marks the model with an empty trainable layer list. The training
+  optimizer respects that metadata and leaves all existing parameters
+  unchanged.
+
+  Returns an updated `%LlmScratch.GPTModel{}`.
+  """
+  @spec freeze(t()) :: t()
+  def freeze(%__MODULE__{} = model), do: %{model | trainable: []}
+
+  @doc """
+  Restores all model parameters to trainable.
+
+  Returns an updated `%LlmScratch.GPTModel{}` with the default `:all`
+  trainable setting.
+  """
+  @spec unfreeze(t()) :: t()
+  def unfreeze(%__MODULE__{} = model), do: %{model | trainable: :all}
+
+  @doc """
+  Marks selected top-level model layers as trainable.
+
+  ## Parameters
+
+    * `model` - `%LlmScratch.GPTModel{}` to update.
+    * `fields` - list of trainable top-level layer names. Accepted fields are
+      `:tok_emb`, `:pos_emb`, `:trf_blocks`, `:final_norm`, and `:out_head`.
+
+  This is useful after calling `freeze/1`, for example to train only a newly
+  added classification head.
+
+  Returns an updated `%LlmScratch.GPTModel{}`.
+  """
+  @spec set_trainable(t(), [atom()]) :: t()
+  def set_trainable(%__MODULE__{} = model, fields) when is_list(fields) do
+    unknown_fields =
+      Enum.reject(fields, fn
+        field when field in @trainable_fields -&gt;
+          true
+
+        {:trf_block, index} when is_integer(index) and index &gt;= 0 -&gt;
+          index &lt; length(model.trf_blocks)
+
+        _field -&gt;
+          false
+      end)
+
+    if unknown_fields != [] do
+      raise ArgumentError,
+            "unknown trainable GPTModel fields #{inspect(unknown_fields)}; expected one or more of #{inspect(@trainable_fields)} or {:trf_block, index}"
+    end
+
+    %{model | trainable: Enum.uniq(fields)}
+  end
+
+  @doc """
+  Returns `true` when no model parameters are trainable.
+  """
+  @spec frozen?(t()) :: boolean()
+  def frozen?(%__MODULE__{trainable: []}), do: true
+  def frozen?(%__MODULE__{}), do: false
+
+  @doc """
+  Replaces the output head with a dense classification or projection head.
+
+  ## Parameters
+
+    * `model` - `%LlmScratch.GPTModel{}` to update.
+    * `out_features` - number of output logits produced for each token.
+    * `opts` - optional keyword list.
+
+  ## Options
+
+    * `:seed` - deterministic initialization seed. Defaults to `123`.
+    * `:bias` - whether to include a bias vector. Defaults to `true`, matching
+      `torch.nn.Linear/2`.
+
+  ## Output
+
+  Returns an updated model whose `:out_head` has kernel shape
+  `{model.cfg.emb_dim, out_features}` and, when enabled, bias shape
+  `{out_features}`.
+  """
+  @spec replace_out_head(t(), pos_integer(), keyword()) :: t()
+  def replace_out_head(%__MODULE__{} = model, out_features, opts \\ [])
+      when is_integer(out_features) and out_features &gt; 0 and is_list(opts) do
+    seed = Keyword.get(opts, :seed, 123)
+    bias = Keyword.get(opts, :bias, true)
+
+    %{
+      model
+      | out_head:
+          SelfAttentionV2.init_dense_weights(
+            model.cfg.emb_dim,
+            out_features,
+            seed,
+            bias,
+            "out_head"
+          )
</span>     }
   end
 
<span class="p">@@ -215,10 +333,17 @@</span> defmodule LlmScratch.GPTModel do
     %{kernel: kernel}
   end
 
<span class="gi">+  defp linear(x, %{kernel: kernel, bias: bias}), do: Nx.dot(x, [-1], kernel, [0]) |&gt; Nx.add(bias)
</span>   defp linear(x, %{kernel: kernel}), do: Nx.dot(x, [-1], kernel, [0])
 
<span class="gd">-  defnp linear_defn(x, %{kernel: kernel}) do
-    Nx.dot(x, [-1], kernel, [0])
</span><span class="gi">+  deftransformp linear_defn(x, layer) do
+    y = Nx.dot(x, [-1], layer.kernel, [0])
+
+    if Map.has_key?(layer, :bias) do
+      Nx.add(y, layer.bias)
+    else
+      y
+    end
</span>   end
</code></pre></div></div>

<p>And <code class="language-plaintext highlighter-rouge">Training</code> module:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="gh">diff --git a/lib/llm_scratch/training.ex b/lib/llm_scratch/training.ex
index 35453ba..88307a9 100644
</span><span class="gd">--- a/lib/llm_scratch/training.ex
</span><span class="gi">+++ b/lib/llm_scratch/training.ex
</span><span class="p">@@ -128,7 +128,7 @@</span> end
 
 defimpl Nx.Container, for: LlmScratch.GPTModel do
   @container_fields [:tok_emb, :pos_emb, :trf_blocks, :final_norm, :out_head]
<span class="gd">-  @metadata_fields [:cfg, :drop_emb]
</span><span class="gi">+  @metadata_fields [:cfg, :drop_emb, :trainable]
</span> 
   def traverse(model, acc, fun) do
     {tok_emb, acc} = fun.(model.tok_emb, acc)
<span class="p">@@ -674,7 +674,7 @@</span> defmodule LlmScratch.Training do
       trainable_tensors(block.norm2)
   end
 
<span class="gd">-  defp trainable_tensors(%GPTModel{} = model) do
</span><span class="gi">+  defp trainable_tensors(%GPTModel{trainable: :all} = model) do
</span>     trainable_tensors(model.tok_emb) ++
       trainable_tensors(model.pos_emb) ++
       Enum.flat_map(model.trf_blocks, &amp;trainable_tensors/1) ++
<span class="p">@@ -682,6 +682,25 @@</span> defmodule LlmScratch.Training do
       dense_tensors(model.out_head, false)
   end
 
<span class="gi">+  defp trainable_tensors(%GPTModel{trainable: trainable} = model) when is_list(trainable) do
+    trainable
+    |&gt; Enum.flat_map(&amp;gpt_model_field_tensors(model, &amp;1))
+  end
+
+  defp gpt_model_field_tensors(model, :tok_emb), do: trainable_tensors(model.tok_emb)
+  defp gpt_model_field_tensors(model, :pos_emb), do: trainable_tensors(model.pos_emb)
+
+  defp gpt_model_field_tensors(model, :trf_blocks),
+    do: Enum.flat_map(model.trf_blocks, &amp;trainable_tensors/1)
+
+  defp gpt_model_field_tensors(model, :final_norm), do: trainable_tensors(model.final_norm)
+
+  defp gpt_model_field_tensors(model, :out_head),
+    do: dense_tensors(model.out_head, Map.has_key?(model.out_head, :bias))
+
+  defp gpt_model_field_tensors(model, {:trf_block, index}),
+    do: model.trf_blocks |&gt; Enum.at(index) |&gt; trainable_tensors()
+
</span>   defp dense_tensors(%{kernel: kernel, bias: bias}, true), do: [kernel, bias]
   defp dense_tensors(%{kernel: kernel}, false), do: [kernel]
 
<span class="p">@@ -729,7 +748,7 @@</span> defmodule LlmScratch.Training do
     {%{block | att: att, ff: ff, norm1: norm1, norm2: norm2}, tensors}
   end
 
<span class="gd">-  defp put_trainable_tensors(%GPTModel{} = model, tensors) do
</span><span class="gi">+  defp put_trainable_tensors(%GPTModel{trainable: :all} = model, tensors) do
</span>     {tok_emb, tensors} = put_trainable_tensors(model.tok_emb, tensors)
     {pos_emb, tensors} = put_trainable_tensors(model.pos_emb, tensors)
     {trf_blocks, tensors} = Enum.map_reduce(model.trf_blocks, tensors, &amp;put_trainable_tensors/2)
<span class="p">@@ -749,6 +768,49 @@</span> defmodule LlmScratch.Training do
     }
   end
 
<span class="gi">+  defp put_trainable_tensors(%GPTModel{trainable: trainable} = model, tensors)
+       when is_list(trainable) do
+    Enum.reduce(trainable, {model, tensors}, fn field, {model, tensors} -&gt;
+      put_gpt_model_field_tensors(model, field, tensors)
+    end)
+  end
+
+  defp put_gpt_model_field_tensors(model, :tok_emb, tensors) do
+    {tok_emb, tensors} = put_trainable_tensors(model.tok_emb, tensors)
+    {%{model | tok_emb: tok_emb}, tensors}
+  end
+
+  defp put_gpt_model_field_tensors(model, :pos_emb, tensors) do
+    {pos_emb, tensors} = put_trainable_tensors(model.pos_emb, tensors)
+    {%{model | pos_emb: pos_emb}, tensors}
+  end
+
+  defp put_gpt_model_field_tensors(model, :trf_blocks, tensors) do
+    {trf_blocks, tensors} = Enum.map_reduce(model.trf_blocks, tensors, &amp;put_trainable_tensors/2)
+    {%{model | trf_blocks: trf_blocks}, tensors}
+  end
+
+  defp put_gpt_model_field_tensors(model, :final_norm, tensors) do
+    {final_norm, tensors} = put_trainable_tensors(model.final_norm, tensors)
+    {%{model | final_norm: final_norm}, tensors}
+  end
+
+  defp put_gpt_model_field_tensors(model, :out_head, tensors) do
+    {out_head, tensors} =
+      put_dense_tensors(model.out_head, Map.has_key?(model.out_head, :bias), tensors)
+
+    {%{model | out_head: out_head}, tensors}
+  end
+
+  defp put_gpt_model_field_tensors(model, {:trf_block, index}, tensors) do
+    {block, tensors} =
+      model.trf_blocks
+      |&gt; Enum.at(index)
+      |&gt; put_trainable_tensors(tensors)
+
+    {%{model | trf_blocks: List.replace_at(model.trf_blocks, index, block)}, tensors}
+  end
+
</span>  defp put_dense_tensors(%{kernel: _kernel, bias: _bias} = dense, true, [kernel, bias | rest]) do
    {%{dense | kernel: kernel, bias: bias}, rest}
  end
</code></pre></div></div>

<p>Here is the test that uses those new features and prepares the model for classification fine-tuning:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">900_000</span>
  <span class="n">test</span> <span class="s2">"6.5 inspects saved GPT-2 small like a PyTorch module tree"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="s2">"124M"</span><span class="p">,</span> <span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>
    <span class="n">inspected</span> <span class="o">=</span> <span class="n">inspect</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"GPTModel("</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(tok_emb): Embedding(50257, 768)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(pos_emb): Embedding(1024, 768)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(drop_emb): Dropout(p=0.0, inplace=false)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(trf_blocks): Sequential("</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(0): TransformerBlock("</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(11): TransformerBlock("</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(W_query): Linear(in_features=768, out_features=768, bias=true)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(W_key): Linear(in_features=768, out_features=768, bias=true)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(W_value): Linear(in_features=768, out_features=768, bias=true)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(out_proj): Linear(in_features=768, out_features=768, bias=true)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(0): Linear(in_features=768, out_features=3072, bias=true)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(1): GELU()"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(2): Linear(in_features=3072, out_features=768, bias=true)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(norm1): LayerNorm()"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(norm2): LayerNorm()"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(drop_resid): Dropout(p=0.0, inplace=false)"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(final_norm): LayerNorm()"</span>
    <span class="n">assert</span> <span class="n">inspected</span> <span class="o">=~</span> <span class="s2">"(out_head): Linear(in_features=768, out_features=50257, bias=false)"</span>

    <span class="n">assert</span> <span class="no">Regex</span><span class="o">.</span><span class="n">scan</span><span class="p">(</span><span class="sr">~r/\(\d+\): TransformerBlock\(/</span><span class="p">,</span> <span class="n">inspected</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="n">length</span><span class="p">()</span> <span class="o">==</span> <span class="mi">12</span>

    <span class="c1"># by frozing model blocks, we skip their training</span>
    <span class="n">frozen_model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">freeze</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">frozen?</span><span class="p">(</span><span class="n">frozen_model</span><span class="p">)</span>
    <span class="n">assert</span> <span class="n">frozen_model</span><span class="o">.</span><span class="n">trainable</span> <span class="o">==</span> <span class="p">[]</span>

    <span class="c1"># we change output block second dimension to two, as this we have two classes, ham/spam</span>
    <span class="c1"># we want to train final_norm block and last transformer block</span>
    <span class="c1"># Sebastian states that based on his experiments, we will get better results</span>
    <span class="n">num_classes</span> <span class="o">=</span> <span class="mi">2</span>
    <span class="n">classification_model</span> <span class="o">=</span>
      <span class="n">frozen_model</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">replace_out_head</span><span class="p">(</span><span class="n">num_classes</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">,</span> <span class="ss">bias:</span> <span class="no">true</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">set_trainable</span><span class="p">([</span><span class="ss">:out_head</span><span class="p">,</span> <span class="ss">:final_norm</span><span class="p">,</span> <span class="p">{</span><span class="ss">:trf_block</span><span class="p">,</span> <span class="mi">11</span><span class="p">}])</span>

    <span class="n">assert</span> <span class="n">classification_model</span><span class="o">.</span><span class="n">trainable</span> <span class="o">==</span> <span class="p">[</span><span class="ss">:out_head</span><span class="p">,</span> <span class="ss">:final_norm</span><span class="p">,</span> <span class="p">{</span><span class="ss">:trf_block</span><span class="p">,</span> <span class="mi">11</span><span class="p">}]</span>

    <span class="n">assert</span> <span class="n">inspect</span><span class="p">(</span><span class="n">classification_model</span><span class="p">)</span> <span class="o">=~</span>
             <span class="s2">"(out_head): Linear(in_features=768, out_features=2, bias=true)"</span>

    <span class="c1">#input message that we want to classify</span>
    <span class="n">inputs</span> <span class="o">=</span>
      <span class="s2">"Do you have time"</span>
      <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_list</span><span class="p">(</span><span class="n">inputs</span><span class="p">)</span> <span class="o">==</span> <span class="p">[[</span><span class="mi">5211</span><span class="p">,</span> <span class="mi">345</span><span class="p">,</span> <span class="mi">423</span><span class="p">,</span> <span class="mi">640</span><span class="p">]]</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">inputs</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">4</span><span class="p">}</span>

    <span class="n">outputs</span> <span class="o">=</span>
      <span class="n">classification_model</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">outputs</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">2</span><span class="p">}</span>

    <span class="n">expected_outputs</span> <span class="o">=</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span>
        <span class="p">[</span>
          <span class="p">[</span><span class="mf">0.41287035</span><span class="p">,</span> <span class="mf">1.2672385</span><span class="p">],</span>
          <span class="p">[</span><span class="o">-</span><span class="mf">3.6814282</span><span class="p">,</span> <span class="mf">4.7981105</span><span class="p">],</span>
          <span class="p">[</span><span class="o">-</span><span class="mf">3.9752169</span><span class="p">,</span> <span class="mf">4.006133</span><span class="p">],</span>
          <span class="p">[</span><span class="o">-</span><span class="mf">1.0933434</span><span class="p">,</span> <span class="mf">3.5170393</span><span class="p">]</span>
        <span class="p">]</span>
      <span class="p">])</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all_close</span><span class="p">(</span><span class="n">outputs</span><span class="p">,</span> <span class="n">expected_outputs</span><span class="p">,</span> <span class="ss">atol:</span> <span class="mf">1.0e-5</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span> <span class="o">==</span> <span class="mi">1</span>

    <span class="c1"># it is enough to do fine tune training only for last input token!</span>
    <span class="c1"># Why? Attention mechanism is using casual attention mask, where we hide from current token its next token</span>
    <span class="c1"># Because of that, only the last message token has attention weights of all previous token.</span>
    <span class="n">last_output_token</span> <span class="o">=</span> <span class="n">outputs</span><span class="p">[[</span><span class="o">..</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">..</span><span class="p">]]</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">last_output_token</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all_close</span><span class="p">(</span><span class="n">last_output_token</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([[</span><span class="o">-</span><span class="mf">1.0933434</span><span class="p">,</span> <span class="mf">3.5170393</span><span class="p">]]),</span> <span class="ss">atol:</span> <span class="mf">1.0e-5</span><span class="p">)</span>
           <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span> <span class="o">==</span> <span class="mi">1</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="66-calculating-the-classification-loss-and-accuracy">6.6 Calculating the classification loss and accuracy</h2>

<p>Next, we need model evaluation functions for training. These are different from the ones we used during model pretraining. The reason is that the classifier logits contain only 2 possible classes instead of 50,257 possible tokens. We also calculate loss and accuracy only from the final input token, because in a causal attention model the final token can attend to all previous tokens.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">LossClassificationUtils</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Classification evaluation helpers.

  These utilities are used after replacing a GPT language-model output head with
  a classification head. For sequence classification, only the logits from the
  final input token are used because that token can attend to all previous
  tokens through causal self-attention.
  """</span>

  <span class="nv">@type</span> <span class="n">device</span> <span class="p">::</span> <span class="no">nil</span> <span class="o">|</span> <span class="ss">:default</span> <span class="o">|</span> <span class="n">atom</span><span class="p">()</span> <span class="o">|</span> <span class="n">tuple</span><span class="p">()</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Calculates classification accuracy over batches from a data loader.

  This is the Nx/Elixir counterpart of the PyTorch loop:

      model.eval()
      logits = model(input_batch)[:, -1, :]
      predicted_labels = torch.argmax(logits, dim=-1)
      accuracy = correct_predictions / num_examples

  `data_loader` should be a map with `:stream` and `:length`, such as the value
  returned by `LlmScratch.DataLoader.new/2`. Each batch may be either:

    * `{input_batch, target_batch}` with already-stacked tensors
    * a list of `{input_tensor, target_tensor}` examples

  `model` must be a struct whose module exports `forward/2`. The model's
  ordinary forward path is used, which is the project's evaluation path for
  GPT-style models.

  `device` accepts an Nx backend, such as `EXLA.Backend` or
  `{EXLA.Backend, client: :cuda}`. Use `nil` or `:default` to keep tensors on
  their current backend.

  Pass `nil` for `num_batches` to evaluate one full pass through the loader.
  When an empty loader, or `num_batches: 0`, produces no examples, returns
  `:nan`.
  """</span>
  <span class="nv">@spec</span> <span class="n">calc_accuracy_loader</span><span class="p">(</span><span class="n">map</span><span class="p">(),</span> <span class="n">struct</span><span class="p">(),</span> <span class="n">device</span><span class="p">(),</span> <span class="no">nil</span> <span class="o">|</span> <span class="n">non_neg_integer</span><span class="p">())</span> <span class="p">::</span> <span class="n">float</span><span class="p">()</span> <span class="o">|</span> <span class="ss">:nan</span>
  <span class="k">def</span> <span class="n">calc_accuracy_loader</span><span class="p">(</span><span class="n">data_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span> <span class="p">\\</span> <span class="ss">:default</span><span class="p">,</span> <span class="n">num_batches</span> <span class="p">\\</span> <span class="no">nil</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">loader_length</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">data_loader</span><span class="p">,</span> <span class="ss">:length</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
    <span class="n">num_batches</span> <span class="o">=</span> <span class="n">normalize_num_batches</span><span class="p">(</span><span class="n">num_batches</span><span class="p">,</span> <span class="n">loader_length</span><span class="p">)</span>

    <span class="c1"># Match the Python implementation's "evaluate at most num_batches" behavior</span>
    <span class="c1"># by taking only that many batches from the loader's epoch stream.</span>
    <span class="p">{</span><span class="n">correct_predictions</span><span class="p">,</span> <span class="n">num_examples</span><span class="p">}</span> <span class="o">=</span>
      <span class="n">data_loader</span><span class="o">.</span><span class="n">stream</span>
      <span class="o">|&gt;</span> <span class="no">Stream</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">num_batches</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">({</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">},</span> <span class="k">fn</span> <span class="n">batch</span><span class="p">,</span> <span class="p">{</span><span class="n">correct_predictions</span><span class="p">,</span> <span class="n">num_examples</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="p">{</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">}</span> <span class="o">=</span> <span class="n">stack_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

        <span class="c1"># Move both tensors together so model computation and label comparison</span>
        <span class="c1"># happen on the requested backend, just like `tensor.to(device)`.</span>
        <span class="n">input_batch</span> <span class="o">=</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>
        <span class="n">target_batch</span> <span class="o">=</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">target_batch</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

        <span class="c1"># A classification head emits logits for every sequence position. With a</span>
        <span class="c1"># causal model, the final token has seen the full message, so its logits</span>
        <span class="c1"># are the sequence-level class scores.</span>
        <span class="n">logits</span> <span class="o">=</span>
          <span class="n">model</span>
          <span class="o">|&gt;</span> <span class="n">forward_model</span><span class="p">(</span><span class="n">input_batch</span><span class="p">)</span>
          <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span> <span class="nv">&amp;1</span><span class="p">[[</span><span class="o">..</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">..</span><span class="p">]])</span>

        <span class="c1"># `argmax(axis: -1)` turns the final-token class scores into integer</span>
        <span class="c1"># labels. Comparing against the target tensor yields one boolean per row.</span>
        <span class="n">predicted_labels</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="c1"># Nx.tensor([0, 1, 1, 0]) and</span>
        <span class="c1"># Nx.tensor([0, 0, 1, 0]) =</span>
        <span class="c1"># Nx.tensor([1, 0, 1, 1], type: {:u, 8})</span>
        <span class="n">batch_correct</span> <span class="o">=</span> <span class="n">predicted_labels</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">target_batch</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span>
        <span class="n">batch_examples</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">axis_size</span><span class="p">(</span><span class="n">predicted_labels</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>

        <span class="p">{</span><span class="n">correct_predictions</span> <span class="o">+</span> <span class="n">batch_correct</span><span class="p">,</span> <span class="n">num_examples</span> <span class="o">+</span> <span class="n">batch_examples</span><span class="p">}</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">num_examples</span> <span class="o">==</span> <span class="mi">0</span> <span class="k">do</span>
      <span class="ss">:nan</span>
    <span class="k">else</span>
      <span class="n">correct_predictions</span> <span class="o">/</span> <span class="n">num_examples</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">tensor</span>
  <span class="k">defp</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="ss">:default</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">tensor</span>
  <span class="k">defp</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">normalize_num_batches</span><span class="p">(</span><span class="no">nil</span><span class="p">,</span> <span class="n">loader_length</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">loader_length</span>
  <span class="k">defp</span> <span class="n">normalize_num_batches</span><span class="p">(</span><span class="n">num_batches</span><span class="p">,</span> <span class="n">loader_length</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">min</span><span class="p">(</span><span class="n">num_batches</span><span class="p">,</span> <span class="n">loader_length</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">stack_batch</span><span class="p">({</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">})</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">stack_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">unzip</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

    <span class="p">{</span><span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">inputs</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">targets</span><span class="p">)}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">forward_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">module</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">__struct__</span>
    <span class="n">apply</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="ss">:forward</span><span class="p">,</span> <span class="p">[</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">])</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And here is a change in the existing <code class="language-plaintext highlighter-rouge">LossUtils</code> function to support evaluating only the final input token:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="gh">diff --git a/lib/llm_scratch/loss_utils.ex b/lib/llm_scratch/loss_utils.ex
index 765ffdb..2fcd385 100644
</span><span class="gd">--- a/lib/llm_scratch/loss_utils.ex
</span><span class="gi">+++ b/lib/llm_scratch/loss_utils.ex
</span><span class="p">@@ -63,22 +63,44 @@</span> defmodule LlmScratch.LossUtils do
   Defn-compatible mean cross entropy loss from logits and target token ids.
   """
   defn cross_entropy_loss_defn(logits, targets) do
<span class="gd">-    vocab_size = Nx.axis_size(logits, 2)
-
-    # step1: logits
-    logits
-    # step2: probabilities
-    |&gt; Axon.Activations.softmax(axis: -1)
-    # step3: target probabilities
-    |&gt; Nx.reshape({:auto, vocab_size})
-    |&gt; Nx.take_along_axis(Nx.reshape(targets, {:auto, 1}), axis: 1)
-    |&gt; Nx.squeeze(axes: [1])
-    # step4: logarithmic probabilities
-    |&gt; Nx.log()
-    # step6: negative average log probabilities
-    |&gt; Nx.negate()
-    # step5: average logarithmic probabilities
-    |&gt; Nx.mean()
</span><span class="gi">+    if Nx.rank(logits) == 3 do
+      # Language-model flow:
+      # logits shape is {batch_size, seq_len, vocab_size}
+      # targets shape is {batch_size, seq_len}
+      vocab_size = Nx.axis_size(logits, 2)
+
+      # step1: logits
+      logits
+      # step2: probabilities
+      |&gt; Axon.Activations.softmax(axis: -1)
+      # step3: target probabilities
+      |&gt; Nx.reshape({:auto, vocab_size})
+      |&gt; Nx.take_along_axis(Nx.reshape(targets, {:auto, 1}), axis: 1)
+      |&gt; Nx.squeeze(axes: [1])
+      # step4: logarithmic probabilities
+      |&gt; Nx.log()
+      # step6: negative average log probabilities
+      |&gt; Nx.negate()
+      # step5: average logarithmic probabilities
+      |&gt; Nx.mean()
+    else
+      # Classification flow:
+      # logits shape is {batch_size, num_classes}
+      # targets shape is {batch_size}
+      # step1: logits
+      logits
+      # step2: probabilities
+      |&gt; Axon.Activations.softmax(axis: -1)
+      # step3: target probabilities
+      |&gt; Nx.take_along_axis(Nx.reshape(targets, {:auto, 1}), axis: 1)
+      |&gt; Nx.squeeze(axes: [1])
+      # step4: logarithmic probabilities
+      |&gt; Nx.log()
+      # step6: negative average log probabilities
+      |&gt; Nx.negate()
+      # step5: average logarithmic probabilities
+      |&gt; Nx.mean()
+    end
</span>   end
 
   @doc """
<span class="p">@@ -98,21 +120,29 @@</span> defmodule LlmScratch.LossUtils do
   The optional `device` argument accepts an Nx backend, such as `EXLA.Backend`
   or `{EXLA.Backend, client: :cuda}`. Use `nil` or `:default` to leave tensors
   on their current backend.
<span class="gi">+
+  Options:
+
+    * `:target` - use `:all_tokens` or omit it for language-model loss over all
+      token positions. Use `:last_token` for sequence classification, where the
+      model output is sliced to `model(input_batch)[:, -1, :]` before computing
+      cross entropy.
</span>   """
<span class="gd">-  @spec calc_loss_batch(Nx.Tensor.t(), Nx.Tensor.t(), struct(), nil | :default | atom() | tuple()) ::
</span><span class="gi">+  @spec calc_loss_batch(
+          Nx.Tensor.t(),
+          Nx.Tensor.t(),
+          struct(),
+          nil | :default | atom() | tuple(),
+          keyword()
+        ) ::
</span>           Nx.Tensor.t()
<span class="gd">-  def calc_loss_batch(
-        %Nx.Tensor{} = input_batch,
-        %Nx.Tensor{} = target_batch,
-        model,
-        device \\ :default
-      )
-      when is_struct(model) do
</span><span class="gi">+  def calc_loss_batch(input_batch, target_batch, model, device \\ :default, opts \\ []) do
</span>     input_batch = maybe_transfer(input_batch, device)
     target_batch = maybe_transfer(target_batch, device)
 
     model
     |&gt; forward_model(input_batch)
<span class="gi">+    |&gt; select_loss_logits(opts)
</span>     |&gt; cross_entropy_loss(target_batch)
   end
 
<span class="p">@@ -127,16 +157,19 @@</span> defmodule LlmScratch.LossUtils do
 
   Pass `nil` for `num_batches` to evaluate one full pass through the loader.
   Returns `:nan` when the loader has no batches.
<span class="gi">+
+  Accepts the same options as `calc_loss_batch/5`; use `target: :last_token` for
+  classification loss.
</span>   """
   @spec calc_loss_loader(
           map(),
           struct(),
           nil | :default | atom() | tuple(),
<span class="gd">-          nil | non_neg_integer()
</span><span class="gi">+          nil | non_neg_integer(),
+          keyword()
</span>         ) ::
           float() | :nan
<span class="gd">-  def calc_loss_loader(data_loader, model, device \\ :default, num_batches \\ nil)
-      when is_map(data_loader) and is_struct(model) do
</span><span class="gi">+  def calc_loss_loader(data_loader, model, device \\ :default, num_batches \\ nil, opts \\ []) do
</span>     loader_length = Map.get(data_loader, :length, 0)
 
     cond do
<span class="p">@@ -150,7 +183,7 @@</span> defmodule LlmScratch.LossUtils do
         |&gt; Stream.take(num_batches)
         |&gt; Enum.reduce(0.0, fn batch, total_loss -&gt;
           {input_batch, target_batch} = stack_batch(batch)
<span class="gd">-          loss = calc_loss_batch(input_batch, target_batch, model, device)
</span><span class="gi">+          loss = calc_loss_batch(input_batch, target_batch, model, device, opts)
</span> 
           total_loss + Nx.to_number(loss)
         end)
<span class="p">@@ -158,17 +191,28 @@</span> defmodule LlmScratch.LossUtils do
     end
   end
 
<span class="gd">-  defp maybe_transfer(tensor, device) when device in [nil, :default], do: tensor
</span><span class="gi">+  defp maybe_transfer(tensor, nil), do: tensor
+  defp maybe_transfer(tensor, :default), do: tensor
</span>   defp maybe_transfer(tensor, device), do: Nx.backend_transfer(tensor, device)
 
   defp normalize_num_batches(nil, loader_length), do: loader_length
   defp normalize_num_batches(num_batches, loader_length), do: min(num_batches, loader_length)
 
<span class="gd">-  defp stack_batch({%Nx.Tensor{} = input_batch, %Nx.Tensor{} = target_batch}) do
</span><span class="gi">+  defp select_loss_logits(logits, opts) do
+    if Keyword.get(opts, :target) == :last_token do
+      # Classification uses only the final token because it has causal attention
+      # over the full input message.
+      logits[[.., -1, ..]]
+    else
+      logits
+    end
+  end
+
+  defp stack_batch({input_batch, target_batch}) do
</span>     {input_batch, target_batch}
   end
 
<span class="gd">-  defp stack_batch(batch) when is_list(batch) do
</span><span class="gi">+  defp stack_batch(batch) do
</span>     {inputs, targets} = Enum.unzip(batch)
 
     {Nx.stack(inputs), Nx.stack(targets)}
<span class="p">@@ -176,17 +220,6 @@</span> defmodule LlmScratch.LossUtils do
 
   defp forward_model(model, input_batch) do
     module = model.__struct__
<span class="gd">-
-    cond do
-      function_exported?(module, :forward, 2) -&gt;
-        apply(module, :forward, [model, input_batch])
-
-      function_exported?(module, :forward, 3) -&gt;
-        apply(module, :forward, [model, input_batch, []])
-
-      true -&gt;
-        raise ArgumentError,
-              "expected #{inspect(module)} to export forward/2 or forward/3"
-    end
</span><span class="gi">+    apply(module, :forward, [model, input_batch])
</span>  end
<span class="p">end
</span></code></pre></div></div>

<p>And here is a test that uses those new utility functions:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">900_000</span>
  <span class="n">test</span> <span class="s2">"6.6 calculates classification accuracy before fine-tuning"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="n">device</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>

    <span class="n">train_dataset</span> <span class="o">=</span> <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"train.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="no">nil</span><span class="p">)</span>

    <span class="n">val_dataset</span> <span class="o">=</span>
      <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"validation.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span><span class="p">)</span>

    <span class="n">test_dataset</span> <span class="o">=</span>
      <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"test.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span><span class="p">)</span>

    <span class="n">batch_size</span> <span class="o">=</span> <span class="mi">8</span>
    <span class="ss">:rand</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="ss">:exsss</span><span class="p">,</span> <span class="p">{</span><span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">})</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="n">train_dataset</span>
      <span class="o">|&gt;</span> <span class="n">dataset_samples</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span> <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">,</span> <span class="ss">num_workers:</span> <span class="mi">0</span><span class="p">,</span> <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="n">val_dataset</span>
      <span class="o">|&gt;</span> <span class="n">dataset_samples</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span> <span class="ss">num_workers:</span> <span class="mi">0</span><span class="p">,</span> <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">)</span>

    <span class="n">test_loader</span> <span class="o">=</span>
      <span class="n">test_dataset</span>
      <span class="o">|&gt;</span> <span class="n">dataset_samples</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span> <span class="ss">num_workers:</span> <span class="mi">0</span><span class="p">,</span> <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span>
      <span class="s2">"124M"</span>
      <span class="o">|&gt;</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">freeze</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">replace_out_head</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">,</span> <span class="ss">bias:</span> <span class="no">true</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">set_trainable</span><span class="p">([</span><span class="ss">:out_head</span><span class="p">,</span> <span class="ss">:final_norm</span><span class="p">,</span> <span class="p">{</span><span class="ss">:trf_block</span><span class="p">,</span> <span class="mi">11</span><span class="p">}])</span>

    <span class="c1"># returns the number of correctly predicted classification divided with number of messages in a batch</span>
    <span class="n">train_accuracy</span> <span class="o">=</span> <span class="no">LossClassificationUtils</span><span class="o">.</span><span class="n">calc_accuracy_loader</span><span class="p">(</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
    <span class="n">val_accuracy</span> <span class="o">=</span> <span class="no">LossClassificationUtils</span><span class="o">.</span><span class="n">calc_accuracy_loader</span><span class="p">(</span><span class="n">val_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
    <span class="n">test_accuracy</span> <span class="o">=</span> <span class="no">LossClassificationUtils</span><span class="o">.</span><span class="n">calc_accuracy_loader</span><span class="p">(</span><span class="n">test_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>

    <span class="n">train_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">10</span><span class="p">,</span> <span class="ss">target:</span> <span class="ss">:last_token</span><span class="p">)</span>
    <span class="n">val_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">val_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">10</span><span class="p">,</span> <span class="ss">target:</span> <span class="ss">:last_token</span><span class="p">)</span>
    <span class="n">test_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">test_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="mi">10</span><span class="p">,</span> <span class="ss">target:</span> <span class="ss">:last_token</span><span class="p">)</span>

    <span class="c1"># we have almost 50:50 chance for ham/spam, as we still have not trained our model for classification</span>
    <span class="n">assert</span> <span class="n">train_accuracy</span> <span class="o">==</span> <span class="mf">0.5125</span>
    <span class="n">assert</span> <span class="n">val_accuracy</span> <span class="o">==</span> <span class="mf">0.525</span>
    <span class="n">assert</span> <span class="n">test_accuracy</span> <span class="o">==</span> <span class="mf">0.5625</span>

    <span class="n">assert_in_delta</span> <span class="n">train_loss</span><span class="p">,</span> <span class="mf">1.5194151997566223</span><span class="p">,</span> <span class="mf">1.0e-6</span>
    <span class="n">assert_in_delta</span> <span class="n">val_loss</span><span class="p">,</span> <span class="mf">1.4858015120029449</span><span class="p">,</span> <span class="mf">1.0e-6</span>
    <span class="n">assert_in_delta</span> <span class="n">test_loss</span><span class="p">,</span> <span class="mf">1.3634233981370927</span><span class="p">,</span> <span class="mf">1.0e-6</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="fine-tuning-the-model-on-supervised-data">Fine-tuning the model on supervised data</h2>

<p>We use a training loop that is very similar to the language-model training loop. The main difference is that we use classification accuracy for evaluation and track the number of examples seen instead of the number of tokens seen.</p>

<p>Here is the <code class="language-plaintext highlighter-rouge">Training</code> module diff:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="gh">diff --git a/lib/llm_scratch/training.ex b/lib/llm_scratch/training.ex
index fadfefd..14f1849 100644
</span><span class="gd">--- a/lib/llm_scratch/training.ex
</span><span class="gi">+++ b/lib/llm_scratch/training.ex
</span><span class="p">@@ -184,7 +184,7 @@</span> defmodule LlmScratch.Training do
 
   import Nx.Defn
 
<span class="gd">-  alias LlmScratch.{GPTModel, LossUtils, TextGeneration, TextUtils}
</span><span class="gi">+  alias LlmScratch.{GPTModel, LossClassificationUtils, LossUtils, TextGeneration, TextUtils}
</span> 
   defmodule AdamW do
     @moduledoc """
<span class="p">@@ -413,6 +413,136 @@</span> defmodule LlmScratch.Training do
     end
   end
 
<span class="gi">+  @spec train_classifier_simple(
+          struct(),
+          map(),
+          map(),
+          optimizer(),
+          nil | :default | atom() | tuple(),
+          non_neg_integer(),
+          pos_integer(),
+          pos_integer(),
+          keyword()
+        ) ::
+          {struct(), [float()], [float()], [float()], [float()], non_neg_integer()}
+          | {struct(), optimizer(), [float()], [float()], [float()], [float()], non_neg_integer()}
+  @doc """
+  Trains a GPT-style sequence classifier with a simple batch loop.
+
+  This mirrors the book's `train_classifier_simple` Python example. The Python
+  model is mutated in place, but Elixir data is immutable, so the updated model
+  is returned.
+
+  ## Parameters
+
+    * `model` - classifier model whose trainable tensors can be traversed by
+      `Nx.Container`. For GPT classification, this is typically a frozen GPT
+      model with a 2-class output head and selected layers marked trainable.
+    * `train_loader` - training data loader map with `:stream`, `:batches`, and
+      `:length`.
+    * `val_loader` - validation data loader map with `:stream`, `:batches`, and
+      `:length`.
+    * `optimizer` - either `Training.adamw/2` output or a two-argument function
+      `(model, gradients -&gt; updated_model)`.
+    * `device` - Nx backend target. Use `:default` or `nil` to keep tensors on
+      their current backend.
+    * `num_epochs` - number of full passes over `train_loader`.
+    * `eval_freq` - evaluate loss every `eval_freq` global training steps.
+    * `eval_iter` - maximum number of batches from each loader used for loss and
+      accuracy evaluation.
+    * `opts` - optional settings:
+      * `:return_optimizer` - when `true`, include final optimizer state in the
+        returned tuple.
+
+  ## Returns
+
+    * `{model, train_losses, val_losses, train_accs, val_accs, examples_seen}` by
+      default.
+    * `{model, optimizer, train_losses, val_losses, train_accs, val_accs,
+      examples_seen}` when `return_optimizer: true`.
+
+  `train_losses` and `val_losses` are captured at `eval_freq` intervals.
+  `train_accs` and `val_accs` are captured once after each epoch. The final
+  `examples_seen` count tracks training examples rather than token counts.
+  """
+  def train_classifier_simple(
+        model,
+        train_loader,
+        val_loader,
+        optimizer,
+        device,
+        num_epochs,
+        eval_freq,
+        eval_iter,
+        opts \\ []
+      ) do
+    # Classifier training still uses dropout in train mode, so we thread an
+    # explicit RNG key through every batch.
+    key =
+      System.unique_integer([:positive])
+      |&gt; Nx.Random.key()
+      |&gt; maybe_transfer_tensor(device)
+
+    # State carries immutable model/optimizer updates plus the metrics collected
+    # during training.
+    state = %{
+      model: maybe_transfer_model(model, device),
+      optimizer: optimizer,
+      key: key,
+      train_losses: [],
+      val_losses: [],
+      train_accs: [],
+      val_accs: [],
+      examples_seen: 0,
+      global_step: -1
+    }
+
+    state =
+      Enum.reduce(1..num_epochs//1, state, fn epoch, state -&gt;
+        state =
+          train_classifier_epoch(
+            state,
+            train_loader,
+            val_loader,
+            device,
+            epoch,
+            eval_freq,
+            eval_iter
+          )
+
+        # Accuracy is non-differentiable, so it is only measured after the epoch
+        # with the current model parameters.
+        train_accuracy =
+          LossClassificationUtils.calc_accuracy_loader(train_loader, state.model, device, eval_iter)
+
+        val_accuracy =
+          LossClassificationUtils.calc_accuracy_loader(val_loader, state.model, device, eval_iter)
+
+        IO.puts(
+          "Training accuracy: #{format_percent(train_accuracy)}% | " &lt;&gt;
+            "Validation accuracy: #{format_percent(val_accuracy)}%"
+        )
+
+        %{
+          state
+          | train_accs: [train_accuracy | state.train_accs],
+            val_accs: [val_accuracy | state.val_accs]
+        }
+      end)
+
+    train_losses = Enum.reverse(state.train_losses)
+    val_losses = Enum.reverse(state.val_losses)
+    train_accs = Enum.reverse(state.train_accs)
+    val_accs = Enum.reverse(state.val_accs)
+
+    if Keyword.get(opts, :return_optimizer, false) do
+      {state.model, state.optimizer, train_losses, val_losses, train_accs, val_accs,
+       state.examples_seen}
+    else
+      {state.model, train_losses, val_losses, train_accs, val_accs, state.examples_seen}
+    end
+  end
+
</span>   @doc """
   Calculates training and validation losses for a fixed number of batches.
 
<span class="p">@@ -440,6 +570,29 @@</span> defmodule LlmScratch.Training do
     }
   end
 
<span class="gi">+  @doc """
+  Calculates classifier training and validation losses for fixed batch counts.
+
+  ## Parameters
+
+    * `model` - classifier model used for inference.
+    * `train_loader` - training data loader map with `:stream` and `:length`.
+    * `val_loader` - validation data loader map with `:stream` and `:length`.
+    * `device` - Nx backend target passed through to `LossUtils.calc_loss_loader/5`.
+    * `eval_iter` - maximum number of batches evaluated from each loader.
+
+  ## Returns
+
+    * `{train_loss, val_loss}` as floats. The loss is computed from final-token
+      logits via `target: :last_token`.
+  """
+  def evaluate_classifier_model(model, train_loader, val_loader, device, eval_iter) do
+    {
+      LossUtils.calc_loss_loader(train_loader, model, device, eval_iter, target: :last_token),
+      LossUtils.calc_loss_loader(val_loader, model, device, eval_iter, target: :last_token)
+    }
+  end
+
</span>   @doc """
   Generates and prints a text sample from `start_context`.
 
<span class="p">@@ -510,6 +663,99 @@</span> defmodule LlmScratch.Training do
     {loss, gradients, key}
   end
 
<span class="gi">+  @doc """
+  Computes classifier loss and gradients for one input/target batch.
+
+  ## Parameters
+
+    * `model` - classifier model whose trainable tensors implement
+      `Nx.Container`.
+    * `input_batch` - token id tensor shaped `{batch_size, seq_len}`.
+    * `target_batch` - class label tensor shaped `{batch_size}`.
+    * `key` - `Nx.Random` key used for dropout.
+
+  ## Returns
+
+    * `{loss, gradients, key}` where `loss` is cross entropy over final-token
+      logits, `gradients` matches the model structure, and `key` is the advanced
+      dropout RNG key.
+  """
+  defn classifier_loss_and_grad(model, input_batch, target_batch, key) do
+    {{loss, key}, {gradients, _input_gradients, _target_gradients, _key_gradients}} =
+      value_and_grad(
+        {model, input_batch, target_batch, key},
+        fn {model, input_batch, target_batch, key} -&gt;
+          {logits, key} = GPTModel.train(model, input_batch, key)
+          logits = logits[[.., -1, ..]]
+          {LossUtils.cross_entropy_loss_defn(logits, target_batch), key}
+        end,
+        &amp;elem(&amp;1, 0)
+      )
+
+    {loss, gradients, key}
+  end
+
+  defp train_classifier_epoch(
+         state,
+         train_loader,
+         val_loader,
+         device,
+         epoch,
+         eval_freq,
+         eval_iter
+       ) do
+    train_loader
+    # Shuffle or enumerate one finite epoch of batches, mirroring the language
+    # model training loop.
+    |&gt; epoch_batches()
+    |&gt; Enum.reduce(state, fn batch, state -&gt;
+      # Split a loader batch into model inputs and class labels, then transfer
+      # both tensors to the requested backend.
+      {input_batch, target_batch} = stack_batch(batch, device)
+
+      # Compute differentiable cross entropy and its gradients for the current
+      # model parameters. This is the Elixir/Nx counterpart to loss.backward().
+      {_loss, gradients, key} =
+        classifier_loss_and_grad(state.model, input_batch, target_batch, state.key)
+
+      # Apply optimizer updates and carry AdamW moment history forward.
+      {model, optimizer} = optimizer_step(state.optimizer, state.model, gradients)
+
+      # Classification progress is tracked by examples, not tokens.
+      examples_seen = state.examples_seen + Nx.axis_size(input_batch, 0)
+      global_step = state.global_step + 1
+
+      state = %{
+        state
+        | model: model,
+          optimizer: optimizer,
+          key: key,
+          examples_seen: examples_seen,
+          global_step: global_step
+      }
+
+      # Periodically evaluate differentiable loss on a fixed number of train and
+      # validation batches. Accuracy is measured after each epoch.
+      if rem(global_step, eval_freq) == 0 do
+        {train_loss, val_loss} =
+          evaluate_classifier_model(model, train_loader, val_loader, device, eval_iter)
+
+        IO.puts(
+          "Ep #{epoch} (Step #{pad_step(global_step)}): " &lt;&gt;
+            "Train loss #{format_loss(train_loss)}, Val loss #{format_loss(val_loss)}"
+        )
+
+        %{
+          state
+          | train_losses: [train_loss | state.train_losses],
+            val_losses: [val_loss | state.val_losses]
+        }
+      else
+        state
+      end
+    end)
+  end
+
</span>   defp train_epoch(
          state,
          train_loader,
<span class="p">@@ -840,6 +1086,9 @@</span> defmodule LlmScratch.Training do
   defp format_loss(:nan), do: "nan"
   defp format_loss(loss), do: :erlang.float_to_binary(loss * 1.0, decimals: 3)
 
<span class="gi">+  defp format_percent(:nan), do: "nan"
+  defp format_percent(value), do: :erlang.float_to_binary(value * 100.0, decimals: 2)
+
</span>   defp pad_step(step) do
     step
     |&gt; Integer.to_string()
</code></pre></div></div>

<p>And here is the test where we do the training:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">:train</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">3_600_000</span>
  <span class="n">test</span> <span class="s2">"6.7 fine-tunes classifier on spam dataset"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="n">device</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">tokenizer</span> <span class="o">=</span> <span class="o">&amp;</span><span class="n">gpt2_compatible_token_ids</span><span class="o">/</span><span class="mi">1</span>

    <span class="n">train_dataset</span> <span class="o">=</span> <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"train.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="no">nil</span><span class="p">)</span>

    <span class="n">val_dataset</span> <span class="o">=</span>
      <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"validation.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span><span class="p">)</span>

    <span class="n">batch_size</span> <span class="o">=</span> <span class="mi">8</span>
    <span class="ss">:rand</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="ss">:exsss</span><span class="p">,</span> <span class="p">{</span><span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">,</span> <span class="mi">123</span><span class="p">})</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="n">train_dataset</span>
      <span class="o">|&gt;</span> <span class="n">dataset_samples</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span> <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">,</span> <span class="ss">num_workers:</span> <span class="mi">0</span><span class="p">,</span> <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="n">val_dataset</span>
      <span class="o">|&gt;</span> <span class="n">dataset_samples</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">DataLoader</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">batch_size:</span> <span class="n">batch_size</span><span class="p">,</span> <span class="ss">num_workers:</span> <span class="mi">0</span><span class="p">,</span> <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span>
      <span class="s2">"124M"</span>
      <span class="o">|&gt;</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">freeze</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">replace_out_head</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">,</span> <span class="ss">bias:</span> <span class="no">true</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">set_trainable</span><span class="p">([</span><span class="ss">:out_head</span><span class="p">,</span> <span class="ss">:final_norm</span><span class="p">,</span> <span class="p">{</span><span class="ss">:trf_block</span><span class="p">,</span> <span class="mi">11</span><span class="p">}])</span>

    <span class="n">optimizer</span> <span class="o">=</span> <span class="no">Training</span><span class="o">.</span><span class="n">adamw</span><span class="p">(</span><span class="mf">5.0e-5</span><span class="p">,</span> <span class="ss">weight_decay:</span> <span class="mf">0.1</span><span class="p">)</span>
    <span class="n">num_epochs</span> <span class="o">=</span> <span class="mi">5</span>

    <span class="p">{</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">train_losses</span><span class="p">,</span> <span class="n">val_losses</span><span class="p">,</span> <span class="n">train_accs</span><span class="p">,</span> <span class="n">val_accs</span><span class="p">,</span>
     <span class="n">examples_seen</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Training</span><span class="o">.</span><span class="n">train_classifier_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">train_loader</span><span class="p">,</span>
        <span class="n">val_loader</span><span class="p">,</span>
        <span class="n">optimizer</span><span class="p">,</span>
        <span class="n">device</span><span class="p">,</span>
        <span class="n">num_epochs</span><span class="p">,</span>
        <span class="mi">50</span><span class="p">,</span>
        <span class="mi">5</span><span class="p">,</span>
        <span class="ss">return_optimizer:</span> <span class="no">true</span>
      <span class="p">)</span>

    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="s2">"ch6_spam_classifier_model_and_optimizer.nx"</span>
    <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">save_training_state!</span><span class="p">(</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="n">metrics_path</span> <span class="o">=</span> <span class="s2">"ch6_spam_classifier_training_metrics.json"</span>

    <span class="n">write_training_metrics!</span><span class="p">(</span>
      <span class="n">metrics_path</span><span class="p">,</span>
      <span class="n">num_epochs</span><span class="p">,</span>
      <span class="n">examples_seen</span><span class="p">,</span>
      <span class="n">train_losses</span><span class="p">,</span>
      <span class="n">val_losses</span><span class="p">,</span>
      <span class="n">train_accs</span><span class="p">,</span>
      <span class="n">val_accs</span>
    <span class="p">)</span>

    <span class="n">assert</span> <span class="p">%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">trained_model</span>
    <span class="n">assert</span> <span class="p">%</span><span class="no">Training</span><span class="o">.</span><span class="no">AdamW</span><span class="p">{}</span> <span class="o">=</span> <span class="n">trained_optimizer</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">)</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_losses</span><span class="p">)</span> <span class="o">==</span> <span class="mi">13</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">val_losses</span><span class="p">)</span> <span class="o">==</span> <span class="mi">13</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_accs</span><span class="p">)</span> <span class="o">==</span> <span class="n">num_epochs</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">val_accs</span><span class="p">)</span> <span class="o">==</span> <span class="n">num_epochs</span>
    <span class="n">assert</span> <span class="n">examples_seen</span> <span class="o">==</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">*</span> <span class="n">batch_size</span> <span class="o">*</span> <span class="n">num_epochs</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">train_losses</span> <span class="o">++</span> <span class="n">val_losses</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_float</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">train_accs</span> <span class="o">++</span> <span class="n">val_accs</span><span class="p">,</span> <span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">&gt;=</span> <span class="mf">0.0</span> <span class="ow">and</span> <span class="nv">&amp;1</span> <span class="o">&lt;=</span> <span class="mf">1.0</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">gpt2_compatible_token_ids</span><span class="p">(</span><span class="n">text</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">token_ids</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">,</span> <span class="n">text</span><span class="p">,</span> <span class="p">[</span><span class="s2">"&lt;|endoftext|&gt;"</span><span class="p">])</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">min</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="mi">50_256</span><span class="p">))</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>That produced the following training output:</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  Ep 1 (Step 000000): Train loss 1.224, Val loss 1.587
Ep 1 (Step 000050): Train loss 0.625, Val loss 0.694
Ep 1 (Step 000100): Train loss 0.542, Val loss 0.569
Training accuracy: 92.50% | Validation accuracy: 92.50%
Ep 2 (Step 000150): Train loss 0.240, Val loss 0.214
Ep 2 (Step 000200): Train loss 0.039, Val loss 0.047
Ep 2 (Step 000250): Train loss 0.022, Val loss 0.066
Training accuracy: 100.00% | Validation accuracy: 100.00%
Ep 3 (Step 000300): Train loss 0.021, Val loss 0.080
Ep 3 (Step 000350): Train loss 0.012, Val loss 0.038
Training accuracy: 100.00% | Validation accuracy: 100.00%
Ep 4 (Step 000400): Train loss 0.012, Val loss 0.036
Ep 4 (Step 000450): Train loss 0.012, Val loss 0.049
Ep 4 (Step 000500): Train loss 0.011, Val loss 0.039
Training accuracy: 100.00% | Validation accuracy: 97.50%
Ep 5 (Step 000550): Train loss 0.007, Val loss 0.015
Ep 5 (Step 000600): Train loss 0.006, Val loss 0.035
Training accuracy: 100.00% | Validation accuracy: 97.50%
</code></pre></div></div>

<p>You can see that we produced very good results using OpenAI’s 124M-parameter GPT-2 model.</p>

<p>Let’s plot the losses in Livebook. First, we use this helper to write the training metrics to a file, so we do not need to train the model again in Livebook:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defp</span> <span class="n">write_training_metrics!</span><span class="p">(</span>
         <span class="n">path</span><span class="p">,</span>
         <span class="n">num_epochs</span><span class="p">,</span>
         <span class="n">examples_seen</span><span class="p">,</span>
         <span class="n">train_losses</span><span class="p">,</span>
         <span class="n">val_losses</span><span class="p">,</span>
         <span class="n">train_accs</span><span class="p">,</span>
         <span class="n">val_accs</span>
       <span class="p">)</span> <span class="k">do</span>
    <span class="n">metrics</span> <span class="o">=</span> <span class="p">%{</span>
      <span class="ss">num_epochs:</span> <span class="n">num_epochs</span><span class="p">,</span>
      <span class="ss">examples_seen:</span> <span class="n">examples_seen</span><span class="p">,</span>
      <span class="ss">losses:</span> <span class="p">%{</span>
        <span class="ss">epochs_seen:</span> <span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">num_epochs</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">length</span><span class="p">(</span><span class="n">train_losses</span><span class="p">)),</span>
        <span class="ss">examples_seen:</span> <span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">examples_seen</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">length</span><span class="p">(</span><span class="n">train_losses</span><span class="p">)),</span>
        <span class="ss">train_values:</span> <span class="n">train_losses</span><span class="p">,</span>
        <span class="ss">val_values:</span> <span class="n">val_losses</span>
      <span class="p">},</span>
      <span class="ss">accuracies:</span> <span class="p">%{</span>
        <span class="ss">epochs_seen:</span> <span class="n">linspace</span><span class="p">(</span><span class="mf">1.0</span><span class="p">,</span> <span class="n">num_epochs</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">length</span><span class="p">(</span><span class="n">train_accs</span><span class="p">)),</span>
        <span class="ss">examples_seen:</span>
          <span class="n">linspace</span><span class="p">(</span><span class="n">examples_seen</span> <span class="o">/</span> <span class="n">num_epochs</span><span class="p">,</span> <span class="n">examples_seen</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">length</span><span class="p">(</span><span class="n">train_accs</span><span class="p">)),</span>
        <span class="ss">train_values:</span> <span class="n">train_accs</span><span class="p">,</span>
        <span class="ss">val_values:</span> <span class="n">val_accs</span>
      <span class="p">}</span>
    <span class="p">}</span>

    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">encoded_metrics</span><span class="p">}</span> <span class="o">=</span> <span class="no">Jason</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">metrics</span><span class="p">,</span> <span class="ss">pretty:</span> <span class="no">true</span><span class="p">)</span>
    <span class="no">File</span><span class="o">.</span><span class="n">write!</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="n">encoded_metrics</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>And here is the Livebook code:</p>

<h2 id="llm-from-scratch-6-classifier-metrics">LLM from Scratch 6 Classifier Metrics</h2>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">project_path</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="s2">".."</span><span class="p">,</span> <span class="n">__DIR__</span><span class="p">)</span>

<span class="no">Mix</span><span class="o">.</span><span class="n">install</span><span class="p">(</span>
  <span class="p">[</span>
    <span class="p">{</span><span class="ss">:jason</span><span class="p">,</span> <span class="s2">"~&gt; 1.4"</span><span class="p">},</span>
    <span class="p">{</span><span class="ss">:kino</span><span class="p">,</span> <span class="s2">"~&gt; 0.12"</span><span class="p">},</span>
    <span class="p">{</span><span class="ss">:kino_vega_lite</span><span class="p">,</span> <span class="s2">"~&gt; 0.1.11"</span><span class="p">},</span>
    <span class="p">{</span><span class="ss">:vega_lite</span><span class="p">,</span> <span class="s2">"~&gt; 0.1.9"</span><span class="p">}</span>
  <span class="p">],</span>
  <span class="ss">consolidate_protocols:</span> <span class="no">false</span>
<span class="p">)</span>
</code></pre></div></div>

<h2 id="load-metrics">Load metrics</h2>

<p>Run test <code class="language-plaintext highlighter-rouge">6.7</code> first to generate <code class="language-plaintext highlighter-rouge">ch6_spam_classifier_training_metrics.json</code>.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">metrics_path</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">project_path</span><span class="p">,</span> <span class="s2">"ch6_spam_classifier_training_metrics.json"</span><span class="p">)</span>

<span class="k">unless</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">metrics_path</span><span class="p">)</span> <span class="k">do</span>
  <span class="k">raise</span> <span class="sd">"""
  Expected classifier training metrics at #{metrics_path}.

  Generate it with:

      mix test test/llm-from-scratch-6_test.exs --only train
  """</span>
<span class="k">end</span>

<span class="n">metrics</span> <span class="o">=</span>
  <span class="n">metrics_path</span>
  <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
  <span class="o">|&gt;</span> <span class="no">Jason</span><span class="o">.</span><span class="n">decode!</span><span class="p">()</span>

<span class="p">%{</span>
  <span class="ss">metrics_path:</span> <span class="n">metrics_path</span><span class="p">,</span>
  <span class="ss">num_epochs:</span> <span class="n">metrics</span><span class="p">[</span><span class="s2">"num_epochs"</span><span class="p">],</span>
  <span class="ss">examples_seen:</span> <span class="n">metrics</span><span class="p">[</span><span class="s2">"examples_seen"</span><span class="p">],</span>
  <span class="ss">loss_points:</span> <span class="n">length</span><span class="p">(</span><span class="n">metrics</span><span class="p">[</span><span class="s2">"losses"</span><span class="p">][</span><span class="s2">"train_values"</span><span class="p">]),</span>
  <span class="ss">accuracy_points:</span> <span class="n">length</span><span class="p">(</span><span class="n">metrics</span><span class="p">[</span><span class="s2">"accuracies"</span><span class="p">][</span><span class="s2">"train_values"</span><span class="p">])</span>
<span class="p">}</span>
</code></pre></div></div>

<h2 id="plot-values">Plot values</h2>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">alias</span> <span class="no">VegaLite</span><span class="p">,</span> <span class="ss">as:</span> <span class="no">Vl</span>

<span class="k">defmodule</span> <span class="no">ClassifierPlot</span> <span class="k">do</span>
  <span class="k">def</span> <span class="n">plot_values</span><span class="p">(</span><span class="n">epochs_seen</span><span class="p">,</span> <span class="n">examples_seen</span><span class="p">,</span> <span class="n">train_values</span><span class="p">,</span> <span class="n">val_values</span><span class="p">,</span> <span class="n">label</span> <span class="p">\\</span> <span class="s2">"loss"</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">value_key</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">to_atom</span><span class="p">(</span><span class="n">label</span><span class="p">)</span>

    <span class="n">value_data</span> <span class="o">=</span>
      <span class="p">[</span>
        <span class="p">{</span><span class="s2">"Training </span><span class="si">#{</span><span class="n">label</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span> <span class="n">train_values</span><span class="p">},</span>
        <span class="p">{</span><span class="s2">"Validation </span><span class="si">#{</span><span class="n">label</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span> <span class="n">val_values</span><span class="p">}</span>
      <span class="p">]</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">flat_map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">series</span><span class="p">,</span> <span class="n">values</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">([</span><span class="n">epochs_seen</span><span class="p">,</span> <span class="n">examples_seen</span><span class="p">,</span> <span class="n">values</span><span class="p">])</span>
        <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">epoch</span><span class="p">,</span> <span class="n">examples</span><span class="p">,</span> <span class="n">value</span><span class="p">}</span> <span class="o">-&gt;</span>
          <span class="p">%{</span>
            <span class="n">value_key</span> <span class="o">=&gt;</span> <span class="n">value</span><span class="p">,</span>
            <span class="ss">series:</span> <span class="n">series</span><span class="p">,</span>
            <span class="ss">epoch:</span> <span class="n">epoch</span><span class="p">,</span>
            <span class="ss">examples_seen:</span> <span class="n">examples</span>
          <span class="p">}</span>
        <span class="k">end</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="n">examples_axis_data</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">([</span><span class="n">epochs_seen</span><span class="p">,</span> <span class="n">examples_seen</span><span class="p">,</span> <span class="n">train_values</span><span class="p">])</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">epoch</span><span class="p">,</span> <span class="n">examples</span><span class="p">,</span> <span class="n">value</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="p">%{</span>
          <span class="n">value_key</span> <span class="o">=&gt;</span> <span class="n">value</span><span class="p">,</span>
          <span class="ss">epoch:</span> <span class="n">epoch</span><span class="p">,</span>
          <span class="ss">examples_seen:</span> <span class="n">examples</span>
        <span class="p">}</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="c1"># The first layer draws the visible train/validation curves against epochs.</span>
    <span class="n">epoch_layer</span> <span class="o">=</span>
      <span class="no">Vl</span><span class="o">.</span><span class="n">new</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">data_from_values</span><span class="p">(</span><span class="n">value_data</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">mark</span><span class="p">(</span><span class="ss">:line</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:x</span><span class="p">,</span> <span class="s2">"epoch"</span><span class="p">,</span>
        <span class="ss">type:</span> <span class="ss">:quantitative</span><span class="p">,</span>
        <span class="ss">title:</span> <span class="s2">"Epochs"</span><span class="p">,</span>
        <span class="ss">axis:</span> <span class="p">[</span><span class="ss">tick_min_step:</span> <span class="mi">1</span><span class="p">]</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:y</span><span class="p">,</span> <span class="n">label</span><span class="p">,</span> <span class="ss">type:</span> <span class="ss">:quantitative</span><span class="p">,</span> <span class="ss">title:</span> <span class="no">String</span><span class="o">.</span><span class="n">capitalize</span><span class="p">(</span><span class="n">label</span><span class="p">))</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:color</span><span class="p">,</span> <span class="s2">"series"</span><span class="p">,</span> <span class="ss">type:</span> <span class="ss">:nominal</span><span class="p">,</span> <span class="ss">title:</span> <span class="no">nil</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:stroke_dash</span><span class="p">,</span> <span class="s2">"series"</span><span class="p">,</span> <span class="ss">type:</span> <span class="ss">:nominal</span><span class="p">,</span> <span class="ss">title:</span> <span class="no">nil</span><span class="p">)</span>

    <span class="c1"># The second layer is invisible; it aligns a top x-axis with the same y-scale</span>
    <span class="c1"># so we can read progress as examples seen.</span>
    <span class="n">examples_layer</span> <span class="o">=</span>
      <span class="no">Vl</span><span class="o">.</span><span class="n">new</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">data_from_values</span><span class="p">(</span><span class="n">examples_axis_data</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">mark</span><span class="p">(</span><span class="ss">:line</span><span class="p">,</span> <span class="ss">opacity:</span> <span class="mi">0</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:x</span><span class="p">,</span> <span class="s2">"examples_seen"</span><span class="p">,</span>
        <span class="ss">type:</span> <span class="ss">:quantitative</span><span class="p">,</span>
        <span class="ss">title:</span> <span class="s2">"Examples seen"</span><span class="p">,</span>
        <span class="ss">axis:</span> <span class="p">[</span><span class="ss">orient:</span> <span class="s2">"top"</span><span class="p">]</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:y</span><span class="p">,</span> <span class="n">label</span><span class="p">,</span> <span class="ss">type:</span> <span class="ss">:quantitative</span><span class="p">,</span> <span class="ss">title:</span> <span class="no">String</span><span class="o">.</span><span class="n">capitalize</span><span class="p">(</span><span class="n">label</span><span class="p">))</span>

    <span class="no">Vl</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">width:</span> <span class="mi">500</span><span class="p">,</span> <span class="ss">height:</span> <span class="mi">300</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">layers</span><span class="p">([</span><span class="n">epoch_layer</span><span class="p">,</span> <span class="n">examples_layer</span><span class="p">])</span>
    <span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">resolve</span><span class="p">(</span><span class="ss">:scale</span><span class="p">,</span> <span class="ss">x:</span> <span class="ss">:independent</span><span class="p">,</span> <span class="ss">y:</span> <span class="ss">:shared</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Kino</span><span class="o">.</span><span class="no">VegaLite</span><span class="o">.</span><span class="n">new</span><span class="p">()</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<h2 id="loss-plot">Loss plot</h2>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">losses</span> <span class="o">=</span> <span class="n">metrics</span><span class="p">[</span><span class="s2">"losses"</span><span class="p">]</span>

<span class="no">ClassifierPlot</span><span class="o">.</span><span class="n">plot_values</span><span class="p">(</span>
  <span class="n">losses</span><span class="p">[</span><span class="s2">"epochs_seen"</span><span class="p">],</span>
  <span class="n">losses</span><span class="p">[</span><span class="s2">"examples_seen"</span><span class="p">],</span>
  <span class="n">losses</span><span class="p">[</span><span class="s2">"train_values"</span><span class="p">],</span>
  <span class="n">losses</span><span class="p">[</span><span class="s2">"val_values"</span><span class="p">],</span>
  <span class="s2">"loss"</span>
<span class="p">)</span>
</code></pre></div></div>

<h2 id="accuracy-plot">Accuracy plot</h2>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">accuracies</span> <span class="o">=</span> <span class="n">metrics</span><span class="p">[</span><span class="s2">"accuracies"</span><span class="p">]</span>

<span class="no">ClassifierPlot</span><span class="o">.</span><span class="n">plot_values</span><span class="p">(</span>
  <span class="n">accuracies</span><span class="p">[</span><span class="s2">"epochs_seen"</span><span class="p">],</span>
  <span class="n">accuracies</span><span class="p">[</span><span class="s2">"examples_seen"</span><span class="p">],</span>
  <span class="n">accuracies</span><span class="p">[</span><span class="s2">"train_values"</span><span class="p">],</span>
  <span class="n">accuracies</span><span class="p">[</span><span class="s2">"val_values"</span><span class="p">],</span>
  <span class="s2">"accuracy"</span>
<span class="p">)</span>
</code></pre></div></div>

<h2 id="classification-loss-plot">Classification loss plot</h2>

<p><img src="/assets/2026/06/classification-loss-plot.png" alt="Classification training and validation loss plot" /></p>

<p>The plot shows how the classifier loss changes during fine-tuning. Training loss and validation loss both fall quickly during the first two epochs, from high initial values to near zero. After that, both curves stay low and close to each other, with only small validation-loss fluctuations. This suggests that the model learned the classification task quickly and that validation performance tracked training performance without a large overfitting gap.</p>

<h2 id="classification-accuracy-plot">Classification accuracy plot</h2>

<p><img src="/assets/2026/06/classification-accuracy-plot.png" alt="Classification training and validation accuracy plot" /></p>

<p>The accuracy plot shows that both training and validation accuracy reach almost perfect performance early in fine-tuning. Training accuracy reaches 1.0 around the second epoch, while validation accuracy stays very close to it with a small dip near the end. This matches the loss plot: the classifier converges quickly and generalizes well on the validation set.</p>

<h2 id="emlx-m3-gpu-support">EMLX M3 GPU support</h2>

<p>During classification training, I checked BEAM process CPU activity with Activity Monitor. GPU usage was 0%. Since my M3 Mac has an 18-core GPU, I explored how to use it from Elixir/Nx. There are several Elixir library options, and <a href="https://github.com/elixir-nx/emlx">EMLX</a> worked for this setup. Classification training dropped from about 1000 seconds to 120 seconds. Here is a small wrapper module for EMLX support:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">EMLXBackend</span> <span class="k">do</span>
  <span class="nv">@compile</span> <span class="p">{</span><span class="ss">:no_warn_undefined</span><span class="p">,</span> <span class="p">{</span><span class="no">EMLX</span><span class="p">,</span> <span class="ss">:__info__</span><span class="p">,</span> <span class="mi">1</span><span class="p">}}</span>
  <span class="nv">@compile</span> <span class="p">{</span><span class="ss">:no_warn_undefined</span><span class="p">,</span> <span class="p">{</span><span class="no">EMLX</span><span class="o">.</span><span class="no">Backend</span><span class="p">,</span> <span class="ss">:__info__</span><span class="p">,</span> <span class="mi">1</span><span class="p">}}</span>

  <span class="nv">@moduledoc</span> <span class="sd">"""
  Experimental EMLX backend selection helpers.

  This module is intentionally separate from the core training/model code. It
  lets tests or scripts opt into EMLX on Apple Silicon without changing the
  existing `llm_scratch` library modules.
  """</span>

  <span class="nv">@type</span> <span class="n">context</span> <span class="p">::</span> <span class="p">%{</span>
          <span class="ss">backend:</span> <span class="n">term</span><span class="p">(),</span>
          <span class="ss">compiler:</span> <span class="n">module</span><span class="p">(),</span>
          <span class="ss">previous_backend:</span> <span class="n">term</span><span class="p">(),</span>
          <span class="ss">previous_defn_options:</span> <span class="n">keyword</span><span class="p">(),</span>
          <span class="ss">selected:</span> <span class="ss">:emlx</span> <span class="o">|</span> <span class="ss">:exla</span>
        <span class="p">}</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Selects EMLX on the Apple GPU for the current process.

  ## Output

  Returns a context map containing the previous Nx backend and `Nx.Defn`
  options. Pass that map to `restore!/1` when the experiment is finished.
  """</span>
  <span class="nv">@spec</span> <span class="n">apple_gpu!</span><span class="p">()</span> <span class="p">::</span> <span class="n">context</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">apple_gpu!</span> <span class="k">do</span>
    <span class="n">ensure_loaded!</span><span class="p">()</span>

    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="n">previous_defn_options</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span><span class="o">.</span><span class="n">default_options</span><span class="p">()</span>
    <span class="n">backend</span> <span class="o">=</span> <span class="p">{</span><span class="no">EMLX</span><span class="o">.</span><span class="no">Backend</span><span class="p">,</span> <span class="ss">device:</span> <span class="ss">:gpu</span><span class="p">}</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">backend</span><span class="p">)</span>
    <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span><span class="o">.</span><span class="n">default_options</span><span class="p">(</span><span class="ss">compiler:</span> <span class="no">EMLX</span><span class="p">)</span>

    <span class="p">%{</span>
      <span class="ss">backend:</span> <span class="n">backend</span><span class="p">,</span>
      <span class="ss">compiler:</span> <span class="no">EMLX</span><span class="p">,</span>
      <span class="ss">previous_backend:</span> <span class="n">previous_backend</span><span class="p">,</span>
      <span class="ss">previous_defn_options:</span> <span class="n">previous_defn_options</span><span class="p">,</span>
      <span class="ss">selected:</span> <span class="ss">:emlx</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Selects the best available training backend for the current process.

  ## Output

  Returns the same context shape as `apple_gpu!/0`. It selects EMLX on Apple GPU
  when supported, otherwise falls back to `EXLA.Backend`.
  """</span>
  <span class="nv">@spec</span> <span class="n">apple_gpu_or_exla!</span><span class="p">()</span> <span class="p">::</span> <span class="n">context</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">apple_gpu_or_exla!</span> <span class="k">do</span>
    <span class="k">if</span> <span class="n">supported?</span><span class="p">()</span> <span class="k">do</span>
      <span class="n">apple_gpu!</span><span class="p">()</span>
    <span class="k">else</span>
      <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
      <span class="n">previous_defn_options</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span><span class="o">.</span><span class="n">default_options</span><span class="p">()</span>
      <span class="n">backend</span> <span class="o">=</span> <span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span>

      <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">backend</span><span class="p">)</span>

      <span class="p">%{</span>
        <span class="ss">backend:</span> <span class="n">backend</span><span class="p">,</span>
        <span class="ss">compiler:</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span><span class="o">.</span><span class="n">default_options</span><span class="p">()[</span><span class="ss">:compiler</span><span class="p">],</span>
        <span class="ss">previous_backend:</span> <span class="n">previous_backend</span><span class="p">,</span>
        <span class="ss">previous_defn_options:</span> <span class="n">previous_defn_options</span><span class="p">,</span>
        <span class="ss">selected:</span> <span class="ss">:exla</span>
      <span class="p">}</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Restores backend settings returned by `apple_gpu!/0`.

  ## Input Parameters

    * `context` - map returned by `apple_gpu!/0`.

  ## Output

  Returns `:ok`.
  """</span>
  <span class="nv">@spec</span> <span class="n">restore!</span><span class="p">(</span><span class="n">context</span><span class="p">())</span> <span class="p">::</span> <span class="ss">:ok</span>
  <span class="k">def</span> <span class="n">restore!</span><span class="p">(</span><span class="n">context</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">previous_backend</span><span class="p">)</span>
    <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span><span class="o">.</span><span class="n">default_options</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">previous_defn_options</span><span class="p">)</span>
    <span class="ss">:ok</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns whether EMLX modules are available.

  ## Output

  Returns `true` when the EMLX dependency is compiled and loadable, otherwise
  `false`.
  """</span>
  <span class="nv">@spec</span> <span class="n">available?</span><span class="p">()</span> <span class="p">::</span> <span class="n">boolean</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">available?</span> <span class="k">do</span>
    <span class="no">Code</span><span class="o">.</span><span class="n">ensure_loaded?</span><span class="p">(</span><span class="no">EMLX</span><span class="p">)</span> <span class="ow">and</span> <span class="no">Code</span><span class="o">.</span><span class="n">ensure_loaded?</span><span class="p">(</span><span class="no">EMLX</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns whether the EMLX Apple GPU backend can run a small tensor operation.

  ## Output

  Returns `true` when EMLX is compiled, the backend can be selected, and a small
  tensor can round-trip through the EMLX backend. Returns `false` when EMLX is
  unavailable or unsupported on the current machine.
  """</span>
  <span class="nv">@spec</span> <span class="n">supported?</span><span class="p">()</span> <span class="p">::</span> <span class="n">boolean</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">supported?</span> <span class="k">do</span>
    <span class="k">try</span> <span class="k">do</span>
      <span class="n">context</span> <span class="o">=</span> <span class="n">apple_gpu!</span><span class="p">()</span>

      <span class="k">try</span> <span class="k">do</span>
        <span class="p">[</span><span class="mf">1.0</span><span class="p">]</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">()</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">backend</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">()</span>

        <span class="no">true</span>
      <span class="k">after</span>
        <span class="n">restore!</span><span class="p">(</span><span class="n">context</span><span class="p">)</span>
      <span class="k">end</span>
    <span class="k">rescue</span>
      <span class="n">_</span> <span class="o">-&gt;</span> <span class="no">false</span>
    <span class="k">catch</span>
      <span class="n">_</span><span class="p">,</span> <span class="n">_</span> <span class="o">-&gt;</span> <span class="no">false</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">ensure_loaded!</span> <span class="k">do</span>
    <span class="k">unless</span> <span class="n">available?</span><span class="p">()</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="sd">"""
      EMLX backend is not available.

      Run:

          mix deps.get
          mix compile
      """</span>
    <span class="k">end</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<h2 id="using-llm-as-spam-filter">Using LLM as spam filter</h2>

<p>Finally, we can use our trained model for ham/spam classification. For that, we need one new helper function, <code class="language-plaintext highlighter-rouge">classify_review</code>. This function first pads or truncates the input text to the maximum length with the end-of-text padding token id. Then it runs the model forward pass to classify the message as ham or spam:</p>

<div class="language-diff highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="gh">diff --git a/lib/llm_scratch/text_generation.ex b/lib/llm_scratch/text_generation.ex
index 27a808d..e547503 100644
</span><span class="gd">--- a/lib/llm_scratch/text_generation.ex
</span><span class="gi">+++ b/lib/llm_scratch/text_generation.ex
</span><span class="p">@@ -206,6 +206,87 @@</span> defmodule LlmScratch.TextGeneration do
     generated_idx
   end
 
<span class="gi">+  @spec classify_review(
+          String.t(),
+          struct(),
+          String.t() | (String.t() -&gt; [integer()]),
+          term(),
+          nil | pos_integer(),
+          integer()
+        ) :: String.t()
+  @doc """
+  Classifies a review/message as `"spam"` or `"not spam"`.
+
+  This mirrors the book's classifier inference helper:
+
+      input_ids = tokenizer.encode(text)
+      input_ids = input_ids[:max_length]
+      input_ids += [pad_token_id] * (max_length - len(input_ids))
+      logits = model(input_tensor)[:, -1, :]
+      predicted_label = torch.argmax(logits, dim=-1).item()
+
+  ## Arguments
+
+    * `text` - message/review text to classify.
+    * `model` - GPT-style classifier model whose output head has two classes.
+    * `tokenizer` - either a Tiktoken model name or a one-argument tokenizer
+      function returning token ids.
+    * `device` - Nx backend target. Use `:default` or `nil` to keep tensors on
+      their current backend.
+    * `max_length` - optional fixed input length. The sequence is truncated to
+      `min(max_length, model context length)` and padded to that same length.
+      When `nil`, only the model context-length truncation is applied.
+    * `pad_token_id` - token id appended when padding is needed. Defaults to
+      GPT-2's end-of-text token id, `50256`.
+
+  ## Returns
+
+  Returns `"spam"` when the predicted label id is `1`; otherwise returns
+  `"not spam"`.
+  """
+  def classify_review(
+        text,
+        model,
+        tokenizer,
+        device \\ :default,
+        max_length \\ nil,
+        pad_token_id \\ 50_256
+      ) do
+    input_ids = encode_text(text, tokenizer)
+    supported_context_length = context_length(model)
+    target_length = target_input_length(input_ids, max_length, supported_context_length)
+
+    # Classifier inference uses the same fixed-length preparation as training:
+    # truncate long messages first, then pad short messages so the model sees
+    # the expected context width.
+    input_ids =
+      input_ids
+      |&gt; Enum.take(target_length)
+      |&gt; then(&amp;(&amp;1 ++ List.duplicate(pad_token_id, target_length - length(&amp;1))))
+
+    input_tensor =
+      input_ids
+      |&gt; Nx.tensor(type: {:s, 64})
+      |&gt; Nx.new_axis(0)
+      |&gt; maybe_transfer_tensor(device)
+
+    model = maybe_transfer_model(model, device)
+
+    # Calling forward/2 is inference-only in this codebase. Dropout is only used
+    # by train/3, and gradients are only created inside explicit value_and_grad
+    # calls, so this is the Nx counterpart to model.eval() and no_grad().
+    predicted_label =
+      model
+      |&gt; forward!(input_tensor)
+      |&gt; last_position_logits()
+      |&gt; Nx.argmax(axis: -1)
+      |&gt; Nx.backend_transfer(Nx.BinaryBackend)
+      |&gt; Nx.to_flat_list()
+      |&gt; List.first()
+
+    if predicted_label == 1, do: "spam", else: "not spam"
+  end
+
</span>   defp last_tokens(idx, context_size) do
     {_batch_size, seq_len} = Nx.shape(idx)
     length = min(seq_len, context_size)
<span class="p">@@ -303,6 +384,44 @@</span> defmodule LlmScratch.TextGeneration do
     end
   end
 
<span class="gi">+  defp encode_text(text, tokenizer) when is_binary(tokenizer) do
+    {:ok, token_ids} = Tiktoken.encode(tokenizer, text, ["&lt;|endoftext|&gt;"])
+    token_ids
+  end
+
+  defp encode_text(text, tokenizer) do
+    tokenizer.(text)
+  end
+
+  defp context_length(%{cfg: %{context_length: context_length}}), do: context_length
+
+  defp context_length(%{pos_emb: %{weight: weight}}) do
+    {context_length, _embedding_dim} = Nx.shape(weight)
+    context_length
+  end
+
+  defp target_input_length(input_ids, nil, supported_context_length) do
+    min(length(input_ids), supported_context_length)
+  end
+
+  defp target_input_length(_input_ids, max_length, supported_context_length) do
+    min(max_length, supported_context_length)
+  end
+
+  defp maybe_transfer_tensor(tensor, nil), do: tensor
+  defp maybe_transfer_tensor(tensor, :default), do: tensor
+
+  defp maybe_transfer_tensor(tensor, device) do
+    Nx.backend_transfer(tensor, device)
+  end
+
+  defp maybe_transfer_model(model, nil), do: model
+  defp maybe_transfer_model(model, :default), do: model
+
+  defp maybe_transfer_model(model, device) do
+    Nx.backend_transfer(model, device)
+  end
+
</span>   defp validate_idx_shape!(idx) do
     case Nx.shape(idx) do
       {_batch_size, _seq_len} -&gt;
</code></pre></div></div>

<p>And here is a test that performs the classification:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@tag</span> <span class="ss">:emlx</span>
  <span class="nv">@tag</span> <span class="ss">:train</span>
  <span class="n">test</span> <span class="s2">"6.8 classifies spam examples with saved classifier checkpoint"</span> <span class="k">do</span>
    <span class="n">device</span> <span class="o">=</span> <span class="n">use_accelerated_backend</span><span class="p">()</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="o">&amp;</span><span class="n">gpt2_compatible_token_ids</span><span class="o">/</span><span class="mi">1</span>
    <span class="n">train_dataset</span> <span class="o">=</span> <span class="no">SpamDataset</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="s2">"train.csv"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="ss">max_length:</span> <span class="no">nil</span><span class="p">)</span>
    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="s2">"ch6_spam_classifier_model_and_optimizer.nx"</span>

    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="p">%{</span><span class="ss">model_state_dict:</span> <span class="n">model</span><span class="p">}</span> <span class="o">=</span> <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">load_training_state!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

    <span class="n">text_1</span> <span class="o">=</span>
      <span class="s2">"You are a winner you have been specially"</span> <span class="o">&lt;&gt;</span>
        <span class="s2">" selected to receive $1000 cash or a $2000 award."</span>

    <span class="n">assert</span> <span class="no">TextGeneration</span><span class="o">.</span><span class="n">classify_review</span><span class="p">(</span>
             <span class="n">text_1</span><span class="p">,</span>
             <span class="n">model</span><span class="p">,</span>
             <span class="n">tokenizer</span><span class="p">,</span>
             <span class="ss">:default</span><span class="p">,</span>
             <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span>
           <span class="p">)</span> <span class="o">==</span> <span class="s2">"spam"</span>

    <span class="n">text_2</span> <span class="o">=</span>
      <span class="s2">"Hey, just wanted to check if we're still on"</span> <span class="o">&lt;&gt;</span>
        <span class="s2">" for dinner tonight? Let me know!"</span>

    <span class="n">assert</span> <span class="no">TextGeneration</span><span class="o">.</span><span class="n">classify_review</span><span class="p">(</span>
             <span class="n">text_2</span><span class="p">,</span>
             <span class="n">model</span><span class="p">,</span>
             <span class="n">tokenizer</span><span class="p">,</span>
             <span class="ss">:default</span><span class="p">,</span>
             <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span>
           <span class="p">)</span> <span class="o">==</span> <span class="s2">"not spam"</span>

    <span class="n">text_3</span> <span class="o">=</span>
      <span class="s2">"Elixir is cool distributed language, would you like to try it?"</span>

    <span class="n">assert</span> <span class="no">TextGeneration</span><span class="o">.</span><span class="n">classify_review</span><span class="p">(</span>
             <span class="n">text_3</span><span class="p">,</span>
             <span class="n">model</span><span class="p">,</span>
             <span class="n">tokenizer</span><span class="p">,</span>
             <span class="ss">:default</span><span class="p">,</span>
             <span class="n">train_dataset</span><span class="o">.</span><span class="n">max_length</span>
           <span class="p">)</span> <span class="o">==</span> <span class="s2">"not spam"</span>

  <span class="k">end</span>
</code></pre></div></div>

<h2 id="exercises">Exercises</h2>

<p>We got three exercises (that I forgot in first version of this post). In first exercise 6.1, instead of using maximal message length as a maximal size of all messages, we padded all messages to GPT model context length, which was 1024 tokens. Result is that training was 30x longer, with worse training accuracy! In exercise 6.2, we retraind for classification whole GPT2 model, not just three layers (last transformer, out head and final linear). I got training accuraccy of 100%, but training was 3x longer. In final exercise, we tuned on first message token, not the last one. Training accuraccy was worse but training duration was almost the same.</p>]]></content><author><name></name></author><category term="llm-from-scratch" /><category term="gpt" /><category term="tutorials" /><category term="llm" /><category term="gpt" /><category term="workbook" /><summary type="html"><![CDATA[This post fine-tunes a pretrained GPT-2 model for SMS spam classification in Elixir.]]></summary></entry><entry><title type="html">Recap of training a GPT model</title><link href="https://karlosmid.com/2026/06/recap-on-training-GPT-model/" rel="alternate" type="text/html" title="Recap of training a GPT model" /><published>2026-06-02T09:00:00+00:00</published><updated>2026-06-02T09:00:00+00:00</updated><id>https://karlosmid.com/2026/06/recap-on-training-GPT-model</id><content type="html" xml:base="https://karlosmid.com/2026/06/recap-on-training-GPT-model/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>I am recapping the blog post <a href="/pretraining-on-unlabeled-data-in-elixir">Pretraining on unlabeled data in Elixir</a> by answering workbook questions from chapter 5 of Build an LLM from Scratch and reading Giles’s blog posts on this chapter:</p>
<ul>
  <li><a href="https://www.gilesthomas.com/2025/10/llm-from-scratch-20-starting-training-cross-entropy-loss">part 20</a></li>
  <li><a href="https://www.gilesthomas.com/2025/10/llm-from-scratch-21-perplexed-by-perplexity">part 21</a></li>
  <li><a href="https://www.gilesthomas.com/2025/10/llm-from-scratch-22-finally-training-our-llm">part 22</a></li>
</ul>

<h2 id="giless-blog-posts">Giles’s blog posts</h2>

<p>In part 20, Giles explains his understanding of the cross-entropy loss function and its input parameters. That function is used to evaluate how good the current token prediction is. Based on that value, we correct the LLM parameters so they start moving toward the correct prediction value. We calculate the loss between the logit value for the token predicted during the training step and the token that should be predicted: the target. We slide through the input token by token by shifting each run’s input one token to the right. The validation value is the target token, which is the next token, but we do not use it as input during training. We want the loss to have a value of zero. We calculate loss with a function called cross-entropy loss. Giles first explains the “how” of the entropy loss function. When the model’s predicted probability for the target token is not 1 (meaning we are not 100% certain that the target token should be next), we want that x value to move toward 1. As we do that, the logit values for the other tokens decrease, out of the box! Giles says that <code class="language-plaintext highlighter-rouge">1 - x</code> must become zero; then our model is correctly trained. But because x is in the range <code class="language-plaintext highlighter-rouge">[0,1]</code>, it is better to use a logarithmic function with base Euler’s number, because <code class="language-plaintext highlighter-rouge">ln 1 = 0</code>. This means that if the target logit value is 1, our loss is zero. So we have <code class="language-plaintext highlighter-rouge">-ln x</code>. We use the minus sign because <code class="language-plaintext highlighter-rouge">ln x</code>, when x is close to zero, is a large negative number. As we apply this function to the logits tensor, in the end, we calculate the mean value of all target token losses.</p>

<p>But what is cross entropy? In thermodynamics, entropy is a measurement of how “messy” a system is. Then we have information theory, which is about messages. What is the quantity of information in a message? <code class="language-plaintext highlighter-rouge">00000</code> has less information than <code class="language-plaintext highlighter-rouge">00001</code>. Since logits represent probability distributions, with entropy we would like to measure how “messy” those distributions are.</p>

<p>Let’s say that our logits probability function is <code class="language-plaintext highlighter-rouge">p(x)</code>, where x is the next token. Then <code class="language-plaintext highlighter-rouge">H(p)</code> is the entropy of the probability distribution, represented as a number. Based on the probability <code class="language-plaintext highlighter-rouge">p(x)</code> of event x, we need to take into account how surprising this event is. For example, the probability that I will win the Croatian lottery (if I play) is very low, so if I win, it would be a big surprise to my friends (of course, I would not tell them that I won). To measure the level of surprise, we can start with <code class="language-plaintext highlighter-rouge">1 - p(x)</code>. The problem with the function <code class="language-plaintext highlighter-rouge">surprise = 1 - p(x)</code> is that it is linear, and the deviation is constant. So the change in surprise for very likely, very unlikely, and <code class="language-plaintext highlighter-rouge">50:50</code> events is the same.</p>

<h3 id="surprise-graph">Surprise graph</h3>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>surprise = 1 - p(x)

surprise
  1.0 |*                   
      |  *                 
  0.8 |    *               
      |      *             
  0.6 |        *           
      |          *         
  0.4 |            *       
      |              *     
  0.2 |                *   
      |                  * 
  0.0 +--------------------*
      0.0  0.2  0.4  0.6  0.8  1.0
                    p(x)
</code></pre></div></div>

<h3 id="negative-natural-logarithm-graph">Negative natural logarithm graph</h3>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>surprise = -ln p(x)

surprise
  inf |*                   
      | *                  
  4.0 |  *                 
      |   *                
  3.0 |    *               
      |      *             
  2.0 |        *           
      |           *        
  1.0 |               *    
      |                  * 
  0.0 +--------------------*
      0.0  0.2  0.4  0.6  0.8  1.0
                    p(x)

As p(x) gets close to 0, -ln p(x) grows toward infinity.
When p(x) is 1, -ln p(x) is 0.
</code></pre></div></div>
<p>We want a different scale, so when p(x) changes by a small amount, the surprise number changes by a much bigger amount. We use a logarithmic function with base Euler’s number e, <code class="language-plaintext highlighter-rouge">surprise = - ln p(x)</code>.
Next, we scale the surprise: <code class="language-plaintext highlighter-rouge">-p(x) ln p(x)</code>
And we sum it across x: <code class="language-plaintext highlighter-rouge">- sum(p(x) ln p(x))</code></p>

<p>But we need cross entropy to calculate loss. We need to cross something with something. The output of the LLM is a set of logits, which we convert into a probability distribution for each input, and this probability distribution also has its own entropy. We mark this probability as <code class="language-plaintext highlighter-rouge">q(x)</code> to distinguish it from <code class="language-plaintext highlighter-rouge">p(x)</code>, the real-world distribution (what the next token is in the training dataset).</p>

<p>We calculate cross entropy between the LLM prediction and the real-world distribution:</p>

<p><code class="language-plaintext highlighter-rouge">H(p,q) = - sum(p(x) ln q(x))</code></p>

<p>In part 21, Giles talks about perplexity. Perplexity was just a side note in the book, but Giles did a side quest to dig deeper into it! Perplexity measures how well the LLM’s token probability distribution is aligned with the actual word distribution. We calculate it by raising the value of the cross-entropy loss to the base that we used as the logarithm base when calculating cross-entropy loss. In our case, this is Euler’s number, e. Giles then calculates the perplexity formula using high school math for logarithms. Read his part 21 for more details.</p>

<p>In part 22, Giles writes about training our GPTModel from scratch (it does not go deep into the simple training algorithm). He does not go on any side quests, but it is worth reading (it is a short read) in order to deepen your understanding of chapter 5.</p>

<h2 id="workbook-questions">Workbook questions</h2>

<p>1 What is the primary purpose of the cross-entropy loss function in LLM training?</p>

<p>A) To generate text samples for evaluation purposes.<br />
B) To prevent the model from overfitting to the training data.<br />
C) To evaluate the model’s performance on a specific task, such as text classification.<br />
D) To measure the difference between the model’s predicted probability distribution of tokens and the actual distribution of tokens in the training data.</p>

<p>My answer: D<br />
Correct answer: +</p>

<p>2 What is the purpose of temperature scaling in text generation?</p>

<p>A) To control the randomness and diversity of the generated text by adjusting the probabiliy distribution of tokens.<br />
B) To improve the model’s accuracy in predicting the next token.<br />
C) To reduce the computational cost of text generation.<br />
D) To prevent the model from overfitting to the training data.</p>

<p>My answer: D<br />
Correct answer: A</p>

<p>3 What is the primary goal of pretraining a large language model (LLM)?</p>

<p>A) To improve the model’s ability to generate coherent and gramatically correct text.<br />
B) To fine-tune the model for a specific task, such as text classification.<br />
C) To learn general language patterns and representations from a massive amount of text data.<br />
D) To reduce the computational cost of training the model.</p>

<p>My answer: C<br />
Correct answer: +</p>

<p>4 What is the main benefit of using pretrained weights from OpenAI for LLMs?</p>

<p>A) It eliminates the need for extensive and expensive pretraining from scratch.<br />
B) It reduces the risk of overfitting to the training data.<br />
C) It guarantees that the model will perform well on any task.<br />
D) It simplifies the process of fine-tuning the model for specific task.</p>

<p>My answer: A<br />
Correct answer: +</p>

<p>5 What is the primary advantage of saving the model’s state dictionary in PyTorch?</p>

<p>A) It improves the model’s performance on unseen data.<br />
B) It allows you to load and reuse the trained model without retraining it from scratch.<br />
C) It prevents the model from overfitting to the training data.<br />
D) It reduces the computational costs of training the model.</p>

<p>My answer: B<br />
Correct answer: +</p>

<h3 id="51-evaluating-generative-text-models">5.1 Evaluating generative text models</h3>

<p>1 What is the purpose of the <code class="language-plaintext highlighter-rouge">generate_text_simple</code> function and how does it work?</p>

<p>Purpose of that function is to predict an <code class="language-plaintext highlighter-rouge">n</code> next tokens based on input context. It is doing that in <code class="language-plaintext highlighter-rouge">n</code> steps, and in each step is predicting next token by using GPTModel prediction (eval function). First step consists of context, next step consist of context + predicted token and so on. We finish when steps number is equal to number of tokens that we should predict.</p>

<p>Correct answer: <em>The generate_text_simple function generates text by taking a starting context, converting it to token IDs, feeding it into the GPT model, and then converting the model’s output logits back into token IDs, which are then decoded into text.</em></p>

<p>2 Explain the concept of text generation loss and its significance in evaluationg the quality of generated text.</p>

<p>Text generation loss if functions that returns number that shows us how far is predicted token from token that is next in training data. How far is calculated using back propagation algorithm of gradients. It is important function, because it shows us how good is GPTModel trained to predict grammatically correct next token.</p>

<p>Correct answer: <em>Text generation loss is a numerical measure used to evaluate the quality of generated text. It quantifies the difference between the model’s predicted output (token probabilities) and the actual desired output (target tokens). A lower loss indicates better text generation quality.</em></p>

<p>3 Describe the role of backpropagation in training an LLM and how it relates to the text generation loss.</p>

<p>In backpropagation algoritm we calculate gradient change as we are moving toward the begining of input text.</p>

<p>Correct answer: <em>Backpropagation is a technique used to update the model’s weights during training. It uses the text generation loss to adjust the weights so that the model produces outputs closer to the target tokens, thereby minimizing the loss and improving the quality of generated text.</em></p>

<p>4 What is corss entropy loss and how is it used in evaluating LLMs?</p>

<p>Cross entropy loss is number that shows the difference between the predicted token distribution and token distribution of trained text.  If cross entropy loss is close to zero, that means that our model is predicting coherent text.</p>

<p>Correct answer: <em>Cross entropy loss is a measure of the difference between two probability distributions, typically the true distribution of labels (tokens) and the predicted distribution from the model. It is used to evaluate the performance of LLMs by quantifying how well the model’s predicted probability distribution matches the actual distribution of tokens in the dataset.</em></p>

<p>5 Label this diagram by assinging the label numbers to the matching description:</p>

<p><img src="/assets/2026/06/gpt-argmax-token-prediction.png" alt="GPT argmax token prediction diagram" /></p>

<table>
  <thead>
    <tr>
      <th>Label</th>
      <th>Description</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>1</td>
      <td>Predicted token IDs are index positions with highest probability</td>
    </tr>
    <tr>
      <td>2</td>
      <td>Input text</td>
    </tr>
    <tr>
      <td>3</td>
      <td>Map input text to token IDs with vocabulary</td>
    </tr>
    <tr>
      <td>4</td>
      <td>Output text generated by LLM</td>
    </tr>
    <tr>
      <td>5</td>
      <td>Map index positions back into text with inverse vocabulary</td>
    </tr>
    <tr>
      <td>6</td>
      <td>Seven-dimensional probability row vector for each input vector</td>
    </tr>
    <tr>
      <td>7</td>
      <td>Index position with highest probability from argmax function</td>
    </tr>
  </tbody>
</table>

<p>My answer: 1 =&gt; 3, 2 =&gt; 7, 3 =&gt; 1, 4 =&gt; 6, 5 =&gt; 5, 6 =&gt; 2, 7 =&gt; 4</p>

<p>Correct answer: As my answers with following change: 2 =&gt; 6, 4 =&gt; 7</p>

<p>6 Explain the concept of perplexity and its relationship to cross entropy loss.</p>

<p>My answer: Perplexity describes how predicted token distributions follows distribution of input text. When we do power of euler number with cross entroply loss, we get perplexity value.</p>

<p>Correct answer: <em>Perplexity is a measure derived from cross entropy loss that provides a more interpretable way to understand the uncertainty of a model in predicting the next token. It represents the effective vocabulary size about which the model is uncertain at each step. A lower perplexity indicates less uncertainty and better model performance.</em></p>

<p>7 Pieces of code have been removed from two places in this listing. Which of these terms have been removed and where should they go?</p>

<p>A softmax B tensor C forward D grad</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="c1"># initialize our gpt model</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>
    <span class="c1"># predict next input token</span>
    <span class="n">logits</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">__1__</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">inputs</span><span class="p">)</span>
    <span class="c1"># using softmax get probability values (this is redundant step)</span>
    <span class="n">probas</span> <span class="o">=</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">__2__</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
    <span class="c1"># get token ids that have maximal logit values</span>
    <span class="n">token_ids</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="ss">keep_axis:</span> <span class="no">true</span><span class="p">)</span>

</code></pre></div></div>

<p>My answer: 1 =&gt; forward, 2 =&gt; softmax</p>

<p>Correct answer: +</p>

<p>I’m still far off in being able to explain these concepts in my own words!</p>

<h3 id="52-training-an-llm">5.2 Training an LLM</h3>

<p>1 What is the purpose of <code class="language-plaintext highlighter-rouge">train_model_simple</code> function and what are its key components?</p>

<p>My answer: This function is training our GPT model based on input data that is organized in batches. Input data is split in 90/10 ratio into traning and validation set. Using validation set we check how good is our training. It repeats the training by the number of epochs. At each epoch step, it calculates entropy loss value for the training step. Based on those backpropagation values and optimizer, it updates model weight parameters so we get better entropy loss in the next epoch step.</p>

<p>Correct answer: <em>The train_model_simple function implements a basic training loop for the LLM. It iterates over epochs, processes batches, calculates loss, updates weights, and evaluates the model’s performance using validation data.</em></p>

<p>My answer is pretty close to correct one.</p>

<p>2 Explain the role of <code class="language-plaintext highlighter-rouge">evaluate_model</code> function in the training process.</p>

<p>My answer: <code class="language-plaintext highlighter-rouge">evaluate_model</code> functions calculates entropy loss for current training step that is split into training and validation part.</p>

<p>Correct answer: <em>The evaluate_model function calculates the loss on both the training and validation sets to assess the model’s performance. It ensures the model is in evaluation mode, disabling gradient tracking and dropout for accurate evaluation.</em></p>

<p>I missed the dropout is disabled as we are in model <code class="language-plaintext highlighter-rouge">eval</code> mode.</p>

<p>3 Pieces of code have been removed from three places in this listing. Which of these terms have been removed and where should they go?</p>

<p>A epochs_seen, B tokens_seen, C losses, D epochs, E loss</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LossPlot</span> <span class="k">do</span>
  <span class="k">def</span> <span class="n">plot_losses</span><span class="p">(</span><span class="n">epochs_seen</span><span class="p">,</span> <span class="n">tokens_seen</span><span class="p">,</span> <span class="n">train_losses</span><span class="p">,</span> <span class="n">val_losses</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">loss_data</span> <span class="o">=</span>
      <span class="p">[</span>
        <span class="p">{</span><span class="s2">"Training loss"</span><span class="p">,</span> <span class="n">train_losses</span><span class="p">},</span>
        <span class="p">{</span><span class="s2">"Validation loss"</span><span class="p">,</span> <span class="n">val_losses</span><span class="p">}</span>
      <span class="p">]</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">flat_map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">series</span><span class="p">,</span> <span class="n">losses</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">([</span><span class="n">__1__</span><span class="p">,</span> <span class="n">__2__</span><span class="p">,</span> <span class="n">__3__</span><span class="p">])</span>
        <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">epoch</span><span class="p">,</span> <span class="n">tokens</span><span class="p">,</span> <span class="n">loss</span><span class="p">}</span> <span class="o">-&gt;</span>
          <span class="p">%{</span>
            <span class="ss">series:</span> <span class="n">series</span><span class="p">,</span>
            <span class="ss">epoch:</span> <span class="n">epoch</span><span class="p">,</span>
            <span class="ss">tokens_seen:</span> <span class="n">tokens</span><span class="p">,</span>
            <span class="ss">loss:</span> <span class="n">loss</span>
          <span class="p">}</span>
        <span class="k">end</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; A, 2 =&gt; B, 3 =&gt; C</p>

<p>Correct answer: +</p>

<p>4 What is the purpose of the <code class="language-plaintext highlighter-rouge">generate_and_print_sample</code> function and how does it work?</p>

<p>My answer: Purpose of the <code class="language-plaintext highlighter-rouge">generate_and_print_sample</code> function is to see how our training is doing by printing for current epoch predicted text of size 50 charachters, based on current epoch context.</p>

<p>Correct answer: <em>The generate_and_print_sample function generates text samples from the model to visually assess its progress during training. It takes a text snippet as input, converts it to token IDs, and uses the generate_text_simple function to generate new text.</em></p>

<p>I missed to describe how actually works this function.</p>

<p>5 What is AdamW and why is it preferred over Adam for training LLMs?</p>

<p>My answer: AdamW is training optimizer that updates GPTModel weights based on current epoch gradients values.</p>

<p>Correct answer: <em>AdamW is a variant of the Adam optimizer that improves weight decay, a technique that penalizes larger weights to prevent overfitting. This makes AdamW more effective for regularizing and generalizing LLMs.</em></p>

<p>I missed that it is about decay and large weight values</p>

<p>6 What is the significance of the training and validation loss curves in Figure 5.12?</p>

<p>My answer: Significance is that it shows that for a several starting epochs steps, validation and training loss are decreasing. But then model is starting to “memorize” the training set, and training loss keeps decreasing, but validation set, that is not part of training set, stagnates. That means that model started just repeating what it learned from the training set. When context becomes something not seen by the model in training set, prediction of next token becomes poor.</p>

<p>Correct answer: <em>The curves show that the model initially learns well, but the training loss continues to decrease while the validation loss stagnates. This indicates overfitting, where the model memorizes the training data instead of generalizing to new data.</em></p>

<p>My answer is correct!</p>

<p>7 Match the terms on the left with the description on the right:</p>

<table>
  <thead>
    <tr>
      <th> </th>
      <th> </th>
      <th> </th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>Text Generation<br />Loss</td>
      <td> </td>
      <td>a common measure in machine learning that quantifies the difference between two probability distributions, typically the true distribution of labels and the predicted distribution from a model.</td>
    </tr>
    <tr>
      <td>Cross Entropy Loss</td>
      <td> </td>
      <td>a standard technique for training deep neural networks that involves updating the model’s weights to minimize the difference between the model’s predicted output and the actual desired output.</td>
    </tr>
    <tr>
      <td>Perplexity</td>
      <td> </td>
      <td>a measure used alongside cross entropy loss to evaluate the performance of models in tasks like language modeling, providing a more interpretable way to understand the uncertainty of a model in predicting the next token in a sequence.</td>
    </tr>
    <tr>
      <td>Backpropagation</td>
      <td> </td>
      <td>a numerical measure used to assess the quality of text generated during training, indicating how well the model’s predictions align with the desired target text.</td>
    </tr>
  </tbody>
</table>

<p>My answer: 1 =&gt; 3, 2 =&gt; 4, 3 =&gt; 1, 4 =&gt; 2</p>

<p>Correct answer: 2 =&gt; 1, 3 =&gt; 4</p>

<p>My answer is partially correct. Cross entropy is about diagrams diff, while perplexity is number (certinity)</p>

<h2 id="53-decoding-strategies-to-control-randomness">5.3 Decoding strategies to control randomness</h2>

<p>1 What is the purpose of temperature scaling in text generation?</p>

<p>My answer: Purpose of temperature scaling in text generation is to make the model more creative. For lower temperatures logits distribution is more stepier and we chose for the next token the most probable one. Model is predictive. But for higher temperatures, logits distribution is more broad, and our model is less predictive.</p>

<p>Correct answer: <em>Temperature scaling adjusts the probability distribution of the next token by dividing the logits by a temperature value. Higher temperatures create a more uniform distribution, leading to more diverse outputs, while lower temperatures sharpen the distribution, favoring the most likely tokens.</em></p>

<p>I did not mention that we are dividing the logits with temperature value.</p>

<p>2 Explain how top-k sampling works and it benefit in text generation.</p>

<p>My answer: top-k sampling is when we narrow the set of probabilites that are candidates for the next token. If we choose for the next token between all possible logits, our module could predict sometimes gibberish text. To avoid this, for example for top-3, we only select from top three logits, and other logits are set to inf value (so probability for inf is 0). Benefit is that model is still creative, but still coherent in generating the next token.</p>

<p>Correct answer: <em>Top-k sampling selects only the top-k most likely tokens for the next token prediction. This helps to reduce the generation of nonsensical or grammatically incorrect text by focusing on the most probable options.</em></p>

<p>3 Describe the difference between the greedy decoding and probabilistic sampling in text generation.</p>

<p>My answer: In greedy encoding, we always select token that has logit with biggest value, the most probable token. In probabilistic sampling, we could select next token from a set of tokens, where probability that we select a token is based on token logit value. Higher the logit value, higher the probability that we would select that token.</p>

<p>Correct answer: <em>Greedy decoding always selects the token with the highest probability, while probabilistic sampling chooses tokens based on their probability distribution. Probabilistic sampling introduces randomness and can lead to more diverse outputs.</em></p>

<p>4 What is the purpose of the Inf Masking step in the top-k sampling process?</p>

<p>My answer: To set logits not in the top-k to negative infinity, removing them from consideration.</p>

<p>Correct answer: +</p>

<p>5 How does <code class="language-plaintext highlighter-rouge">generate</code> function incorporate temperature scaling and top-k sampling?</p>

<p>My answer: <code class="language-plaintext highlighter-rouge">generate</code> function first evaluates model with current context text in order to get logits for the last token in the input context. Then it masks values in those logits to inf value, for all values except for the top-k logits. Afer that, logits are scalled for temperature value, and in the end we do probability sampling.</p>

<p>Correct answer: <em>The generate function allows for the specification of temperature and top-k values. If a temperature is provided, the logits are scaled accordingly. If a top-k value is given, the logits are masked to exclude tokens outside the top-k most likely options.</em></p>

<h2 id="54-loading-and-saving-model-weights">5.4 Loading and saving model weights</h2>

<p>1 Why is it important to save the model weights after training a large language model?</p>

<p>My answer: It is important because training the large language model requires large hardware resouces. By saving it, do not need to traing it again. We could also continue training that same model.</p>

<p>Correct answer: <em>Saving the model weights allows you to reuse the trained model without retraining it from scratch, saving significant computational time and resources. This is especially important for large language models, which can take a long time to train.</em></p>

<p>2 What is recommended way to save a Nx model’s weights?</p>

<p>My answer: Recommended way is to use <code class="language-plaintext highlighter-rouge">Nx.Container</code> behaviour.</p>

<p>Correct answer: +</p>

<p>3 What is the purpose of the <code class="language-plaintext highlighter-rouge">model.eval()</code> method?</p>

<p>My answer: Purpose of this method is to predict the next token based on the given input context.</p>

<p>Correct answer: <em>The model.eval() method switches the model to evaluation mode, disabling dropout layers. This is important for inference, as we don’t want to randomly drop out information during prediction</em></p>

<p>I did not mention that dropout is skipped and why is skipped.</p>

<p>4 Why is important to save the optimizer state when saving a model?</p>

<p>My answer: It is important because we also need hardware to compute optimizer weights.</p>

<p>Correct answer: <em>Saving the optimizer state allows you to resume training from where you left off. This is important for adaptive optimizers like AdamW, which store historical data to adjust learning rates dynamically. Without the optimizer state, the model may learn suboptimally or fail to converge properly.</em></p>

<p>AdamW has historicaly data during the training.</p>

<p>5 Pieces of the code have been removed from two places in this listing. Which of these terms have been removed and where should they go?</p>

<p>A write_serialized B load C save</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">def</span> <span class="n">save_training_state!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optimizer</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">__1__!</span><span class="p">(</span>
      <span class="p">%{</span>
        <span class="ss">model_state_dict:</span> <span class="n">model</span><span class="p">,</span>
        <span class="ss">optimizer_state_dict:</span> <span class="n">optimizer</span>
      <span class="p">},</span>
      <span class="n">path</span>
    <span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; A</p>

<p>Correct answer: +</p>

<p>6 Match the terms on the left with the description on the right:</p>

<table>
  <thead>
    <tr>
      <th> </th>
      <th> </th>
      <th> </th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>AdamW</td>
      <td> </td>
      <td>the process of iterating over the training data multiple times, calculating the loss for each batch, and updating the model weights to minimize the loss</td>
    </tr>
    <tr>
      <td>Training<br />Loop</td>
      <td> </td>
      <td>a variant of the Adam optimizer that improves the weight decay approach, which aims to minimize model complexity and prevent overfitting by penalizing larger weights</td>
    </tr>
    <tr>
      <td>Overfitting</td>
      <td> </td>
      <td>the loss calculated on a separate dataset that is not used for training, which provides an estimate of the model’s performance on unseen data</td>
    </tr>
    <tr>
      <td>Validation<br />Loss</td>
      <td> </td>
      <td>a phenomenon that occurs when a model learns the training data too well and performs poorly on unseen data</td>
    </tr>
  </tbody>
</table>

<p>My answer: 1 =&gt; 2, 2 =&gt; 1, 3 =&gt; 4, 4 =&gt; 3</p>

<p>Correct answer: +</p>

<h2 id="55-loading-pretrained-weights-from-openai">5.5 Loading pretrained weights from OpenAI</h2>

<p>1 What is the primary advantage of using pretrained weights from OpenAI for a GPT-2 model?</p>

<p>My answer: The primary advantage is that we get GPT-2 model that is ready for usage and we do not need expensive hardware to pretrain it.</p>

<p>Correct answer: <em>Using pretrained weights from OpenAI eliminates the need for extensive training on a large corpus, saving significant time and computational resources.</em></p>

<p>2 What are the key components of the <code class="language-plaintext highlighter-rouge">settings</code> and <code class="language-plaintext highlighter-rouge">params</code> dictionaries obtained from OpenAI’s GPT-2 model weights?</p>

<p>My answer: Key components of <code class="language-plaintext highlighter-rouge">settings</code> are model settings: vocabulary size, context size, no of transformer layers, no of data heads, and embedings dimension. params contain weights for embedings, positions, …</p>

<p>Correct answer: <em>The settings dictionary contains the LLM architecture settings, while the params dictionary holds the actual weight tensors for the model’s layers.</em></p>

<p>3 How are the pretrained weights from OpenAI loaded into a custom GPTModel instance?</p>

<p>My answer: Weights are loaded from files where weights are in tensor flow format.</p>

<p>Correct answer: <em>The load_weights_into_gpt function carefully matches the weights from OpenAI’s implementation with the corresponding layers in the custom GPTModel instance, ensuring consistency and functionality.</em></p>

<p>4 What is the significance of the <code class="language-plaintext highlighter-rouge">model_configs</code> dictionary in the context of loading pretrained weights?</p>

<p>My answer: <code class="language-plaintext highlighter-rouge">model_configs</code> is a set of GPT model settings grouped by GPT model size.</p>

<p>Correct answer:  <em>The model_configs dictionary provides the specific architectural settings for different GPT-2 model sizes, allowing for the selection and loading of weights for the desired model.</em></p>

<p>5 Why is necessary to update the NEW_CONFIG dictionary with <code class="language-plaintext highlighter-rouge">context_length</code> and <code class="language-plaintext highlighter-rouge">qkv_bias</code> settings?</p>

<p>My answer: It is necessary because OpenAi used different values for those two parameters than what we used in book examples.</p>

<p>Correct answer: <em>The pretrained GPT-2 models from OpenAI were trained with a different context_length and used bias vectors in the attention module, requiring these settings to be updated for compatibility.</em></p>

<p>6 Match the terms on the left with the description on the right:</p>

<table>
  <thead>
    <tr>
      <th> </th>
      <th> </th>
      <th> </th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>State Dict</td>
      <td> </td>
      <td>a dictionary containing the parameters of each layer in a PyTorch model.</td>
    </tr>
    <tr>
      <td>Model<br />Weights</td>
      <td> </td>
      <td>a dictionary containing the internal state of the optimizer, such as the learning rate and momentum.</td>
    </tr>
    <tr>
      <td>Optimizer<br />State</td>
      <td> </td>
      <td>the parameters that are learned during the training process and are used to make predictions.</td>
    </tr>
    <tr>
      <td>Evaluation<br />Mode</td>
      <td> </td>
      <td>a mode in which the model is used for inference, and dropout layers are disabled.</td>
    </tr>
  </tbody>
</table>

<p>My answer: 1 =&gt; 1, 2 =&gt; 3, 3 =&gt; 2, 4 =&gt; 4</p>

<p>Correct answer: +</p>]]></content><author><name></name></author><category term="llm-from-scratch" /><category term="gpt" /><category term="tutorials" /><category term="llm" /><category term="gpt" /><category term="workbook" /><summary type="html"><![CDATA[Workbook notes on GPT training, cross-entropy loss, perplexity, and model weights.]]></summary></entry><entry><title type="html">Pretraining on unlabeled data in Elixir</title><link href="https://karlosmid.com/2026/05/pretraining-on-unlabeled-data-in-elixir/" rel="alternate" type="text/html" title="Pretraining on unlabeled data in Elixir" /><published>2026-05-29T09:00:00+00:00</published><updated>2026-05-29T09:00:00+00:00</updated><id>https://karlosmid.com/2026/05/pretraining-on-unlabeled-data-in-elixir</id><content type="html" xml:base="https://karlosmid.com/2026/05/pretraining-on-unlabeled-data-in-elixir/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>This post follows Chapter 5 of <em>Build a Large Language Model (From Scratch)</em> and ports the pretraining workflow to Elixir. We start by evaluating an untrained GPT model, calculating cross-entropy loss for generated text, and measuring training and validation losses on <em>The Verdict</em>. Then we implement a simple immutable Elixir training loop with <code class="language-plaintext highlighter-rouge">Nx.Container</code>, <code class="language-plaintext highlighter-rouge">Nx.Defn</code>, dropout-aware training, and AdamW updates. The model can learn the tiny dataset, but the validation loss shows clear overfitting, so we add decoding controls with temperature scaling and top-k sampling. Finally, we save and reload model and optimizer checkpoints, load OpenAI GPT-2 weights, and compare the much better losses and generated text from pretrained GPT-2 models.</p>

<p>You can find the source code in my <a href="https://github.com/karlosmid/llm-from-scratch">GitHub repository</a>.</p>

<h2 id="evaluating-generative-text-models">Evaluating generative text models</h2>

<p>Ah, testing. You cannot do anything these days without testing. But how do we test an LLM? For now, we know how to generate text with an LLM. To evaluate generative text models, we need to evaluate the generated text and calculate training and validation losses.</p>

<h3 id="using-gpt-to-generate-text">Using GPT to generate text</h3>

<p>Sebastian starts where we finished chapter 4: text generation with an untrained model, which produces gibberish text. The GPTModel configuration uses a context length of 256 tokens so we can train that model on our laptops, in my case an M3 Pro chip and 36 GB of RAM. We introduce two helper functions, tokens =&gt; ids and ids =&gt; tokens, because we will use those operations a lot in this chapter.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">TextUtils</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Token/text conversion helpers for GPT-style token id tensors.
  """</span>

  <span class="nv">@end_of_text</span> <span class="s2">"&lt;|endoftext|&gt;"</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Encodes `text` with `tokenizer` and returns token ids shaped `{1, seq_len}`.

  ## Examples

      iex&gt; token_ids = LlmScratch.TextUtils.text_to_token_ids("Every effort moves you", "code-davinci-002")
      iex&gt; Nx.shape(token_ids)
      {1, 4}
      iex&gt; Nx.to_flat_list(token_ids)
      [6109, 3626, 6100, 345]
  """</span>
  <span class="nv">@spec</span> <span class="n">text_to_token_ids</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">text</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">text</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">encoded</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="n">text</span><span class="p">,</span> <span class="p">[</span><span class="nv">@end_of_text</span><span class="p">])</span>

    <span class="n">encoded</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">new_axis</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Decodes token ids shaped `{1, seq_len}` back to text with `tokenizer`.

  ## Examples

      iex&gt; token_ids = Nx.tensor([[6109, 3626, 6100, 345]], type: {:s, 64})
      iex&gt; LlmScratch.TextUtils.token_ids_to_text(token_ids, "code-davinci-002")
      "Every effort moves you"
  """</span>
  <span class="nv">@spec</span> <span class="n">token_ids_to_text</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">token_ids_to_text</span><span class="p">(%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">flat</span> <span class="o">=</span>
      <span class="n">token_ids</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="ss">axes:</span> <span class="p">[</span><span class="mi">0</span><span class="p">])</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">()</span>

    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">text</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">,</span> <span class="n">flat</span><span class="p">)</span>
    <span class="n">text</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>Here is a test where we use this model with those two helper functions:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"5.1.1 generate_text_simple generates text from a GPT-124M start context"</span> <span class="k">do</span>
    <span class="c1"># set EXLA for faster computing</span>

    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="c1"># we set context_length to 256 so we can train this model in this century on m3 chip</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">256</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.0</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">start_context</span> <span class="o">=</span> <span class="s2">"Every effort moves you"</span>
    <span class="n">encoded_tensor</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">start_context</span><span class="p">,</span> <span class="s2">"code-davinci-002"</span><span class="p">)</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="n">out</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate_text_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">encoded_tensor</span><span class="p">,</span>
        <span class="mi">6</span><span class="p">,</span>
        <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span>
      <span class="p">)</span>

    <span class="n">output_tokens</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">(</span><span class="n">out</span><span class="p">)</span>
    <span class="n">decoded_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">out</span><span class="p">,</span> <span class="s2">"code-davinci-002"</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">(</span><span class="n">encoded_tensor</span><span class="p">)</span> <span class="o">==</span> <span class="p">[</span><span class="mi">6109</span><span class="p">,</span> <span class="mi">3626</span><span class="p">,</span> <span class="mi">6100</span><span class="p">,</span> <span class="mi">345</span><span class="p">]</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">encoded_tensor</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">4</span><span class="p">}</span>
    <span class="n">assert</span> <span class="n">output_tokens</span> <span class="o">==</span> <span class="p">[</span><span class="mi">6109</span><span class="p">,</span> <span class="mi">3626</span><span class="p">,</span> <span class="mi">6100</span><span class="p">,</span> <span class="mi">345</span><span class="p">,</span> <span class="mi">16244</span><span class="p">,</span> <span class="mi">42570</span><span class="p">,</span> <span class="mi">46784</span><span class="p">,</span> <span class="mi">9329</span><span class="p">,</span> <span class="mi">30840</span><span class="p">,</span> <span class="mi">27005</span><span class="p">]</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">out</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">10</span><span class="p">}</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">output_tokens</span><span class="p">)</span> <span class="o">==</span> <span class="mi">10</span>
    <span class="n">assert</span> <span class="n">decoded_text</span> <span class="o">==</span> <span class="s2">"Every effort moves youdevelopclimategra Lind Healthy Dread"</span>
  <span class="k">end</span>
</code></pre></div></div>

<h3 id="calculating-the-text-generation-loss">Calculating the text generation loss</h3>

<p>The text generation loss is used for assessing the text quality generated during training, and it is used in the backpropagation algorithm during training. Here is an example of correct text generation with a vocabulary of seven tokens:</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
Vocabulary                    Probability rows from softmax                          Prediction
----------                    -----------------------------                          ----------
"a"       -&gt; 0
"effort"  -&gt; 1                index:      0      1      2      3      4      5      6
"every"   -&gt; 2                every  [2]  0.10  [0.60]  0.20   0.05   0.00   0.02   0.01  -&gt; argmax 1 -&gt; effort
"forward" -&gt; 3                effort [1]  0.06   0.07   0.01   0.26  [0.35]  0.13   0.12  -&gt; argmax 4 -&gt; moves
"moves"   -&gt; 4                moves  [4]  0.01   0.10   0.10   0.20   0.12  [0.34]  0.13  -&gt; argmax 5 -&gt; you
"you"     -&gt; 5
"zoo"     -&gt; 6

Inverse vocabulary:
  0 -&gt; "a"
  1 -&gt; "effort"
  2 -&gt; "every"
  3 -&gt; "forward"
  4 -&gt; "moves"
  5 -&gt; "you"
  6 -&gt; "zoo"

Flow:
  input text           token IDs      predicted IDs      output text
  ----------           ---------      -------------      -----------
  every effort moves   [2, 1, 4]  -&gt;  [1, 4, 5]      -&gt;  effort moves you
</code></pre></div></div>

<p>So when the last token is “moves”, the max value for the “moves” logits row is [0.34], which corresponds to the 5th vocabulary index, and this is the token “you”. Our goal is to train the model so the values in [] are exactly at those indexes, because these are the valid next tokens. How do we do that? We will use the following batches as training data: “every effort moves”, “effort moves you”.</p>

<p>This is a helper module that calculates the loss:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">LossUtils</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Loss helpers for GPT-style logits and target token tensors.
  """</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns the predicted probabilities for the expected target token ids.

  `probas` should be shaped `{batch_size, seq_len, vocab_size}` and `targets`
  should be shaped `{batch_size, seq_len}`. The returned tensor is flattened to
  `{batch_size * seq_len}`.

  ## Examples

      iex&gt; probas = Nx.tensor([[[0.5, 0.25, 0.25], [0.125, 0.375, 0.5]]])
      iex&gt; targets = Nx.tensor([[0, 2]])
      iex&gt; LlmScratch.LossUtils.target_token_probas(probas, targets) |&gt; Nx.to_flat_list()
      [0.5, 0.5]
  """</span>
  <span class="nv">@spec</span> <span class="n">target_token_probas</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">target_token_probas</span><span class="p">(%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">probas</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">targets</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">vocab_size</span> <span class="o">=</span> <span class="n">elem</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">probas</span><span class="p">),</span> <span class="mi">2</span><span class="p">)</span>

    <span class="c1">#{2, 3, 50257} =&gt; {6, 50257}</span>
    <span class="n">probas</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">reshape</span><span class="p">({</span><span class="ss">:auto</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">})</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">take_along_axis</span><span class="p">(</span>
      <span class="c1">#{2, 3} =&gt; {6, 1}</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="p">{</span><span class="ss">:auto</span><span class="p">,</span> <span class="mi">1</span><span class="p">}),</span> <span class="ss">axis:</span> <span class="mi">1</span>
    <span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="ss">axes:</span> <span class="p">[</span><span class="mi">1</span><span class="p">])</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Calculates mean cross entropy loss from logits and target token ids.
  step1: logits
  step2: probabilities
  step3: target probabilities
  step4: logarithmic probabilities
  step5: average logarithmic probabilities
  step6: negative average log probabilities

  `logits` should be shaped `{batch_size, seq_len, vocab_size}` and `targets`
  should be shaped `{batch_size, seq_len}`.

  ## Examples

      iex&gt; logits = Nx.log(Nx.tensor([[[0.5, 0.25, 0.25], [0.125, 0.375, 0.5]]]))
      iex&gt; targets = Nx.tensor([[0, 2]])
      iex&gt; loss = LlmScratch.LossUtils.cross_entropy_loss(logits, targets)
      iex&gt; Float.round(Nx.to_number(loss), 6)
      0.693147
  """</span>
  <span class="nv">@spec</span> <span class="n">cross_entropy_loss</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">cross_entropy_loss</span><span class="p">(%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">logits</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">targets</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">logits</span> <span class="c1">#step1</span>
    <span class="o">|&gt;</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="c1">#step2</span>
    <span class="o">|&gt;</span> <span class="n">target_token_probas</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span> <span class="c1">#step3</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">log</span><span class="p">()</span> <span class="c1">#step4</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">negate</span><span class="p">()</span> <span class="c1">#step6</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span> <span class="c1">#step5</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>Here is a test that calculates the loss for our simple example:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"5.1.2 calculating the text generation loss"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">256</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.0</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">input_texts</span> <span class="o">=</span> <span class="p">[</span><span class="s2">"every effort moves"</span><span class="p">,</span> <span class="s2">"I really like"</span><span class="p">]</span>
    <span class="n">target_texts</span> <span class="o">=</span> <span class="p">[</span><span class="s2">" effort moves you"</span><span class="p">,</span> <span class="s2">" really like chocolate"</span><span class="p">]</span>

    <span class="n">inputs</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">texts_to_token_ids</span><span class="p">(</span><span class="n">input_texts</span><span class="p">,</span> <span class="s2">"code-davinci-002"</span><span class="p">)</span>
    <span class="n">targets</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">texts_to_token_ids</span><span class="p">(</span><span class="n">target_texts</span><span class="p">,</span> <span class="s2">"code-davinci-002"</span><span class="p">)</span>

    <span class="c1"># initialize our gpt model</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>
    <span class="c1"># predict next input token</span>
    <span class="n">logits</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">inputs</span><span class="p">)</span>
    <span class="c1"># use softmax to get probability values (this is a redundant step)</span>
    <span class="n">probas</span> <span class="o">=</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
    <span class="c1"># get token ids that have maximum logit values</span>
    <span class="n">token_ids</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="ss">keep_axis:</span> <span class="no">true</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_list</span><span class="p">(</span><span class="n">inputs</span><span class="p">)</span> <span class="o">==</span> <span class="p">[</span>
             <span class="p">[</span><span class="mi">16_833</span><span class="p">,</span> <span class="mi">3626</span><span class="p">,</span> <span class="mi">6100</span><span class="p">],</span>
             <span class="p">[</span><span class="mi">40</span><span class="p">,</span> <span class="mi">1107</span><span class="p">,</span> <span class="mi">588</span><span class="p">]</span>
           <span class="p">]</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_list</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span> <span class="o">==</span> <span class="p">[</span>
             <span class="p">[</span><span class="mi">3626</span><span class="p">,</span> <span class="mi">6100</span><span class="p">,</span> <span class="mi">345</span><span class="p">],</span>
             <span class="p">[</span><span class="mi">1107</span><span class="p">,</span> <span class="mi">588</span><span class="p">,</span> <span class="mi">11_311</span><span class="p">]</span>
           <span class="p">]</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">50_257</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">}</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_list</span><span class="p">(</span><span class="n">token_ids</span><span class="p">)</span> <span class="o">==</span> <span class="p">[</span>
             <span class="p">[[</span><span class="mi">40_524</span><span class="p">],</span> <span class="p">[</span><span class="mi">8520</span><span class="p">],</span> <span class="p">[</span><span class="mi">4436</span><span class="p">]],</span>
             <span class="p">[[</span><span class="mi">22_666</span><span class="p">],</span> <span class="p">[</span><span class="mi">2829</span><span class="p">],</span> <span class="p">[</span><span class="mi">5387</span><span class="p">]]</span>
           <span class="p">]</span>

    <span class="c1"># targets the model should predict</span>
    <span class="n">assert</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">targets</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">new_axis</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span> <span class="s2">"code-davinci-002"</span><span class="p">)</span> <span class="o">==</span>
             <span class="s2">" effort moves you"</span>

    <span class="c1"># what the model actually predicted</span>
    <span class="n">assert</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span>
             <span class="n">token_ids</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">flatten</span><span class="p">()</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">new_axis</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
             <span class="s2">"code-davinci-002"</span>
           <span class="p">)</span> <span class="o">==</span>
             <span class="s2">" Clarksonerved hospital"</span>

    <span class="c1"># probabilities that the model generated for targets</span>
    <span class="n">target_probas</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">target_token_probas</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span>

    <span class="n">assert_close</span><span class="p">(</span>
      <span class="n">target_probas</span><span class="p">,</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span>
        <span class="mf">2.2307187e-5</span><span class="p">,</span>
        <span class="mf">2.2980759e-5</span><span class="p">,</span>
        <span class="mf">2.1722459e-5</span><span class="p">,</span>
        <span class="mf">2.1338150e-5</span><span class="p">,</span>
        <span class="mf">1.7680910e-5</span><span class="p">,</span>
        <span class="mf">1.4839508e-5</span>
      <span class="p">]),</span>
      <span class="ss">atol:</span> <span class="mf">1.0e-10</span>
    <span class="p">)</span>

    <span class="c1"># we calculate loss as the difference between what the model actually predicted</span>
    <span class="c1"># and the target tokens</span>
    <span class="c1"># this loss should be minimized</span>
    <span class="c1"># we have a big loss result, which is expected because we did not train the model</span>
    <span class="c1"># the goal is to have loss close to zero</span>
    <span class="n">loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">cross_entropy_loss</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">loss</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">10.824145</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-6</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<h3 id="calculating-the-training-and-validation-set-losses">Calculating the training and validation set losses</h3>

<p>We are training our GPTModel using The Verdict story, which we used in the <a href="/build-llm-from-scratch-chapter-2">post about text tokenization</a>. We are splitting that story into two parts: training and validation. It is a small training dataset that we can run on our laptops. For example, the Llama 2 model with just 7 billion parameters required 184,320 GPU hours on A100 GPU hardware!</p>

<p>We first extended our <code class="language-plaintext highlighter-rouge">LossUtil</code> module with two loss calculation functions, one for a single batch, and one for calculating loss over all batches:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@doc</span> <span class="sd">"""
  Calculates cross entropy loss for an input/target batch and model.

  This is the Nx counterpart to:

      logits = model(input_batch)
      loss = torch.nn.functional.cross_entropy(
        logits.flatten(0, 1),
        target_batch.flatten()
      )

  `input_batch` and `target_batch` should be shaped `{batch_size, seq_len}`.
  The model must be a struct whose module exports `forward/2` or `forward/3`.

  The optional `device` argument accepts an Nx backend, such as `EXLA.Backend`
  or `{EXLA.Backend, client: :cuda}`. Use `nil` or `:default` to leave tensors
  on their current backend.
  """</span>
  <span class="nv">@spec</span> <span class="n">calc_loss_batch</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">struct</span><span class="p">(),</span> <span class="no">nil</span> <span class="o">|</span> <span class="ss">:default</span> <span class="o">|</span> <span class="n">atom</span><span class="p">()</span> <span class="o">|</span> <span class="n">tuple</span><span class="p">())</span> <span class="p">::</span>
          <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">calc_loss_batch</span><span class="p">(</span>
        <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">input_batch</span><span class="p">,</span>
        <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">target_batch</span><span class="p">,</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">device</span> <span class="p">\\</span> <span class="ss">:default</span>
      <span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">input_batch</span> <span class="o">=</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>
    <span class="n">target_batch</span> <span class="o">=</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">target_batch</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

    <span class="n">model</span>
    <span class="o">|&gt;</span> <span class="n">forward_model</span><span class="p">(</span><span class="n">input_batch</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">cross_entropy_loss</span><span class="p">(</span><span class="n">target_batch</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Calculates average loss over batches from a data loader.

  `data_loader` should be a `%{stream: stream, length: length}` map, such as the
  value returned by `LlmScratch.DataLoader.new/2`. Each batch may be either:

    * `{input_batch, target_batch}` with already-stacked tensors
    * a list of `{input_tensor, target_tensor}` examples

  Pass `nil` for `num_batches` to evaluate one full pass through the loader.
  Returns `:nan` when the loader has no batches.
  """</span>
  <span class="nv">@spec</span> <span class="n">calc_loss_loader</span><span class="p">(</span>
          <span class="n">map</span><span class="p">(),</span>
          <span class="n">struct</span><span class="p">(),</span>
          <span class="no">nil</span> <span class="o">|</span> <span class="ss">:default</span> <span class="o">|</span> <span class="n">atom</span><span class="p">()</span> <span class="o">|</span> <span class="n">tuple</span><span class="p">(),</span>
          <span class="no">nil</span> <span class="o">|</span> <span class="n">non_neg_integer</span><span class="p">()</span>
        <span class="p">)</span> <span class="p">::</span>
          <span class="n">float</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">data_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span> <span class="p">\\</span> <span class="ss">:default</span><span class="p">,</span> <span class="n">num_batches</span> <span class="p">\\</span> <span class="no">nil</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">data_loader</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">loader_length</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">data_loader</span><span class="p">,</span> <span class="ss">:length</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>

    <span class="k">cond</span> <span class="k">do</span>
      <span class="n">loader_length</span> <span class="o">==</span> <span class="mi">0</span> <span class="o">-&gt;</span>
        <span class="ss">:nan</span>

      <span class="no">true</span> <span class="o">-&gt;</span>
        <span class="n">num_batches</span> <span class="o">=</span> <span class="n">normalize_num_batches</span><span class="p">(</span><span class="n">num_batches</span><span class="p">,</span> <span class="n">loader_length</span><span class="p">)</span>

        <span class="n">data_loader</span><span class="o">.</span><span class="n">stream</span>
        <span class="o">|&gt;</span> <span class="no">Stream</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">num_batches</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="k">fn</span> <span class="n">batch</span><span class="p">,</span> <span class="n">total_loss</span> <span class="o">-&gt;</span>
          <span class="p">{</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">}</span> <span class="o">=</span> <span class="n">stack_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>
          <span class="n">loss</span> <span class="o">=</span> <span class="n">calc_loss_batch</span><span class="p">(</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

          <span class="n">total_loss</span> <span class="o">+</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">(</span><span class="n">loss</span><span class="p">)</span>
        <span class="k">end</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">./</span><span class="p">(</span><span class="n">num_batches</span><span class="p">)</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span> <span class="ow">when</span> <span class="n">device</span> <span class="ow">in</span> <span class="p">[</span><span class="no">nil</span><span class="p">,</span> <span class="ss">:default</span><span class="p">],</span> <span class="k">do</span><span class="p">:</span> <span class="n">tensor</span>
  <span class="k">defp</span> <span class="n">maybe_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">normalize_num_batches</span><span class="p">(</span><span class="no">nil</span><span class="p">,</span> <span class="n">loader_length</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">loader_length</span>
  <span class="k">defp</span> <span class="n">normalize_num_batches</span><span class="p">(</span><span class="n">num_batches</span><span class="p">,</span> <span class="n">loader_length</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">min</span><span class="p">(</span><span class="n">num_batches</span><span class="p">,</span> <span class="n">loader_length</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">stack_batch</span><span class="p">({%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">input_batch</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">target_batch</span><span class="p">})</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">stack_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">unzip</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

    <span class="p">{</span><span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">inputs</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">targets</span><span class="p">)}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">forward_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">module</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">__struct__</span>

    <span class="k">cond</span> <span class="k">do</span>
      <span class="n">function_exported?</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="ss">:forward</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span> <span class="o">-&gt;</span>
        <span class="n">apply</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="ss">:forward</span><span class="p">,</span> <span class="p">[</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">])</span>

      <span class="n">function_exported?</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="ss">:forward</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span> <span class="o">-&gt;</span>
        <span class="n">apply</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="ss">:forward</span><span class="p">,</span> <span class="p">[</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">,</span> <span class="p">[]])</span>

      <span class="no">true</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
              <span class="s2">"expected </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">module</span><span class="p">)</span><span class="si">}</span><span class="s2"> to export forward/2 or forward/3"</span>
    <span class="k">end</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Here is a test where we calculate loss over loaded batches. The loss is &gt; 0 because the model has not been trained. Note that we are using <em>The Verdict</em> story.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"5.1.3 Calculating the training and validation set losses"</span> <span class="k">do</span>
    <span class="c1"># EXLA backend to speed up GPT model initialization</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="n">device</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="c1"># read The Verdict story</span>
    <span class="n">file_content</span> <span class="o">=</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">(</span><span class="s2">"the-verdict.txt"</span><span class="p">)</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">encoded_tokens</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">,</span> <span class="n">file_content</span><span class="p">)</span>

    <span class="c1"># number of characters in The Verdict</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">file_content</span><span class="p">)</span> <span class="o">==</span> <span class="mi">20_479</span>
    <span class="c1"># number of The Verdict tokens</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">encoded_tokens</span><span class="p">)</span> <span class="o">==</span> <span class="mi">5_145</span>

    <span class="c1"># the smaller GPT-2 model</span>
    <span class="c1"># context length is even smaller, 256 tokens</span>
    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">256</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.0</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="c1"># 90% of The Verdict is training data, 10% is validation</span>
    <span class="c1"># we need validation to see how accurate our loss calculation is</span>
    <span class="n">train_ratio</span> <span class="o">=</span> <span class="mf">0.90</span>
    <span class="n">split_idx</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">train_ratio</span> <span class="o">*</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">file_content</span><span class="p">))</span>
    <span class="n">train_data</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">file_content</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">split_idx</span><span class="p">)</span>
    <span class="n">val_data</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">file_content</span><span class="p">,</span> <span class="n">split_idx</span><span class="p">,</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">file_content</span><span class="p">)</span> <span class="o">-</span> <span class="n">split_idx</span><span class="p">)</span>

    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">train_tokens</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">,</span> <span class="n">train_data</span><span class="p">,</span> <span class="p">[</span><span class="s2">"&lt;|endoftext|&gt;"</span><span class="p">])</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">val_tokens</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">,</span> <span class="n">val_data</span><span class="p">,</span> <span class="p">[</span><span class="s2">"&lt;|endoftext|&gt;"</span><span class="p">])</span>

    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">train_data</span><span class="p">)</span> <span class="o">==</span> <span class="mi">18_431</span>
    <span class="n">assert</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">val_data</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2_048</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_tokens</span><span class="p">)</span> <span class="o">==</span> <span class="mi">4_612</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">val_tokens</span><span class="p">)</span> <span class="o">==</span> <span class="mi">534</span>

    <span class="c1"># load training data into batches of two</span>
    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="no">GptDatasetV1</span><span class="o">.</span><span class="n">create_dataloader_v1</span><span class="p">(</span>
        <span class="ss">raw_text:</span> <span class="n">train_data</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span>
        <span class="ss">max_length:</span> <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">stride:</span> <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="mi">0</span>
      <span class="p">)</span>

    <span class="c1"># load validation data in batches of two</span>
    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="no">GptDatasetV1</span><span class="o">.</span><span class="n">create_dataloader_v1</span><span class="p">(</span>
        <span class="ss">raw_text:</span> <span class="n">val_data</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span>
        <span class="ss">max_length:</span> <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">stride:</span> <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="mi">0</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="p">%{</span><span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span> <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">,</span> <span class="ss">num_workers:</span> <span class="mi">0</span><span class="p">}</span> <span class="o">=</span> <span class="n">train_loader</span>
    <span class="n">assert</span> <span class="p">%{</span><span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span> <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span> <span class="ss">num_workers:</span> <span class="mi">0</span><span class="p">}</span> <span class="o">=</span> <span class="n">val_loader</span>

    <span class="c1"># we have 9 training batches and 1 validation batch</span>
    <span class="n">train_batches</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">train_loader</span><span class="o">.</span><span class="n">stream</span><span class="p">,</span> <span class="mi">9</span><span class="p">)</span>
    <span class="n">val_batches</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">val_loader</span><span class="o">.</span><span class="n">stream</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_batches</span><span class="p">)</span> <span class="o">==</span> <span class="mi">9</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">val_batches</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">each</span><span class="p">(</span><span class="n">train_batches</span><span class="p">,</span> <span class="k">fn</span> <span class="n">batch</span> <span class="o">-&gt;</span>
      <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">unzip</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

      <span class="c1"># each training batch has two token sequences with the context length dimension</span>
      <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">inputs</span><span class="p">))</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">256</span><span class="p">}</span>
      <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">targets</span><span class="p">))</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">256</span><span class="p">}</span>
    <span class="k">end</span><span class="p">)</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">each</span><span class="p">(</span><span class="n">val_batches</span><span class="p">,</span> <span class="k">fn</span> <span class="n">batch</span> <span class="o">-&gt;</span>
      <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">unzip</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

      <span class="c1"># each validation batch has two token sequences with the context length dimension</span>
      <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">inputs</span><span class="p">))</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">256</span><span class="p">}</span>
      <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">targets</span><span class="p">))</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">256</span><span class="p">}</span>
    <span class="k">end</span><span class="p">)</span>

    <span class="c1"># our GPTModel initialization</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="c1"># calculate loss for loaded training and validation data</span>
    <span class="n">train_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>
    <span class="n">val_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">val_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

    <span class="n">assert_in_delta</span> <span class="n">train_loss</span><span class="p">,</span> <span class="mf">10.847595</span><span class="p">,</span> <span class="mf">1.0e-5</span>
    <span class="n">assert_in_delta</span> <span class="n">val_loss</span><span class="p">,</span> <span class="mf">10.849722</span><span class="p">,</span> <span class="mf">1.0e-5</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Equipped with loss functions, we are ready to train our model!</p>

<h2 id="52-training-an-llm">5.2 Training an LLM</h2>

<p>We are implementing the following simple training loop. There are a number of more advanced techniques, but for learning purposes, simple training (which is actually not so simple) is enough for a start.</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>Training loop

  +---------------------------------------+        One epoch is one complete
  | 1. Iterate over training epochs       | &lt;----  pass over the training set.
  +---------------------------------------+
             \
              v
        +---------------------------------------+  Number of batches =
        | 2. Iterate over batches in each epoch |  training set size /
        +---------------------------------------+  batch size.
                     \
                      v
              +---------------------------------------+
              | 3. Reset loss gradients from the      |
              |    previous batch iteration           |
              +---------------------------------------+
                              |
                              v
              +---------------------------------------+
              | 4. Calculate loss on current batch    |
              +---------------------------------------+      These are the usual
                              |                              training steps for
                              v                              neural networks.
              +---------------------------------------+
              | 5. Backward pass to calculate loss    |
              |    gradients                          |
              +---------------------------------------+
                              |
                              v
              +---------------------------------------+
              | 6. Update model weights using loss    |
              |    gradients                          |
              +---------------------------------------+
                              |
                              v
              +---------------------------------------+
              | 7. Print training and validation      |      Optional steps for
              |    losses                             | &lt;--- tracking progress.
              +---------------------------------------+
                       /
                      v
        +---------------------------------------+
        | 8. Generate sample text for visual    | &lt;--- Optional progress check.
        |    inspection                         |
        +---------------------------------------+
             |
             v
  repeat with the next epoch
</code></pre></div></div>

<p>The training loop in Elixir is similar to the Python version, but with one important
difference: Elixir data is immutable, so the optimizer must return the updated model
after each training step.</p>

<p>But before we implement this simple training algorithm, we first need to implement the Nx.Container protocol for the structs that will be used in the training algorithm. That allows Nx to operate on Nx.Tensors that are part of those structs. Here is an example for <code class="language-plaintext highlighter-rouge">LlmScratch.EmbeddingNative</code>; the principle is the same for <code class="language-plaintext highlighter-rouge">LlmScratch.DummyLayerNorm</code>, <code class="language-plaintext highlighter-rouge">LlmScratch.GELU</code>, <code class="language-plaintext highlighter-rouge">LlmScratch.FeedForward</code>, <code class="language-plaintext highlighter-rouge">LlmScratch.MultiheadAttention</code>, <code class="language-plaintext highlighter-rouge">LlmScratch.TransformerBlock</code>, and <code class="language-plaintext highlighter-rouge">LlmScratch.GPTModel</code>.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defimpl</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Container</span><span class="p">,</span> <span class="ss">for:</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">EmbeddingNative</span> <span class="k">do</span>
  <span class="k">def</span> <span class="n">traverse</span><span class="p">(</span><span class="n">embedding</span><span class="p">,</span> <span class="n">acc</span><span class="p">,</span> <span class="n">fun</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">weight</span><span class="p">,</span> <span class="n">acc</span><span class="p">}</span> <span class="o">=</span> <span class="n">fun</span><span class="o">.</span><span class="p">(</span><span class="n">embedding</span><span class="o">.</span><span class="n">weight</span><span class="p">,</span> <span class="n">acc</span><span class="p">)</span>
    <span class="p">{%{</span><span class="n">embedding</span> <span class="o">|</span> <span class="ss">weight:</span> <span class="n">weight</span><span class="p">},</span> <span class="n">acc</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">def</span> <span class="n">reduce</span><span class="p">(</span><span class="n">embedding</span><span class="p">,</span> <span class="n">acc</span><span class="p">,</span> <span class="n">fun</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">fun</span><span class="o">.</span><span class="p">(</span><span class="n">embedding</span><span class="o">.</span><span class="n">weight</span><span class="p">,</span> <span class="n">acc</span><span class="p">)</span>

  <span class="k">def</span> <span class="n">serialize</span><span class="p">(</span><span class="n">embedding</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">metadata</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">embedding</span><span class="p">,</span> <span class="p">[</span><span class="ss">:vocab_size</span><span class="p">,</span> <span class="ss">:embedding_dim</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">])</span>
    <span class="p">{</span><span class="bp">__MODULE__</span><span class="p">,</span> <span class="p">[</span><span class="ss">weight:</span> <span class="n">embedding</span><span class="o">.</span><span class="n">weight</span><span class="p">],</span> <span class="n">metadata</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">def</span> <span class="n">deserialize</span><span class="p">([</span><span class="ss">weight:</span> <span class="n">weight</span><span class="p">],</span> <span class="n">metadata</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">struct!</span><span class="p">(</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">EmbeddingNative</span><span class="p">,</span> <span class="no">Map</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="n">metadata</span><span class="p">,</span> <span class="ss">:weight</span><span class="p">,</span> <span class="n">weight</span><span class="p">))</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>Another change in existing modules is that each forward function uses the <code class="language-plaintext highlighter-rouge">Nx.defn</code> macro, so <code class="language-plaintext highlighter-rouge">Nx</code> is able to track training gradients.</p>

<p>Here is the not-so-simple Training module. I added code comments so you can track the training algorithm image:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">Training</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Simple GPT training loop helpers mirroring the book's `train_model_simple`.
  """</span>

  <span class="kn">import</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span><span class="no">GPTModel</span><span class="p">,</span> <span class="no">LossUtils</span><span class="p">,</span> <span class="no">TextGeneration</span><span class="p">,</span> <span class="no">TextUtils</span><span class="p">}</span>

  <span class="k">defmodule</span> <span class="no">AdamW</span> <span class="k">do</span>
    <span class="nv">@moduledoc</span> <span class="sd">"""
    Minimal AdamW optimizer state.

    In the training loop, AdamW is responsible for turning the gradients from
    `loss_and_grad/4` into a new model. Each batch computes how the loss changes
    with respect to the model's trainable tensors; AdamW decides how large each
    parameter update should be, applies regularization through weight decay, and
    carries optimizer state forward to the next batch.

    AdamW keeps two moving averages for each trainable tensor:

      * `m` tracks the exponentially decayed average gradient.
      * `v` tracks the exponentially decayed average squared gradient.

    During each optimization step, the averages are bias-corrected with the
    current `step`, normalized with `eps`, combined with decoupled weight
    decay, and subtracted from the parameter using `learning_rate`.

    ## Fields

      * `:learning_rate` - positive scalar that controls the size of each
        parameter update.
      * `:weight_decay` - decoupled decay factor applied directly to
        parameters, separate from the gradient moments.
      * `:beta1` - exponential decay rate for `m`, the first-moment gradient
        estimate.
      * `:beta2` - exponential decay rate for `v`, the second-moment squared
        gradient estimate.
      * `:eps` - small constant added to the denominator for numerical
        stability.
      * `:step` - number of optimizer updates already applied; used for bias
        correction.
      * `:m` - first-moment tensors, lazily initialized to zeros with the same
        shapes as the trainable parameters.
      * `:v` - second-moment tensors, lazily initialized to zeros with the same
        shapes as the trainable parameters.
    """</span>

    <span class="k">defstruct</span> <span class="ss">learning_rate:</span> <span class="no">nil</span><span class="p">,</span>
              <span class="ss">weight_decay:</span> <span class="mf">0.0</span><span class="p">,</span>
              <span class="ss">beta1:</span> <span class="mf">0.9</span><span class="p">,</span>
              <span class="ss">beta2:</span> <span class="mf">0.999</span><span class="p">,</span>
              <span class="ss">eps:</span> <span class="mf">1.0e-8</span><span class="p">,</span>
              <span class="ss">step:</span> <span class="mi">0</span><span class="p">,</span>
              <span class="ss">m:</span> <span class="no">nil</span><span class="p">,</span>
              <span class="ss">v:</span> <span class="no">nil</span>
  <span class="k">end</span>

  <span class="nv">@type</span> <span class="n">optimizer</span> <span class="p">::</span> <span class="p">%</span><span class="no">AdamW</span><span class="p">{}</span> <span class="o">|</span> <span class="p">(</span><span class="n">struct</span><span class="p">(),</span> <span class="n">struct</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="n">struct</span><span class="p">())</span>

  <span class="nv">@spec</span> <span class="n">adamw</span><span class="p">(</span><span class="n">number</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="no">AdamW</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Creates a minimal AdamW optimizer.

  This mirrors the optimizer used by the Python chapter example. Moment state is
  initialized lazily on the first optimization step because it must match the
  model's trainable parameter structure.

  AdamW adapts each parameter update using the history of recent gradients:

      m = beta1 * m + (1 - beta1) * gradient
      v = beta2 * v + (1 - beta2) * gradient ** 2
      m_hat = m / (1 - beta1 ** step)
      v_hat = v / (1 - beta2 ** step)
      update = m_hat / (sqrt(v_hat) + eps) + weight_decay * parameter
      parameter = parameter - learning_rate * update

  The `m` and `v` tensors start as zeros with the same shapes as the model's
  trainable tensors. `weight_decay` is decoupled from the gradient moments,
  which is the distinguishing behavior of AdamW compared to classic Adam with
  L2 regularization folded into the gradient.

  ## Options

    * `:weight_decay` - decoupled weight decay. Defaults to `0.0`.
    * `:beta1` - first-moment decay. Defaults to `0.9`.
    * `:beta2` - second-moment decay. Defaults to `0.999`.
    * `:eps` - numerical stability constant. Defaults to `1.0e-8`.
  """</span>
  <span class="k">def</span> <span class="n">adamw</span><span class="p">(</span><span class="n">learning_rate</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_number</span><span class="p">(</span><span class="n">learning_rate</span><span class="p">)</span> <span class="ow">and</span> <span class="n">learning_rate</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="p">%</span><span class="no">AdamW</span><span class="p">{</span>
      <span class="ss">learning_rate:</span> <span class="n">learning_rate</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span>
      <span class="ss">weight_decay:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:weight_decay</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">)</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span>
      <span class="ss">beta1:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:beta1</span><span class="p">,</span> <span class="mf">0.9</span><span class="p">)</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span>
      <span class="ss">beta2:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:beta2</span><span class="p">,</span> <span class="mf">0.999</span><span class="p">)</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span>
      <span class="ss">eps:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:eps</span><span class="p">,</span> <span class="mf">1.0e-8</span><span class="p">)</span> <span class="o">*</span> <span class="mf">1.0</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">train_model_simple</span><span class="p">(</span>
          <span class="n">struct</span><span class="p">(),</span>
          <span class="n">map</span><span class="p">(),</span>
          <span class="n">map</span><span class="p">(),</span>
          <span class="n">optimizer</span><span class="p">(),</span>
          <span class="no">nil</span> <span class="o">|</span> <span class="ss">:default</span> <span class="o">|</span> <span class="n">atom</span><span class="p">()</span> <span class="o">|</span> <span class="n">tuple</span><span class="p">(),</span>
          <span class="n">non_neg_integer</span><span class="p">(),</span>
          <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
        <span class="p">)</span> <span class="p">::</span> <span class="p">{</span><span class="n">struct</span><span class="p">(),</span> <span class="p">[</span><span class="n">float</span><span class="p">()],</span> <span class="p">[</span><span class="n">float</span><span class="p">()],</span> <span class="p">[</span><span class="n">non_neg_integer</span><span class="p">()]}</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Trains a GPT-style model with a simple batch loop.

  This mirrors the book's Python `train_model_simple` example. Since Elixir
  data is immutable, the updated model is returned instead of being mutated in
  place. The training path calls `GPTModel.train/3`, which enables dropout and
  threads an explicit `Nx.Random` key through the model.

  ## Parameters

    * `model` - GPT-style model struct whose trainable tensors can be traversed
      by `Nx.Container`.
    * `train_loader` - data loader map with `:stream` and `:length`.
    * `val_loader` - validation data loader map with `:stream` and `:length`.
    * `optimizer` - either `Training.adamw/2` output or a two-argument function
      `(model, gradients -&gt; updated_model)`. The built-in AdamW optimizer
      returns updated optimizer state after each batch, so its moment estimates
      are carried through the full training loop.
    * `device` - Nx backend target. Use `:default` or `nil` to keep tensors on
      their current backend.
    * `num_epochs` - number of full passes over `train_loader`.
    * `eval_freq` - evaluate every `eval_freq` global training steps.
    * `eval_iter` - maximum number of batches from each loader during
      evaluation.
    * `start_context` - prompt used by `generate_and_print_sample/4`.
    * `tokenizer` - tokenizer name passed to `Tiktoken`.

  ## Returns

    * `{model, train_losses, val_losses, track_tokens_seen}`.
  """</span>
  <span class="k">def</span> <span class="n">train_model_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">train_loader</span><span class="p">,</span>
        <span class="n">val_loader</span><span class="p">,</span>
        <span class="n">optimizer</span><span class="p">,</span>
        <span class="n">device</span><span class="p">,</span>
        <span class="n">num_epochs</span><span class="p">,</span>
        <span class="n">eval_freq</span><span class="p">,</span>
        <span class="n">eval_iter</span><span class="p">,</span>
        <span class="n">start_context</span><span class="p">,</span>
        <span class="n">tokenizer</span>
      <span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_map</span><span class="p">(</span><span class="n">train_loader</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_map</span><span class="p">(</span><span class="n">val_loader</span><span class="p">)</span> <span class="ow">and</span>
             <span class="n">is_integer</span><span class="p">(</span><span class="n">num_epochs</span><span class="p">)</span> <span class="ow">and</span> <span class="n">num_epochs</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">eval_freq</span><span class="p">)</span> <span class="ow">and</span>
             <span class="n">eval_freq</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">eval_iter</span><span class="p">)</span> <span class="ow">and</span> <span class="n">eval_iter</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="c1"># Start dropout RNG from a fresh positive runtime integer, then move the</span>
    <span class="c1"># key to the same backend requested for the model and batches.</span>
    <span class="n">key</span> <span class="o">=</span>
      <span class="no">System</span><span class="o">.</span><span class="n">unique_integer</span><span class="p">([</span><span class="ss">:positive</span><span class="p">])</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">key</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="c1"># Training state is carried through the epoch/batch reducers because Elixir</span>
    <span class="c1"># data is immutable. It keeps the current model parameters, AdamW optimizer</span>
    <span class="c1"># memory, dropout RNG key, progress counters, and evaluation metrics that</span>
    <span class="c1"># are accumulated while learning.</span>
    <span class="n">state</span> <span class="o">=</span> <span class="p">%{</span>
      <span class="c1"># Current model parameters; replaced after each optimizer step.</span>
      <span class="ss">model:</span> <span class="n">maybe_transfer_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">),</span>
      <span class="c1"># Periodic evaluation losses, stored newest-first during training.</span>
      <span class="ss">train_losses:</span> <span class="p">[],</span>
      <span class="ss">val_losses:</span> <span class="p">[],</span>
      <span class="c1"># Token counts captured at the same evaluation points as the losses.</span>
      <span class="ss">track_tokens_seen:</span> <span class="p">[],</span>
      <span class="c1"># Running count of all tokens consumed by training batches.</span>
      <span class="ss">tokens_seen:</span> <span class="mi">0</span><span class="p">,</span>
      <span class="c1"># Batch counter; starts at -1 so the first batch becomes step 0.</span>
      <span class="ss">global_step:</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span>
      <span class="c1"># Dropout RNG key; GPTModel.train/3 returns the next key for the next batch.</span>
      <span class="ss">key:</span> <span class="n">key</span><span class="p">,</span>
      <span class="c1"># Optimizer state; AdamW carries step count and moment tensors across batches.</span>
      <span class="ss">optimizer:</span> <span class="n">optimizer</span>
    <span class="p">}</span>
    <span class="c1"># Main training loop by number of epochs</span>
    <span class="n">state</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">1</span><span class="o">..</span><span class="n">num_epochs</span><span class="o">//</span><span class="mi">1</span><span class="p">,</span> <span class="n">state</span><span class="p">,</span> <span class="k">fn</span> <span class="n">epoch</span><span class="p">,</span> <span class="n">state</span> <span class="o">-&gt;</span>
        <span class="n">train_epoch</span><span class="p">(</span>
          <span class="n">state</span><span class="p">,</span>
          <span class="n">train_loader</span><span class="p">,</span>
          <span class="n">val_loader</span><span class="p">,</span>
          <span class="n">device</span><span class="p">,</span>
          <span class="n">epoch</span><span class="p">,</span>
          <span class="n">eval_freq</span><span class="p">,</span>
          <span class="n">eval_iter</span><span class="p">,</span>
          <span class="n">start_context</span><span class="p">,</span>
          <span class="n">tokenizer</span>
        <span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="n">state</span><span class="o">.</span><span class="n">model</span><span class="p">,</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">(</span><span class="n">state</span><span class="o">.</span><span class="n">train_losses</span><span class="p">),</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">(</span><span class="n">state</span><span class="o">.</span><span class="n">val_losses</span><span class="p">),</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">(</span><span class="n">state</span><span class="o">.</span><span class="n">track_tokens_seen</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Calculates training and validation losses for a fixed number of batches.

  Evaluation uses the model's dropout-free forward path, matching PyTorch
  `model.eval()`.

  ## Parameters

    * `model` - GPT-style model used for inference.
    * `train_loader` - training data loader map with `:stream` and `:length`.
    * `val_loader` - validation data loader map with `:stream` and `:length`.
    * `device` - Nx backend target passed through to
      `LossUtils.calc_loss_loader/4`.
    * `eval_iter` - maximum number of batches evaluated from each loader.

  ## Returns

    * `{train_loss, val_loss}` as floats.
  """</span>
  <span class="k">def</span> <span class="n">evaluate_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">train_loader</span><span class="p">,</span> <span class="n">val_loader</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">eval_iter</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">eval_iter</span><span class="p">)</span> <span class="ow">and</span> <span class="n">eval_iter</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="p">{</span>
      <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">eval_iter</span><span class="p">),</span>
      <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">val_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">eval_iter</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Generates and prints a text sample from `start_context`.

  The prompt is encoded with `tokenizer`, moved to `device` when requested, and
  decoded back to text after greedy generation. The generated text is printed
  and also returned.

  ## Parameters

    * `model` - GPT-style model used for greedy text generation.
    * `tokenizer` - tokenizer name passed to `Tiktoken`.
    * `device` - Nx backend target for the encoded prompt.
    * `start_context` - prompt text to continue.

  ## Returns

    * generated text as a string.
  """</span>
  <span class="k">def</span> <span class="n">generate_and_print_sample</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">start_context</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">start_context</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">idx</span> <span class="o">=</span>
      <span class="n">start_context</span>
      <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="n">context_size</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span>

    <span class="n">generated</span> <span class="o">=</span>
      <span class="n">model</span>
      <span class="o">|&gt;</span> <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate_text_simple</span><span class="p">(</span><span class="n">idx</span><span class="p">,</span> <span class="mi">50</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">generated</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>
    <span class="no">IO</span><span class="o">.</span><span class="n">puts</span><span class="p">(</span><span class="n">text</span><span class="p">)</span>
    <span class="n">text</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Computes training loss and gradients for one input/target batch.

  This is the `loss.backward()` counterpart from the Python example. It calls
  `GPTModel.train/3`, so dropout is enabled, and returns the next RNG key for
  the following batch.

  ## Parameters

    * `model` - GPT model whose trainable tensors implement `Nx.Container`.
    * `input_batch` - token id tensor shaped `{batch_size, seq_len}`.
    * `target_batch` - target token id tensor shaped `{batch_size, seq_len}`.
    * `key` - `Nx.Random` key used for dropout.

  ## Returns

    * `{loss, gradients, key}` where `loss` is a scalar tensor, `gradients`
      matches the model structure, and `key` is the advanced RNG key.
  """</span>
  <span class="n">defn</span> <span class="n">loss_and_grad</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{{</span><span class="n">loss</span><span class="p">,</span> <span class="n">key</span><span class="p">},</span> <span class="p">{</span><span class="n">gradients</span><span class="p">,</span> <span class="n">_input_gradients</span><span class="p">,</span> <span class="n">_target_gradients</span><span class="p">,</span> <span class="n">_key_gradients</span><span class="p">}}</span> <span class="o">=</span>
      <span class="n">value_and_grad</span><span class="p">(</span>
        <span class="p">{</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">,</span> <span class="n">key</span><span class="p">},</span>
        <span class="k">fn</span> <span class="p">{</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">,</span> <span class="n">key</span><span class="p">}</span> <span class="o">-&gt;</span>
          <span class="p">{</span><span class="n">logits</span><span class="p">,</span> <span class="n">key</span><span class="p">}</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">train</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span>
          <span class="p">{</span><span class="no">LossUtils</span><span class="o">.</span><span class="n">cross_entropy_loss_defn</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">),</span> <span class="n">key</span><span class="p">}</span>
        <span class="k">end</span><span class="p">,</span>
        <span class="o">&amp;</span><span class="n">elem</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
      <span class="p">)</span>

    <span class="p">{</span><span class="n">loss</span><span class="p">,</span> <span class="n">gradients</span><span class="p">,</span> <span class="n">key</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">train_epoch</span><span class="p">(</span>
         <span class="n">state</span><span class="p">,</span>
         <span class="n">train_loader</span><span class="p">,</span>
         <span class="n">val_loader</span><span class="p">,</span>
         <span class="n">device</span><span class="p">,</span>
         <span class="n">epoch</span><span class="p">,</span>
         <span class="n">eval_freq</span><span class="p">,</span>
         <span class="n">eval_iter</span><span class="p">,</span>
         <span class="n">start_context</span><span class="p">,</span>
         <span class="n">tokenizer</span>
       <span class="p">)</span> <span class="k">do</span>
    <span class="n">train_loader</span>
    <span class="c1">#shuffle batches</span>
    <span class="o">|&gt;</span> <span class="n">epoch_batches</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">state</span><span class="p">,</span> <span class="k">fn</span> <span class="n">batch</span><span class="p">,</span> <span class="n">state</span> <span class="o">-&gt;</span>
      <span class="c1">#separate batch inputs and targets</span>
      <span class="p">{</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">}</span> <span class="o">=</span> <span class="n">stack_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>
      <span class="c1">#calculate loss and gradients</span>
      <span class="p">{</span><span class="n">_loss</span><span class="p">,</span> <span class="n">gradients</span><span class="p">,</span> <span class="n">key</span><span class="p">}</span> <span class="o">=</span> <span class="n">loss_and_grad</span><span class="p">(</span><span class="n">state</span><span class="o">.</span><span class="n">model</span><span class="p">,</span> <span class="n">input_batch</span><span class="p">,</span> <span class="n">target_batch</span><span class="p">,</span> <span class="n">state</span><span class="o">.</span><span class="n">key</span><span class="p">)</span>
      <span class="c1">#avoid overfitting and penalized larger weights</span>
      <span class="p">{</span><span class="n">model</span><span class="p">,</span> <span class="n">optimizer</span><span class="p">}</span> <span class="o">=</span> <span class="n">optimizer_step</span><span class="p">(</span><span class="n">state</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">state</span><span class="o">.</span><span class="n">model</span><span class="p">,</span> <span class="n">gradients</span><span class="p">)</span>
      <span class="c1">#tokens that we have processed so far</span>
      <span class="n">tokens_seen</span> <span class="o">=</span> <span class="n">state</span><span class="o">.</span><span class="n">tokens_seen</span> <span class="o">+</span> <span class="no">Nx</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">input_batch</span><span class="p">)</span>
      <span class="n">global_step</span> <span class="o">=</span> <span class="n">state</span><span class="o">.</span><span class="n">global_step</span> <span class="o">+</span> <span class="mi">1</span>

      <span class="n">state</span> <span class="o">=</span> <span class="p">%{</span>
        <span class="n">state</span>
        <span class="o">|</span> <span class="ss">model:</span> <span class="n">model</span><span class="p">,</span>
          <span class="ss">tokens_seen:</span> <span class="n">tokens_seen</span><span class="p">,</span>
          <span class="ss">global_step:</span> <span class="n">global_step</span><span class="p">,</span>
          <span class="ss">key:</span> <span class="n">key</span><span class="p">,</span>
          <span class="ss">optimizer:</span> <span class="n">optimizer</span>
      <span class="p">}</span>
      <span class="c1">#evaluate model on evaluation frequency</span>
      <span class="k">if</span> <span class="n">rem</span><span class="p">(</span><span class="n">global_step</span><span class="p">,</span> <span class="n">eval_freq</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span> <span class="k">do</span>
        <span class="p">{</span><span class="n">train_loss</span><span class="p">,</span> <span class="n">val_loss</span><span class="p">}</span> <span class="o">=</span>
          <span class="n">evaluate_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">train_loader</span><span class="p">,</span> <span class="n">val_loader</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">eval_iter</span><span class="p">)</span>

        <span class="no">IO</span><span class="o">.</span><span class="n">puts</span><span class="p">(</span>
          <span class="s2">"Ep </span><span class="si">#{</span><span class="n">epoch</span><span class="si">}</span><span class="s2"> (Step </span><span class="si">#{</span><span class="n">pad_step</span><span class="p">(</span><span class="n">global_step</span><span class="p">)</span><span class="si">}</span><span class="s2">): "</span> <span class="o">&lt;&gt;</span>
            <span class="s2">"Train loss </span><span class="si">#{</span><span class="n">format_loss</span><span class="p">(</span><span class="n">train_loss</span><span class="p">)</span><span class="si">}</span><span class="s2">, Val loss </span><span class="si">#{</span><span class="n">format_loss</span><span class="p">(</span><span class="n">val_loss</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
        <span class="p">)</span>
        <span class="c1">#generate tokens based on current model to see what model actually predicts</span>
        <span class="c1">#we do not want gibberish text!</span>
        <span class="n">generate_and_print_sample</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">start_context</span><span class="p">)</span>

        <span class="p">%{</span>
          <span class="n">state</span>
          <span class="o">|</span> <span class="ss">train_losses:</span> <span class="p">[</span><span class="n">train_loss</span> <span class="o">|</span> <span class="n">state</span><span class="o">.</span><span class="n">train_losses</span><span class="p">],</span>
            <span class="ss">val_losses:</span> <span class="p">[</span><span class="n">val_loss</span> <span class="o">|</span> <span class="n">state</span><span class="o">.</span><span class="n">val_losses</span><span class="p">],</span>
            <span class="ss">track_tokens_seen:</span> <span class="p">[</span><span class="n">tokens_seen</span> <span class="o">|</span> <span class="n">state</span><span class="o">.</span><span class="n">track_tokens_seen</span><span class="p">]</span>
        <span class="p">}</span>
      <span class="k">else</span>
        <span class="n">state</span>
      <span class="k">end</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">epoch_batches</span><span class="p">(%{</span><span class="ss">batches:</span> <span class="n">batches</span><span class="p">,</span> <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">})</span> <span class="k">do</span>
    <span class="no">Enum</span><span class="o">.</span><span class="n">shuffle</span><span class="p">(</span><span class="n">batches</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">epoch_batches</span><span class="p">(%{</span><span class="ss">batches:</span> <span class="n">batches</span><span class="p">})</span> <span class="k">do</span>
    <span class="n">batches</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">epoch_batches</span><span class="p">(%{</span><span class="ss">stream:</span> <span class="n">stream</span><span class="p">,</span> <span class="ss">length:</span> <span class="n">length</span><span class="p">})</span> <span class="k">do</span>
    <span class="no">Enum</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">stream</span><span class="p">,</span> <span class="n">length</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">stack_batch</span><span class="p">({%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">input_batch</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">target_batch</span><span class="p">},</span> <span class="n">device</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">input_batch</span><span class="p">,</span> <span class="n">device</span><span class="p">),</span> <span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">target_batch</span><span class="p">,</span> <span class="n">device</span><span class="p">)}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">stack_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">}</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">unzip</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="n">inputs</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">()</span> <span class="o">|&gt;</span> <span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">device</span><span class="p">),</span>
      <span class="n">targets</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">()</span> <span class="o">|&gt;</span> <span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>
  <span class="c1"># here we update model weights, central part of training algorithm</span>
  <span class="c1"># we have one gradient for each model parameter (weight)x</span>
  <span class="k">defp</span> <span class="n">optimizer_step</span><span class="p">(%</span><span class="no">AdamW</span><span class="p">{}</span> <span class="o">=</span> <span class="n">optimizer</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">gradients</span><span class="p">)</span> <span class="k">do</span>
    <span class="c1"># extract model weights</span>
    <span class="n">parameter_tensors</span> <span class="o">=</span> <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
    <span class="c1"># extract gradient weights</span>
    <span class="n">gradient_tensors</span> <span class="o">=</span> <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">gradients</span><span class="p">)</span>
    <span class="c1"># adamw m and v tensors</span>
    <span class="p">{</span><span class="n">m_tensors</span><span class="p">,</span> <span class="n">v_tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">adamw_moments</span><span class="p">(</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">gradient_tensors</span><span class="p">)</span>
    <span class="n">step</span> <span class="o">=</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">step</span> <span class="o">+</span> <span class="mi">1</span>

    <span class="c1"># update parameters using AdamW optimization</span>
    <span class="p">{</span><span class="n">updated_parameters</span><span class="p">,</span> <span class="p">{</span><span class="n">updated_m</span><span class="p">,</span> <span class="n">updated_v</span><span class="p">}}</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">([</span><span class="n">parameter_tensors</span><span class="p">,</span> <span class="n">gradient_tensors</span><span class="p">,</span> <span class="n">m_tensors</span><span class="p">,</span> <span class="n">v_tensors</span><span class="p">])</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map_reduce</span><span class="p">({[],</span> <span class="p">[]},</span> <span class="k">fn</span> <span class="p">{</span><span class="n">parameter</span><span class="p">,</span> <span class="n">gradient</span><span class="p">,</span> <span class="n">m</span><span class="p">,</span> <span class="n">v</span><span class="p">},</span> <span class="p">{</span><span class="n">updated_m</span><span class="p">,</span> <span class="n">updated_v</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="n">m</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">optimizer</span><span class="o">.</span><span class="n">beta1</span><span class="p">,</span> <span class="n">m</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">beta1</span><span class="p">,</span> <span class="n">gradient</span><span class="p">))</span>

        <span class="n">v</span> <span class="o">=</span>
          <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span>
            <span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">optimizer</span><span class="o">.</span><span class="n">beta2</span><span class="p">,</span> <span class="n">v</span><span class="p">),</span>
            <span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">beta2</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">gradient</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
          <span class="p">)</span>

        <span class="n">m_hat</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">divide</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="mf">1.0</span> <span class="o">-</span> <span class="ss">:math</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">optimizer</span><span class="o">.</span><span class="n">beta1</span><span class="p">,</span> <span class="n">step</span><span class="p">))</span>
        <span class="n">v_hat</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">divide</span><span class="p">(</span><span class="n">v</span><span class="p">,</span> <span class="mf">1.0</span> <span class="o">-</span> <span class="ss">:math</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">optimizer</span><span class="o">.</span><span class="n">beta2</span><span class="p">,</span> <span class="n">step</span><span class="p">))</span>

        <span class="n">update</span> <span class="o">=</span>
          <span class="n">m_hat</span>
          <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">divide</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">v_hat</span><span class="p">),</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">eps</span><span class="p">))</span>
          <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">optimizer</span><span class="o">.</span><span class="n">weight_decay</span><span class="p">,</span> <span class="n">parameter</span><span class="p">))</span>

        <span class="n">parameter</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">subtract</span><span class="p">(</span><span class="n">parameter</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">optimizer</span><span class="o">.</span><span class="n">learning_rate</span><span class="p">,</span> <span class="n">update</span><span class="p">))</span>

        <span class="p">{</span><span class="n">parameter</span><span class="p">,</span> <span class="p">{[</span><span class="n">m</span> <span class="o">|</span> <span class="n">updated_m</span><span class="p">],</span> <span class="p">[</span><span class="n">v</span> <span class="o">|</span> <span class="n">updated_v</span><span class="p">]}}</span>
      <span class="k">end</span><span class="p">)</span>
    <span class="c1"># update model with new weight values</span>
    <span class="p">{</span><span class="n">model</span><span class="p">,</span> <span class="p">[]}</span> <span class="o">=</span> <span class="n">put_trainable_tensors</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">updated_parameters</span><span class="p">)</span>

    <span class="c1"># the next batch can continue from the accumulated AdamW moment history.</span>
    <span class="n">optimizer</span> <span class="o">=</span> <span class="p">%{</span>
      <span class="n">optimizer</span>
      <span class="o">|</span> <span class="ss">step:</span> <span class="n">step</span><span class="p">,</span>
        <span class="ss">m:</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">(</span><span class="n">updated_m</span><span class="p">),</span>
        <span class="ss">v:</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reverse</span><span class="p">(</span><span class="n">updated_v</span><span class="p">)</span>
    <span class="p">}</span>

    <span class="p">{</span><span class="n">model</span><span class="p">,</span> <span class="n">optimizer</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">optimizer_step</span><span class="p">(</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">gradients</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_function</span><span class="p">(</span><span class="n">optimizer</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">optimizer</span><span class="o">.</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">gradients</span><span class="p">),</span> <span class="n">optimizer</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">adamw_moments</span><span class="p">(%</span><span class="no">AdamW</span><span class="p">{</span><span class="ss">m:</span> <span class="no">nil</span><span class="p">,</span> <span class="ss">v:</span> <span class="no">nil</span><span class="p">},</span> <span class="n">gradient_tensors</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">moments</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">gradient_tensors</span><span class="p">,</span> <span class="o">&amp;</span><span class="no">Nx</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">)))</span>
    <span class="p">{</span><span class="n">moments</span><span class="p">,</span> <span class="n">moments</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">adamw_moments</span><span class="p">(%</span><span class="no">AdamW</span><span class="p">{</span><span class="ss">m:</span> <span class="n">m</span><span class="p">,</span> <span class="ss">v:</span> <span class="n">v</span><span class="p">},</span> <span class="n">_gradient_tensors</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="p">{</span><span class="n">m</span><span class="p">,</span> <span class="n">v</span><span class="p">}</span>

  <span class="k">defp</span> <span class="n">trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">EmbeddingNative</span><span class="p">{</span><span class="ss">weight:</span> <span class="n">weight</span><span class="p">}),</span> <span class="k">do</span><span class="p">:</span> <span class="p">[</span><span class="n">weight</span><span class="p">]</span>

  <span class="k">defp</span> <span class="n">trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">DummyLayerNorm</span><span class="p">{</span><span class="ss">scale:</span> <span class="n">scale</span><span class="p">,</span> <span class="ss">shift:</span> <span class="n">shift</span><span class="p">})</span> <span class="k">do</span>
    <span class="p">[</span><span class="n">scale</span><span class="p">,</span> <span class="n">shift</span><span class="p">]</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">GELU</span><span class="p">{}),</span> <span class="k">do</span><span class="p">:</span> <span class="p">[]</span>

  <span class="k">defp</span> <span class="n">trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">FeedForward</span><span class="p">{</span><span class="ss">layers:</span> <span class="n">layers</span><span class="p">})</span> <span class="k">do</span>
    <span class="n">dense_tensors</span><span class="p">(</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span> <span class="o">++</span> <span class="n">dense_tensors</span><span class="p">(</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">MultiheadAttention</span><span class="p">{}</span> <span class="o">=</span> <span class="n">attention</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">dense_tensors</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_q</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">dense_tensors</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_k</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">dense_tensors</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_v</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">dense_tensors</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">out_proj</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">TransformerBlock</span><span class="p">{}</span> <span class="o">=</span> <span class="n">block</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">norm1</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">norm2</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">trainable_tensors</span><span class="p">(%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">tok_emb</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span><span class="p">)</span> <span class="o">++</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">flat_map</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">trainable_tensors</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">trainable_tensors</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">final_norm</span><span class="p">)</span> <span class="o">++</span>
      <span class="n">dense_tensors</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="p">,</span> <span class="no">false</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">dense_tensors</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">,</span> <span class="ss">bias:</span> <span class="n">bias</span><span class="p">},</span> <span class="no">true</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="p">[</span><span class="n">kernel</span><span class="p">,</span> <span class="n">bias</span><span class="p">]</span>
  <span class="k">defp</span> <span class="n">dense_tensors</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">},</span> <span class="no">false</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="p">[</span><span class="n">kernel</span><span class="p">]</span>

  <span class="k">defp</span> <span class="n">put_trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">EmbeddingNative</span><span class="p">{}</span> <span class="o">=</span> <span class="n">embedding</span><span class="p">,</span> <span class="p">[</span><span class="n">weight</span> <span class="o">|</span> <span class="n">rest</span><span class="p">])</span> <span class="k">do</span>
    <span class="p">{%{</span><span class="n">embedding</span> <span class="o">|</span> <span class="ss">weight:</span> <span class="n">weight</span><span class="p">},</span> <span class="n">rest</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">DummyLayerNorm</span><span class="p">{}</span> <span class="o">=</span> <span class="n">layer_norm</span><span class="p">,</span> <span class="p">[</span>
         <span class="n">scale</span><span class="p">,</span>
         <span class="n">shift</span> <span class="o">|</span> <span class="n">rest</span>
       <span class="p">])</span> <span class="k">do</span>
    <span class="p">{%{</span><span class="n">layer_norm</span> <span class="o">|</span> <span class="ss">scale:</span> <span class="n">scale</span><span class="p">,</span> <span class="ss">shift:</span> <span class="n">shift</span><span class="p">},</span> <span class="n">rest</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">GELU</span><span class="p">{}</span> <span class="o">=</span> <span class="n">gelu</span><span class="p">,</span> <span class="n">rest</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="p">{</span><span class="n">gelu</span><span class="p">,</span> <span class="n">rest</span><span class="p">}</span>

  <span class="k">defp</span> <span class="n">put_trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">FeedForward</span><span class="p">{}</span> <span class="o">=</span> <span class="n">feed_forward</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">first</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_dense_tensors</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="p">,</span> <span class="no">true</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">second</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_dense_tensors</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="p">,</span> <span class="no">true</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="p">%{</span><span class="n">feed_forward</span> <span class="o">|</span> <span class="ss">layers:</span> <span class="p">%{</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span> <span class="o">|</span> <span class="ss">first:</span> <span class="n">first</span><span class="p">,</span> <span class="ss">second:</span> <span class="n">second</span><span class="p">}},</span>
      <span class="n">tensors</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">MultiheadAttention</span><span class="p">{}</span> <span class="o">=</span> <span class="n">attention</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">w_q</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_dense_tensors</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_q</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">w_k</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_dense_tensors</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_k</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">w_v</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_dense_tensors</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_v</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">out_proj</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_dense_tensors</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">out_proj</span><span class="p">,</span> <span class="no">true</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="p">%{</span><span class="n">attention</span> <span class="o">|</span> <span class="ss">w_q:</span> <span class="n">w_q</span><span class="p">,</span> <span class="ss">w_k:</span> <span class="n">w_k</span><span class="p">,</span> <span class="ss">w_v:</span> <span class="n">w_v</span><span class="p">,</span> <span class="ss">out_proj:</span> <span class="n">out_proj</span><span class="p">},</span>
      <span class="n">tensors</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_trainable_tensors</span><span class="p">(%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">TransformerBlock</span><span class="p">{}</span> <span class="o">=</span> <span class="n">block</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">att</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_trainable_tensors</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">ff</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_trainable_tensors</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">norm1</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_trainable_tensors</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">norm1</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">norm2</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_trainable_tensors</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">norm2</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>

    <span class="p">{%{</span><span class="n">block</span> <span class="o">|</span> <span class="ss">att:</span> <span class="n">att</span><span class="p">,</span> <span class="ss">ff:</span> <span class="n">ff</span><span class="p">,</span> <span class="ss">norm1:</span> <span class="n">norm1</span><span class="p">,</span> <span class="ss">norm2:</span> <span class="n">norm2</span><span class="p">},</span> <span class="n">tensors</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_trainable_tensors</span><span class="p">(%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">tok_emb</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_trainable_tensors</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">tok_emb</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">pos_emb</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_trainable_tensors</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">trf_blocks</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map_reduce</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">,</span> <span class="n">tensors</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">put_trainable_tensors</span><span class="o">/</span><span class="mi">2</span><span class="p">)</span>
    <span class="p">{</span><span class="n">final_norm</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_trainable_tensors</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">final_norm</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
    <span class="p">{</span><span class="n">out_head</span><span class="p">,</span> <span class="n">tensors</span><span class="p">}</span> <span class="o">=</span> <span class="n">put_dense_tensors</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="p">,</span> <span class="no">false</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="p">%{</span>
        <span class="n">model</span>
        <span class="o">|</span> <span class="ss">tok_emb:</span> <span class="n">tok_emb</span><span class="p">,</span>
          <span class="ss">pos_emb:</span> <span class="n">pos_emb</span><span class="p">,</span>
          <span class="ss">trf_blocks:</span> <span class="n">trf_blocks</span><span class="p">,</span>
          <span class="ss">final_norm:</span> <span class="n">final_norm</span><span class="p">,</span>
          <span class="ss">out_head:</span> <span class="n">out_head</span>
      <span class="p">},</span>
      <span class="n">tensors</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_dense_tensors</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">_kernel</span><span class="p">,</span> <span class="ss">bias:</span> <span class="n">_bias</span><span class="p">}</span> <span class="o">=</span> <span class="n">dense</span><span class="p">,</span> <span class="no">true</span><span class="p">,</span> <span class="p">[</span><span class="n">kernel</span><span class="p">,</span> <span class="n">bias</span> <span class="o">|</span> <span class="n">rest</span><span class="p">])</span> <span class="k">do</span>
    <span class="p">{%{</span><span class="n">dense</span> <span class="o">|</span> <span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">,</span> <span class="ss">bias:</span> <span class="n">bias</span><span class="p">},</span> <span class="n">rest</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_dense_tensors</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">_kernel</span><span class="p">}</span> <span class="o">=</span> <span class="n">dense</span><span class="p">,</span> <span class="no">false</span><span class="p">,</span> <span class="p">[</span><span class="n">kernel</span> <span class="o">|</span> <span class="n">rest</span><span class="p">])</span> <span class="k">do</span>
    <span class="p">{%{</span><span class="n">dense</span> <span class="o">|</span> <span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">},</span> <span class="n">rest</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">maybe_transfer_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span> <span class="ow">when</span> <span class="n">device</span> <span class="ow">in</span> <span class="p">[</span><span class="no">nil</span><span class="p">,</span> <span class="ss">:default</span><span class="p">],</span> <span class="k">do</span><span class="p">:</span> <span class="n">model</span>
  <span class="k">defp</span> <span class="n">maybe_transfer_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span> <span class="ow">when</span> <span class="n">device</span> <span class="ow">in</span> <span class="p">[</span><span class="no">nil</span><span class="p">,</span> <span class="ss">:default</span><span class="p">],</span> <span class="k">do</span><span class="p">:</span> <span class="n">tensor</span>
  <span class="k">defp</span> <span class="n">maybe_transfer_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">format_loss</span><span class="p">(</span><span class="ss">:nan</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="s2">"nan"</span>
  <span class="k">defp</span> <span class="n">format_loss</span><span class="p">(</span><span class="n">loss</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="ss">:erlang</span><span class="o">.</span><span class="n">float_to_binary</span><span class="p">(</span><span class="n">loss</span> <span class="o">*</span> <span class="mf">1.0</span><span class="p">,</span> <span class="ss">decimals:</span> <span class="mi">3</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">pad_step</span><span class="p">(</span><span class="n">step</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">step</span>
    <span class="o">|&gt;</span> <span class="no">Integer</span><span class="o">.</span><span class="n">to_string</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">pad_leading</span><span class="p">(</span><span class="mi">6</span><span class="p">,</span> <span class="s2">"0"</span><span class="p">)</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And a test that exercises our model training algorithm:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
  <span class="nv">@tag</span> <span class="ss">:train</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">720_000</span>
  <span class="n">test</span> <span class="s2">"5.2 train an llm"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="n">device</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">file_content</span> <span class="o">=</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">(</span><span class="s2">"the-verdict.txt"</span><span class="p">)</span>
    <span class="n">train_ratio</span> <span class="o">=</span> <span class="mf">0.90</span>
    <span class="n">split_idx</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">train_ratio</span> <span class="o">*</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">file_content</span><span class="p">))</span>
    <span class="n">train_data</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">file_content</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">split_idx</span><span class="p">)</span>
    <span class="n">val_data</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">file_content</span><span class="p">,</span> <span class="n">split_idx</span><span class="p">,</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">file_content</span><span class="p">)</span> <span class="o">-</span> <span class="n">split_idx</span><span class="p">)</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">256</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.0</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="no">GptDatasetV1</span><span class="o">.</span><span class="n">create_dataloader_v1</span><span class="p">(</span>
        <span class="ss">raw_text:</span> <span class="n">train_data</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span>
        <span class="ss">max_length:</span> <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">stride:</span> <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="mi">0</span>
      <span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="no">GptDatasetV1</span><span class="o">.</span><span class="n">create_dataloader_v1</span><span class="p">(</span>
        <span class="ss">raw_text:</span> <span class="n">val_data</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span>
        <span class="ss">max_length:</span> <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">stride:</span> <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="mi">0</span>
      <span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="n">optimizer</span> <span class="o">=</span>
      <span class="no">Training</span><span class="o">.</span><span class="n">adamw</span><span class="p">(</span>
        <span class="mf">0.0004</span><span class="p">,</span>
        <span class="ss">weight_decay:</span> <span class="mf">0.1</span>
      <span class="p">)</span>

    <span class="p">{</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">train_losses</span><span class="p">,</span> <span class="n">val_losses</span><span class="p">,</span> <span class="n">tokens_seen</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Training</span><span class="o">.</span><span class="n">train_model_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">train_loader</span><span class="p">,</span>
        <span class="n">val_loader</span><span class="p">,</span>
        <span class="n">optimizer</span><span class="p">,</span>
        <span class="n">device</span><span class="p">,</span>
        <span class="mi">10</span><span class="p">,</span>
        <span class="mi">5</span><span class="p">,</span>
        <span class="mi">5</span><span class="p">,</span>
        <span class="s2">"Every effort moves you"</span><span class="p">,</span>
        <span class="s2">"code-davinci-002"</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">train_losses</span><span class="p">)</span> <span class="o">==</span> <span class="mi">18</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">val_losses</span><span class="p">)</span> <span class="o">==</span> <span class="mi">18</span>

    <span class="n">assert</span> <span class="n">tokens_seen</span> <span class="o">==</span> <span class="p">[</span>
             <span class="mi">512</span><span class="p">,</span>
             <span class="mi">3072</span><span class="p">,</span>
             <span class="mi">5632</span><span class="p">,</span>
             <span class="mi">8192</span><span class="p">,</span>
             <span class="mi">10_752</span><span class="p">,</span>
             <span class="mi">13_312</span><span class="p">,</span>
             <span class="mi">15_872</span><span class="p">,</span>
             <span class="mi">18_432</span><span class="p">,</span>
             <span class="mi">20_992</span><span class="p">,</span>
             <span class="mi">23_552</span><span class="p">,</span>
             <span class="mi">26_112</span><span class="p">,</span>
             <span class="mi">28_672</span><span class="p">,</span>
             <span class="mi">31_232</span><span class="p">,</span>
             <span class="mi">33_792</span><span class="p">,</span>
             <span class="mi">36_352</span><span class="p">,</span>
             <span class="mi">38_912</span><span class="p">,</span>
             <span class="mi">41_472</span><span class="p">,</span>
             <span class="mi">44_032</span>
           <span class="p">]</span>

    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">train_losses</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_float</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">val_losses</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_float</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">train_losses</span> <span class="o">++</span> <span class="n">val_losses</span><span class="p">,</span> <span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;1</span> <span class="o">&gt;</span> <span class="mf">0.0</span><span class="p">))</span>
    <span class="n">assert</span> <span class="p">%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">trained_model</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Here are training steps:</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
  Ep 1 (Step 000000): Train loss 10.054, Val loss 10.147
Every effort moves you, the, the the

Ep 1 (Step 000005): Train loss 8.111, Val loss 8.347
Every effort moves you,.

Ep 2 (Step 000010): Train loss 6.654, Val loss 7.084
Every effort moves you, the,, the, the, the, the, the,, the, the, the,, the,, the, the, the, the,, the, the.

Ep 2 (Step 000015): Train loss 5.949, Val loss 6.575
Every effort moves you, I had, I had I had I had I had I had I had, I had, I, I had I had I had, I had, I had, I had I had I had I had, I, I had I had

Ep 3 (Step 000020): Train loss 5.607, Val loss 6.451
Every effort moves you the" the the I had the I had.
" the the of the the I had the I, I had I had I had, the I had the I had of the I, the the of the the I, the the I had

Ep 3 (Step 000025): Train loss 5.468, Val loss 6.549
Every effort moves you sun

Ep 4 (Step 000030): Train loss 5.113, Val loss 6.408
Every effort moves you his

Ep 4 (Step 000035): Train loss 4.730, Val loss 6.403
Every effort moves you of the fact.

"I.


"I--and to have to have to have to have to the _--as Jack's the _.

Ep 5 (Step 000040): Train loss 4.427, Val loss 6.335
Every effort moves you.
"Oh, and I had a little a little.
"Oh, I had a little, I had been his

Ep 6 (Step 000045): Train loss 3.885, Val loss 6.219
Every effort moves you, and he had been the picture, and he had been the picture of the picture-rooms, and I felt, and I had been, and he had been of the, and I felt it, and I felt--as, and I had

Ep 6 (Step 000050): Train loss 3.446, Val loss 6.261
Every effort moves you know.

Ep 7 (Step 000055): Train loss 2.972, Val loss 6.232
Every effort moves you of the fact, and I felt it to have that lifted the fact of a self-rooms, I had been fact, the picture to me.
"I had been his pictures.
"Oh, I

Ep 7 (Step 000060): Train loss 2.444, Val loss 6.171
Every effort moves you in the fact, and I felt it to have been a good-room a self-confident moustache, I felt to see it.

Ep 8 (Step 000065): Train loss 1.944, Val loss 6.202
Every effort moves you?"
"I looked the Riviera him, I felt it--I had always--and by me--I had to the man of the moment--as Jack himself, I had the honour of the donkey, and he had been his own

Ep 8 (Step 000070): Train loss 1.613, Val loss 6.207
Every effort moves you?"

"Yes--I glanced after him, and Mrs.

"Oh, you know, in fact, you know; and the moment--as Jack himself, one might put it, and up and down the room, with your

Ep 9 (Step 000075): Train loss 1.243, Val loss 6.285
Every effort moves you?"

"Yes--I glanced after him, struck by his last word. Gisburn's an awful simpleton, and muddling; then I looked at the donkey--as I had the secret, and were the hour.

Ep 9 (Step 000080): Train loss 0.979, Val loss 6.284
Every effort moves you?"

"Yes--quite insensible to the fact with his last word. Victor Grindle was, in fact, the cigars you like."
"I looked at the donkey--if I saw that, and down the room, I found

Ep 10 (Step 000085): Train loss 0.817, Val loss 6.394
Every effort moves you?"

"Yes--quite insensible to the fact with equanimity.

" to what a degree he had the same quality as his pictures--the quality of looking cleverer than he was prince
.
Finished in 681.6 seconds (0.00s async, 681.6s sync)
</code></pre></div></div>

<p>You can see that it took us 11 minutes to train the model on this small set of training data. It is important to note the discrepancy between training loss and validation loss. Training loss has an acceptable value, but validation loss does not. Our prediction is grammatically correct, but it makes no sense. What is the reason for that?</p>

<p>During training, the model sees only the training split. In this project, that is the first 90% of the-verdict.txt. The validation split is the remaining 10%, which the model does not use for training.</p>

<p>So if after training you see:</p>

<ul>
  <li>training loss low</li>
  <li>validation loss much higher</li>
  <li>that means:</li>
</ul>

<p>the model predicts the training text well but predicts unseen validation text poorly. That is overfitting.
The “memorizes the training data verbatim” part means: instead of learning broadly useful language patterns, the model starts reproducing exact phrases from the training text. For example, if the generated text contains:</p>

<p>quite insensible to the irony</p>

<p>and you search the-verdict.txt and find that exact phrase in the training portion, then the model probably did not invent a generally good continuation. It remembered a phrase from the data it trained on.</p>

<p>A simpler analogy:</p>

<p>Training data:
“The cat sat on the mat.”</p>

<p>Validation data:
“The dog slept by the fire.”<br />
An overfit model might become very good at continuing:</p>

<p>“The cat sat…”
with:</p>

<p>“on the mat”
because it memorized that exact sentence.</p>

<p>But when asked to predict unseen text like:</p>

<p>“The dog slept…”
it performs badly, because it did not really learn the broader pattern. It mostly memorized examples.</p>

<p>So the book is saying:</p>

<p>The model’s training loss decreases because it gets good at the text it has already seen.
The validation loss stays higher because it does not perform as well on unseen text.
If generated snippets appear verbatim in the original training file, that is evidence of memorization.
Memorization plus poor validation performance is overfitting.</p>

<p>And here is a graph that shows the difference between training and validation loss:</p>

<p><img src="/assets/2026/05/training_validation_loss.png" alt="Training and validation loss over epochs" /></p>

<h2 id="decoding-strategies-to-control-randomness">Decoding strategies to control randomness</h2>

<p>Decoding strategies are text generation strategies whose goal is to generate more original text. We first started with the <code class="language-plaintext highlighter-rouge">generate_text_simple</code> function, which generates the next token based on the highest probabilities in the logits matrix. This is called the greedy algorithm. But we can do better with temperature scaling and top-k sampling.</p>

<p>Let’s recap the greedy algorithm. We also implemented a helper module that loads and saves a trained GPT model, so we do not need to run the training algorithm again. The size of our GPT model is 623 MB!</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">ModelCheckpoint</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Saves and loads Nx model containers.

  This is intended for trained models such as `LlmScratch.GPTModel`, whose
  tensors are serializable through `Nx.Container`.
  """</span>

  <span class="nv">@type</span> <span class="n">path</span> <span class="p">::</span> <span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Serializes `model` and writes it to `path`.

  Parent directories are created automatically. The model must implement
  `Nx.Container`; `LlmScratch.GPTModel` does through its container
  implementation in `LlmScratch.Training`.

  ## Example

      checkpoint_path = "ch5_2_gpt_124m.nx"
      LlmScratch.ModelCheckpoint.save!(trained_model, checkpoint_path)

  """</span>
  <span class="nv">@spec</span> <span class="n">save!</span><span class="p">(</span><span class="n">struct</span><span class="p">(),</span> <span class="n">path</span><span class="p">())</span> <span class="p">::</span> <span class="ss">:ok</span>
  <span class="k">def</span> <span class="n">save!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">path</span>
    <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">dirname</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">mkdir_p!</span><span class="p">()</span>

    <span class="no">File</span><span class="o">.</span><span class="n">write!</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">serialize</span><span class="p">(</span><span class="n">model</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Reads a checkpoint from `path` and deserializes it back into a model struct.

  The checkpoint must have been written with `save!/2` or another compatible
  `Nx.serialize/1` call.

  ## Example

      trained_model = LlmScratch.ModelCheckpoint.load!("ch5_2_gpt_124m.nx")

  """</span>
  <span class="nv">@spec</span> <span class="n">load!</span><span class="p">(</span><span class="n">path</span><span class="p">())</span> <span class="p">::</span> <span class="n">struct</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">load!</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">path</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">deserialize</span><span class="p">()</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And a test that loads the trained model and does greedy text generation. We repeat text generation to show that it is currently deterministic when we use the same GPT model:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="n">test</span> <span class="s2">"5.3 load trained model checkpoint and generate text"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="s2">"ch5_2_gpt_124m.nx"</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="n">trained_model</span> <span class="o">=</span> <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">load!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>

    <span class="n">token_ids</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate_text_simple</span><span class="p">(</span>
        <span class="n">trained_model</span><span class="p">,</span>
        <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="s2">"Every effort moves you"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">),</span>
        <span class="mi">25</span><span class="p">,</span>
        <span class="n">trained_model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">repeated_token_ids</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate_text_simple</span><span class="p">(</span>
        <span class="n">trained_model</span><span class="p">,</span>
        <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="s2">"Every effort moves you"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">),</span>
        <span class="mi">25</span><span class="p">,</span>
        <span class="n">trained_model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">decoded_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>
    <span class="n">repeated_decoded_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">repeated_token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">token_ids</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">29</span><span class="p">}</span>

    <span class="n">assert</span> <span class="n">decoded_text</span> <span class="o">==</span>
             <span class="sd">"""
             Every effort moves you?"

             "Yes--quite insensible to the fact with a laugh: "Yes--and by me to me to\
             """</span>

    <span class="n">assert</span> <span class="n">decoded_text</span> <span class="o">==</span> <span class="n">repeated_decoded_text</span>
  <span class="k">end</span>
</code></pre></div></div>

<h3 id="531-temperature-scaling">5.3.1 Temperature scaling</h3>

<p>Let’s recap the greedy algorithm in the <code class="language-plaintext highlighter-rouge">generate_text_simple</code> function, where we select the next token based on the highest logits value using <code class="language-plaintext highlighter-rouge">Nx.argmax</code>:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
  <span class="k">def</span> <span class="n">generate_text_simple</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">idx</span><span class="p">,</span> <span class="n">max_new_tokens</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">max_new_tokens</span><span class="p">)</span> <span class="ow">and</span> <span class="n">max_new_tokens</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">context_size</span><span class="p">)</span> <span class="ow">and</span>
             <span class="n">context_size</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="n">validate_idx_shape!</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">1</span><span class="o">..</span><span class="n">max_new_tokens</span><span class="o">//</span><span class="mi">1</span><span class="p">,</span> <span class="n">idx</span><span class="p">,</span> <span class="k">fn</span> <span class="n">_step</span><span class="p">,</span> <span class="n">acc</span> <span class="o">-&gt;</span>
      <span class="n">idx_cond</span> <span class="o">=</span> <span class="n">last_tokens</span><span class="p">(</span><span class="n">acc</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
      <span class="n">logits</span> <span class="o">=</span> <span class="n">forward!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">idx_cond</span><span class="p">)</span>
      <span class="n">logits</span> <span class="o">=</span> <span class="n">last_position_logits</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span>
      <span class="n">probas</span> <span class="o">=</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>

      <span class="n">idx_next</span> <span class="o">=</span>
        <span class="n">probas</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="ss">keep_axis:</span> <span class="no">true</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">as_type</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">acc</span><span class="p">))</span>

      <span class="no">Nx</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">acc</span><span class="p">,</span> <span class="n">idx_next</span><span class="p">],</span> <span class="ss">axis:</span> <span class="mi">1</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>What would happen when we replace <code class="language-plaintext highlighter-rouge">Nx.argmax</code> with a probabilistic sampling process? First, we need a module with some helper functions that will be used in the test.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
  <span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">TemperatureScaling</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Sampling helpers used by temperature-scaled text generation examples.

  `multinomial/2` mirrors `torch.multinomial(probas, num_samples: n)` for a
  one-dimensional probability tensor. It samples token ids from `0` to
  `Nx.size(probas) - 1`.
  """</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Applies temperature scaling to `logits`, then returns softmax probabilities.

  This mirrors:

      scaled_logits = logits / temperature
      torch.softmax(scaled_logits, dim=0)

  A lower temperature sharpens the distribution; a higher temperature flattens
  it. `temperature` must be greater than zero.

  ## Example

      logits = Nx.tensor([1.0, 2.0, 3.0])
      LlmScratch.TemperatureScaling.softmax_with_temperature(logits, 1.0)

  """</span>
  <span class="nv">@spec</span> <span class="n">softmax_with_temperature</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">number</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">softmax_with_temperature</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">temperature</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_number</span><span class="p">(</span><span class="n">temperature</span><span class="p">)</span> <span class="ow">and</span> <span class="n">temperature</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="n">logits</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">divide</span><span class="p">(</span><span class="n">temperature</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Samples token ids from `probas`.

  Options:

    * `:samples` - number of token ids to sample.
    * `:seed` - deterministic seed for `Nx.Random.key/1`.

  ## Example

      probas = Nx.tensor([0.25, 0.75])
      LlmScratch.TemperatureScaling.multinomial(probas, samples: 3, seed: 123)

  """</span>
  <span class="nv">@spec</span> <span class="n">multinomial</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="p">[</span><span class="n">integer</span><span class="p">()]</span>
  <span class="k">def</span> <span class="n">multinomial</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">samples</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:samples</span><span class="p">)</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">)</span>

    <span class="p">{</span><span class="n">sampled_token_ids</span><span class="p">,</span> <span class="n">_key</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="n">seed</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">iota</span><span class="p">({</span><span class="no">Nx</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">probas</span><span class="p">)}),</span> <span class="n">probas</span><span class="p">,</span> <span class="ss">samples:</span> <span class="n">samples</span><span class="p">)</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">(</span><span class="n">sampled_token_ids</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Samples repeatedly from `probas` and returns `{frequency, token}` pairs.

  `inverse_vocab` maps token ids to token strings, for example
  `%{0 =&gt; "closer", 1 =&gt; "every"}`.
  """</span>
  <span class="nv">@spec</span> <span class="n">sampled_token_frequencies</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="p">%{</span><span class="n">integer</span><span class="p">()</span> <span class="o">=&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()},</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="p">[</span>
          <span class="p">{</span><span class="n">non_neg_integer</span><span class="p">(),</span> <span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">()}</span>
        <span class="p">]</span>
  <span class="k">def</span> <span class="n">sampled_token_frequencies</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="n">inverse_vocab</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">inverse_vocab</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">sampled_ids</span> <span class="o">=</span> <span class="n">multinomial</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="n">frequencies</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">frequencies</span><span class="p">(</span><span class="n">sampled_ids</span><span class="p">)</span>

    <span class="mi">0</span><span class="o">..</span><span class="p">(</span><span class="n">map_size</span><span class="p">(</span><span class="n">inverse_vocab</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="n">token_id</span> <span class="o">-&gt;</span> <span class="p">{</span><span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">frequencies</span><span class="p">,</span> <span class="n">token_id</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">inverse_vocab</span><span class="p">[</span><span class="n">token_id</span><span class="p">]}</span> <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>We will show this in a test that uses a simple vocabulary with probability values that are not real, but are good enough to explain the concept.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
  <span class="n">test</span> <span class="s2">"5.3 samples next token from vocabulary probabilities"</span> <span class="k">do</span>
    <span class="n">vocab</span> <span class="o">=</span> <span class="p">%{</span>
      <span class="s2">"closer"</span> <span class="o">=&gt;</span> <span class="mi">0</span><span class="p">,</span>
      <span class="s2">"every"</span> <span class="o">=&gt;</span> <span class="mi">1</span><span class="p">,</span>
      <span class="s2">"effort"</span> <span class="o">=&gt;</span> <span class="mi">2</span><span class="p">,</span>
      <span class="s2">"forward"</span> <span class="o">=&gt;</span> <span class="mi">3</span><span class="p">,</span>
      <span class="s2">"inches"</span> <span class="o">=&gt;</span> <span class="mi">4</span><span class="p">,</span>
      <span class="s2">"moves"</span> <span class="o">=&gt;</span> <span class="mi">5</span><span class="p">,</span>
      <span class="s2">"pizza"</span> <span class="o">=&gt;</span> <span class="mi">6</span><span class="p">,</span>
      <span class="s2">"toward"</span> <span class="o">=&gt;</span> <span class="mi">7</span><span class="p">,</span>
      <span class="s2">"you"</span> <span class="o">=&gt;</span> <span class="mi">8</span>
    <span class="p">}</span>

    <span class="n">inverse_vocab</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">vocab</span><span class="p">,</span> <span class="k">fn</span> <span class="p">{</span><span class="n">token</span><span class="p">,</span> <span class="n">token_id</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="p">{</span><span class="n">token_id</span><span class="p">,</span> <span class="n">token</span><span class="p">}</span> <span class="k">end</span><span class="p">)</span>

    <span class="c1"># Fictional logits for input text "every effort moves you"</span>
    <span class="n">next_token_logits</span> <span class="o">=</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">4.51</span><span class="p">,</span> <span class="mf">0.89</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.90</span><span class="p">,</span> <span class="mf">6.75</span><span class="p">,</span> <span class="mf">1.63</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.62</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.89</span><span class="p">,</span> <span class="mf">6.28</span><span class="p">,</span> <span class="mf">1.79</span><span class="p">])</span>

    <span class="c1"># logits to probability values</span>
    <span class="n">probas</span> <span class="o">=</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">next_token_logits</span><span class="p">)</span>

    <span class="c1"># greedy algorithm</span>
    <span class="n">next_token_id</span> <span class="o">=</span>
      <span class="n">probas</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span>

    <span class="n">assert</span> <span class="n">inverse_vocab</span><span class="p">[</span><span class="n">next_token_id</span><span class="p">]</span> <span class="o">==</span> <span class="s2">"forward"</span>

    <span class="c1"># this is equivalent to torch.multinomial</span>
    <span class="c1"># we are using probability sampling process instead of Nx.argmax</span>
    <span class="n">sampled_token_id</span> <span class="o">=</span> <span class="no">TemperatureScaling</span><span class="o">.</span><span class="n">multinomial</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="ss">samples:</span> <span class="mi">1</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="n">hd</span><span class="p">()</span>

    <span class="n">assert</span> <span class="n">inverse_vocab</span><span class="p">[</span><span class="n">sampled_token_id</span><span class="p">]</span> <span class="o">==</span> <span class="s2">"toward"</span>

    <span class="c1"># let's repeat sampling more times to get the token distribution</span>
    <span class="c1"># Our GPTModel becomes more creative, or better to say non-deterministic,</span>
    <span class="c1"># because it will generate according to this distribution:</span>
    <span class="c1"># Every effort moves you forward</span>
    <span class="c1"># Every effort moves you toward</span>
    <span class="c1"># Every effort move you closer</span>
    <span class="c1"># But it will never generate Every effort moves you pizza.</span>
    <span class="n">assert</span> <span class="no">TemperatureScaling</span><span class="o">.</span><span class="n">sampled_token_frequencies</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="n">inverse_vocab</span><span class="p">,</span>
             <span class="ss">seed:</span> <span class="mi">123</span><span class="p">,</span>
             <span class="ss">samples:</span> <span class="mi">1_000</span>
           <span class="p">)</span> <span class="o">==</span> <span class="p">[</span>
             <span class="p">{</span><span class="mi">64</span><span class="p">,</span> <span class="s2">"closer"</span><span class="p">},</span>
             <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="s2">"every"</span><span class="p">},</span>
             <span class="p">{</span><span class="mi">0</span><span class="p">,</span> <span class="s2">"effort"</span><span class="p">},</span>
             <span class="p">{</span><span class="mi">572</span><span class="p">,</span> <span class="s2">"forward"</span><span class="p">},</span>
             <span class="p">{</span><span class="mi">3</span><span class="p">,</span> <span class="s2">"inches"</span><span class="p">},</span>
             <span class="p">{</span><span class="mi">0</span><span class="p">,</span> <span class="s2">"moves"</span><span class="p">},</span>
             <span class="p">{</span><span class="mi">0</span><span class="p">,</span> <span class="s2">"pizza"</span><span class="p">},</span>
             <span class="p">{</span><span class="mi">357</span><span class="p">,</span> <span class="s2">"toward"</span><span class="p">},</span>
             <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="s2">"you"</span><span class="p">}</span>
           <span class="p">]</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Using <code class="language-plaintext highlighter-rouge">multinomial</code> just explains the concept of probability distribution, but we will not use it in the real case. Instead, we scale softmax logits values with temperature using <code class="language-plaintext highlighter-rouge">LlmScratch.TemperatureScaling.softmax_with_temperature</code> in the following Livebook code:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">vocab</span> <span class="o">=</span> <span class="p">[</span>
  <span class="p">{</span><span class="s2">"closer"</span><span class="p">,</span> <span class="mi">0</span><span class="p">},</span>
  <span class="p">{</span><span class="s2">"every"</span><span class="p">,</span> <span class="mi">1</span><span class="p">},</span>
  <span class="p">{</span><span class="s2">"effort"</span><span class="p">,</span> <span class="mi">2</span><span class="p">},</span>
  <span class="p">{</span><span class="s2">"forward"</span><span class="p">,</span> <span class="mi">3</span><span class="p">},</span>
  <span class="p">{</span><span class="s2">"inches"</span><span class="p">,</span> <span class="mi">4</span><span class="p">},</span>
  <span class="p">{</span><span class="s2">"moves"</span><span class="p">,</span> <span class="mi">5</span><span class="p">},</span>
  <span class="p">{</span><span class="s2">"pizza"</span><span class="p">,</span> <span class="mi">6</span><span class="p">},</span>
  <span class="p">{</span><span class="s2">"toward"</span><span class="p">,</span> <span class="mi">7</span><span class="p">},</span>
  <span class="p">{</span><span class="s2">"you"</span><span class="p">,</span> <span class="mi">8</span><span class="p">}</span>
<span class="p">]</span>

<span class="n">next_token_logits</span> <span class="o">=</span>
  <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">4.51</span><span class="p">,</span> <span class="mf">0.89</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.90</span><span class="p">,</span> <span class="mf">6.75</span><span class="p">,</span> <span class="mf">1.63</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.62</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.89</span><span class="p">,</span> <span class="mf">6.28</span><span class="p">,</span> <span class="mf">1.79</span><span class="p">])</span>

<span class="n">temperatures</span> <span class="o">=</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mi">5</span><span class="p">]</span>

<span class="n">scaled_probas</span> <span class="o">=</span>
  <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">temperatures</span><span class="p">,</span> <span class="k">fn</span> <span class="n">temperature</span> <span class="o">-&gt;</span>
    <span class="no">TemperatureScaling</span><span class="o">.</span><span class="n">softmax_with_temperature</span><span class="p">(</span><span class="n">next_token_logits</span><span class="p">,</span> <span class="n">temperature</span><span class="p">)</span>
  <span class="k">end</span><span class="p">)</span>
</code></pre></div></div>

<h2 id="plot">Plot</h2>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">alias</span> <span class="no">VegaLite</span><span class="p">,</span> <span class="ss">as:</span> <span class="no">Vl</span>

<span class="n">plot_data</span> <span class="o">=</span>
  <span class="n">temperatures</span>
  <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">(</span><span class="n">scaled_probas</span><span class="p">)</span>
  <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">flat_map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">temperature</span><span class="p">,</span> <span class="n">probas</span><span class="p">}</span> <span class="o">-&gt;</span>
    <span class="n">probas</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">(</span><span class="n">probas</span><span class="p">)</span>

    <span class="n">vocab</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">token</span><span class="p">,</span> <span class="n">token_id</span><span class="p">}</span> <span class="o">-&gt;</span>
      <span class="p">%{</span>
        <span class="ss">token:</span> <span class="n">token</span><span class="p">,</span>
        <span class="ss">token_id:</span> <span class="n">token_id</span><span class="p">,</span>
        <span class="ss">temperature:</span> <span class="s2">"Temperature = </span><span class="si">#{</span><span class="n">temperature</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span>
        <span class="ss">probability:</span> <span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="n">token_id</span><span class="p">)</span>
      <span class="p">}</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span><span class="p">)</span>

<span class="no">Vl</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">width:</span> <span class="mi">500</span><span class="p">,</span> <span class="ss">height:</span> <span class="mi">300</span><span class="p">)</span>
<span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">data_from_values</span><span class="p">(</span><span class="n">plot_data</span><span class="p">)</span>
<span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">mark</span><span class="p">(</span><span class="ss">:bar</span><span class="p">)</span>
<span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:x</span><span class="p">,</span> <span class="s2">"token"</span><span class="p">,</span>
  <span class="ss">type:</span> <span class="ss">:nominal</span><span class="p">,</span>
  <span class="ss">title:</span> <span class="no">nil</span><span class="p">,</span>
  <span class="ss">sort:</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="n">vocab</span><span class="p">,</span> <span class="k">fn</span> <span class="p">{</span><span class="n">token</span><span class="p">,</span> <span class="n">_token_id</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="n">token</span> <span class="k">end</span><span class="p">),</span>
  <span class="ss">axis:</span> <span class="p">[</span><span class="ss">labelAngle:</span> <span class="o">-</span><span class="mi">90</span><span class="p">]</span>
<span class="p">)</span>
<span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:x_offset</span><span class="p">,</span> <span class="s2">"temperature"</span><span class="p">,</span> <span class="ss">type:</span> <span class="ss">:nominal</span><span class="p">)</span>
<span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:y</span><span class="p">,</span> <span class="s2">"probability"</span><span class="p">,</span> <span class="ss">type:</span> <span class="ss">:quantitative</span><span class="p">,</span> <span class="ss">title:</span> <span class="s2">"Probability"</span><span class="p">)</span>
<span class="o">|&gt;</span> <span class="no">Vl</span><span class="o">.</span><span class="n">encode_field</span><span class="p">(</span><span class="ss">:color</span><span class="p">,</span> <span class="s2">"temperature"</span><span class="p">,</span> <span class="ss">type:</span> <span class="ss">:nominal</span><span class="p">,</span> <span class="ss">title:</span> <span class="no">nil</span><span class="p">)</span>
<span class="o">|&gt;</span> <span class="no">Kino</span><span class="o">.</span><span class="no">VegaLite</span><span class="o">.</span><span class="n">new</span><span class="p">()</span>
</code></pre></div></div>

<p>That gives us the following graph:</p>

<p><img src="/assets/2026/05/temperature_scaling.png" alt="Temperature scaling probability distribution" /></p>

<p>You can see that a lower temperature means a spikier distribution. The model becomes more predictable and less creative.</p>

<p>And we have another exercise:</p>

<p>Exercise 5.1</p>

<p><em>Use the print_sampled_tokens function to print the sampling frequencies of the softmax probabilities scaled with the temperatures shown in figure 5.14. How often is the word pizza sampled in each case? Can you think of a faster and more accurate way to determine how often the word pizza is sampled?</em></p>

<p>Faster way to calculate temperature scaling probability value is to multiply probs value for token <code class="language-plaintext highlighter-rouge">pizza</code> with sampling size:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"exercise 5.1 counts pizza sampling frequency across temperatures"</span> <span class="k">do</span>
    <span class="n">vocab</span> <span class="o">=</span> <span class="p">%{</span>
      <span class="s2">"closer"</span> <span class="o">=&gt;</span> <span class="mi">0</span><span class="p">,</span>
      <span class="s2">"every"</span> <span class="o">=&gt;</span> <span class="mi">1</span><span class="p">,</span>
      <span class="s2">"effort"</span> <span class="o">=&gt;</span> <span class="mi">2</span><span class="p">,</span>
      <span class="s2">"forward"</span> <span class="o">=&gt;</span> <span class="mi">3</span><span class="p">,</span>
      <span class="s2">"inches"</span> <span class="o">=&gt;</span> <span class="mi">4</span><span class="p">,</span>
      <span class="s2">"moves"</span> <span class="o">=&gt;</span> <span class="mi">5</span><span class="p">,</span>
      <span class="s2">"pizza"</span> <span class="o">=&gt;</span> <span class="mi">6</span><span class="p">,</span>
      <span class="s2">"toward"</span> <span class="o">=&gt;</span> <span class="mi">7</span><span class="p">,</span>
      <span class="s2">"you"</span> <span class="o">=&gt;</span> <span class="mi">8</span>
    <span class="p">}</span>

    <span class="n">inverse_vocab</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">vocab</span><span class="p">,</span> <span class="k">fn</span> <span class="p">{</span><span class="n">token</span><span class="p">,</span> <span class="n">token_id</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="p">{</span><span class="n">token_id</span><span class="p">,</span> <span class="n">token</span><span class="p">}</span> <span class="k">end</span><span class="p">)</span>
    <span class="n">next_token_logits</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">4.51</span><span class="p">,</span> <span class="mf">0.89</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.90</span><span class="p">,</span> <span class="mf">6.75</span><span class="p">,</span> <span class="mf">1.63</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.62</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.89</span><span class="p">,</span> <span class="mf">6.28</span><span class="p">,</span> <span class="mf">1.79</span><span class="p">])</span>

    <span class="n">pizza_sample_counts</span> <span class="o">=</span>
      <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mi">5</span><span class="p">]</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="n">temperature</span> <span class="o">-&gt;</span>
        <span class="n">probas</span> <span class="o">=</span> <span class="no">TemperatureScaling</span><span class="o">.</span><span class="n">softmax_with_temperature</span><span class="p">(</span><span class="n">next_token_logits</span><span class="p">,</span> <span class="n">temperature</span><span class="p">)</span>

        <span class="p">{</span><span class="n">pizza_count</span><span class="p">,</span> <span class="s2">"pizza"</span><span class="p">}</span> <span class="o">=</span>
          <span class="n">probas</span>
          <span class="o">|&gt;</span> <span class="no">TemperatureScaling</span><span class="o">.</span><span class="n">sampled_token_frequencies</span><span class="p">(</span><span class="n">inverse_vocab</span><span class="p">,</span>
            <span class="ss">seed:</span> <span class="mi">123</span><span class="p">,</span>
            <span class="ss">samples:</span> <span class="mi">1_000</span>
          <span class="p">)</span>
          <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">_count</span><span class="p">,</span> <span class="n">token</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="n">token</span> <span class="o">==</span> <span class="s2">"pizza"</span> <span class="k">end</span><span class="p">)</span>

        <span class="p">{</span><span class="n">temperature</span><span class="p">,</span> <span class="n">pizza_count</span><span class="p">}</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">pizza_sample_counts</span> <span class="o">==</span> <span class="p">[{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">},</span> <span class="p">{</span><span class="mf">0.1</span><span class="p">,</span> <span class="mi">0</span><span class="p">},</span> <span class="p">{</span><span class="mi">5</span><span class="p">,</span> <span class="mi">39</span><span class="p">}]</span>

    <span class="c1"># faster and accurate count for pizza</span>
    <span class="n">pizza_expected_counts</span> <span class="o">=</span>
      <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mi">5</span><span class="p">]</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="n">temperature</span> <span class="o">-&gt;</span>
        <span class="n">probas</span> <span class="o">=</span> <span class="no">TemperatureScaling</span><span class="o">.</span><span class="n">softmax_with_temperature</span><span class="p">(</span><span class="n">next_token_logits</span><span class="p">,</span> <span class="n">temperature</span><span class="p">)</span>
        <span class="c1"># multiply the probability value for pizza by the sampling size</span>
        <span class="p">{</span><span class="n">temperature</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">(</span><span class="n">probas</span><span class="p">[</span><span class="n">vocab</span><span class="p">[</span><span class="s2">"pizza"</span><span class="p">]])</span> <span class="o">*</span> <span class="mi">1_000</span><span class="p">}</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="n">assert</span> <span class="p">[{</span><span class="mi">1</span><span class="p">,</span> <span class="n">temperature_1_count</span><span class="p">},</span> <span class="p">{</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">nearly_zero</span><span class="p">},</span> <span class="p">{</span><span class="mi">5</span><span class="p">,</span> <span class="n">temperature_5_count</span><span class="p">}]</span> <span class="o">=</span>
             <span class="n">pizza_expected_counts</span>

    <span class="n">assert_in_delta</span> <span class="n">temperature_1_count</span><span class="p">,</span> <span class="mf">0.10120050865225494</span><span class="p">,</span> <span class="mf">1.0e-10</span>
    <span class="n">assert</span> <span class="n">nearly_zero</span> <span class="o">&lt;</span> <span class="mf">1.0e-30</span>
    <span class="n">assert_in_delta</span> <span class="n">temperature_5_count</span><span class="p">,</span> <span class="mf">42.99795627593994</span><span class="p">,</span> <span class="mf">1.0e-10</span>
  <span class="k">end</span>
</code></pre></div></div>

<h3 id="532-top-k-sampling">5.3.2 Top-k sampling</h3>

<p>If we just use temperature scaling and probabilistic sampling, we could hit the problem that our GPT model becomes too creative, so it outputs something like “Every effort moves you pizza.” To avoid this, we implement top-k sampling.</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>Vocabulary:             "closer"      "every"      "effort"     "forward"     "inches"     "moves"      "pizza"      "toward"     "you"
Index position:             0             1            2             3             4            5            6             7           8

---------------------------------------------------------------------------------------------------------------------------------------

Logits                  = [ 4.51,        0.89,       -1.90,        6.75,         1.63,       -1.62,       -1.89,        6.28,       1.79 ]
                         |
                         v
Top-k (k = 3)           = [ [4.51],      0.89,       -1.90,       [6.75],       1.63,       -1.62,       -1.89,       [6.28],      1.79 ]
                         |
                         v
-inf mask               = [ [4.51],      -inf,       -inf,       [6.75],       -inf,       -inf,       -inf,       [6.28],      -inf ]
                         |
                         v
Softmax                 = [ [0.06],      0.00,        0.00,       [0.57],       0.00,        0.00,        0.00,       [0.36],      0.00 ]
</code></pre></div></div>

<p>In top-k sampling, we select only the top k logits values as candidates for probabilistic sampling and temperature scaling. The rest of the values are set to -inf, so softmax sets them to 0.00.</p>

<p>Here is a test that shows the concept of top-k sampling:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"5.3.2 top-k sampling masks logits outside top k"</span> <span class="k">do</span>
  <span class="n">next_token_logits</span> <span class="o">=</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">4.51</span><span class="p">,</span> <span class="mf">0.89</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.90</span><span class="p">,</span> <span class="mf">6.75</span><span class="p">,</span> <span class="mf">1.63</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.62</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.89</span><span class="p">,</span> <span class="mf">6.28</span><span class="p">,</span> <span class="mf">1.79</span><span class="p">])</span>

  <span class="n">top_k</span> <span class="o">=</span> <span class="mi">3</span>
  <span class="p">{</span><span class="n">top_logits</span><span class="p">,</span> <span class="n">top_pos</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">top_k</span><span class="p">(</span><span class="n">next_token_logits</span><span class="p">,</span> <span class="ss">k:</span> <span class="n">top_k</span><span class="p">)</span>

  <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all_close</span><span class="p">(</span><span class="n">top_logits</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">6.75</span><span class="p">,</span> <span class="mf">6.28</span><span class="p">,</span> <span class="mf">4.51</span><span class="p">]))</span>
  <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">(</span><span class="n">top_pos</span><span class="p">)</span> <span class="o">==</span> <span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">7</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>

  <span class="n">min_top_logit</span> <span class="o">=</span> <span class="n">top_logits</span><span class="p">[</span><span class="n">top_k</span> <span class="o">-</span> <span class="mi">1</span><span class="p">]</span>

  <span class="n">new_logits</span> <span class="o">=</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">select</span><span class="p">(</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">next_token_logits</span><span class="p">,</span> <span class="n">min_top_logit</span><span class="p">),</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="ss">:neg_infinity</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">next_token_logits</span><span class="p">)),</span>
      <span class="n">next_token_logits</span>
    <span class="p">)</span>

  <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all_close</span><span class="p">(</span>
           <span class="n">new_logits</span><span class="p">,</span>
           <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span>
             <span class="mf">4.51</span><span class="p">,</span>
             <span class="ss">:neg_infinity</span><span class="p">,</span>
             <span class="ss">:neg_infinity</span><span class="p">,</span>
             <span class="mf">6.75</span><span class="p">,</span>
             <span class="ss">:neg_infinity</span><span class="p">,</span>
             <span class="ss">:neg_infinity</span><span class="p">,</span>
             <span class="ss">:neg_infinity</span><span class="p">,</span>
             <span class="mf">6.28</span><span class="p">,</span>
             <span class="ss">:neg_infinity</span>
           <span class="p">])</span>
         <span class="p">)</span>

  <span class="n">topk_probas</span> <span class="o">=</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">new_logits</span><span class="p">)</span>

  <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all_close</span><span class="p">(</span>
           <span class="n">topk_probas</span><span class="p">,</span>
           <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span>
             <span class="mf">0.0615</span><span class="p">,</span>
             <span class="mf">0.0000</span><span class="p">,</span>
             <span class="mf">0.0000</span><span class="p">,</span>
             <span class="mf">0.5775</span><span class="p">,</span>
             <span class="mf">0.0000</span><span class="p">,</span>
             <span class="mf">0.0000</span><span class="p">,</span>
             <span class="mf">0.0000</span><span class="p">,</span>
             <span class="mf">0.3610</span><span class="p">,</span>
             <span class="mf">0.0000</span>
           <span class="p">]),</span>
           <span class="ss">atol:</span> <span class="mf">1.0e-4</span>
         <span class="p">)</span>
<span class="k">end</span>
</code></pre></div></div>

<h3 id="533-modify-the-text-generation-function">5.3.3 Modify the text generation function</h3>

<p>We need a new <code class="language-plaintext highlighter-rouge">LlmScratch.TextGeneration.generate</code> function that generates the next token based on the model and current token while implementing temperature sampling and top-k sampling.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">TextGeneration</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Token generation helpers for GPT-style models.

  `generate_text_simple/4` mirrors the Python loop from the book:

      idx_cond = idx[:, -context_size:]
      logits = model(idx_cond)
      logits = logits[:, -1, :]
      probas = torch.softmax(logits, dim=-1)
      idx_next = torch.argmax(probas, dim=-1, keepdim=True)
      idx = torch.cat((idx, idx_next), dim=1)

  The model is expected to be a struct whose module exports `forward/2`.
  Existing modules such as `LlmScratch.GPTModel` and
  `LlmScratch.DummyGPTModel` satisfy that contract.
  """</span>

  <span class="nv">@default_sampling_seed</span> <span class="mi">123</span>

  <span class="nv">@spec</span> <span class="n">generate_text_simple</span><span class="p">(</span><span class="n">struct</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">non_neg_integer</span><span class="p">(),</span> <span class="n">pos_integer</span><span class="p">())</span> <span class="p">::</span>
          <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Generates new token ids with greedy decoding.

  ## Arguments

    * `model` - a GPT-style model struct whose module exports `forward/2`.
      `forward/2` receives token ids shaped `{batch_size, seq_len}` and must
      return logits shaped `{batch_size, seq_len, vocab_size}`.

    * `idx` - token ids shaped `{batch_size, seq_len}`. The returned tensor
      keeps the same batch size and token id type.

    * `max_new_tokens` - number of new tokens to generate and append. Use `0`
      to return `idx` unchanged.

    * `context_size` - maximum number of latest tokens to pass to the model on
      each generation step.

  ## Returns

  A token id tensor shaped `{batch_size, seq_len + max_new_tokens}`.

  The next token is selected with `argmax` over the final-position softmax
  probabilities, then appended to `idx` along the sequence axis.

  ## Examples

      generated =
        LlmScratch.TextGeneration.generate_text_simple(
          model,
          Nx.tensor([[6109, 3626]], type: {:s, 64}),
          5,
          1024
        )

      Nx.shape(generated)
      #=&gt; {1, 7}
   ## Algorithm:

      Start with prompt tokens.
      Keep only the last context_size tokens.
      Run the model.
      Take logits from the last position.
      Convert logits to probabilities.
      Pick the highest-probability token.
      Append it.
      Repeat max_new_tokens times.
  """</span>
  <span class="k">def</span> <span class="n">generate_text_simple</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">idx</span><span class="p">,</span> <span class="n">max_new_tokens</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">max_new_tokens</span><span class="p">)</span> <span class="ow">and</span> <span class="n">max_new_tokens</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">context_size</span><span class="p">)</span> <span class="ow">and</span>
             <span class="n">context_size</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="n">validate_idx_shape!</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">1</span><span class="o">..</span><span class="n">max_new_tokens</span><span class="o">//</span><span class="mi">1</span><span class="p">,</span> <span class="n">idx</span><span class="p">,</span> <span class="k">fn</span> <span class="n">_step</span><span class="p">,</span> <span class="n">acc</span> <span class="o">-&gt;</span>
      <span class="n">idx_cond</span> <span class="o">=</span> <span class="n">last_tokens</span><span class="p">(</span><span class="n">acc</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
      <span class="n">logits</span> <span class="o">=</span> <span class="n">forward!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">idx_cond</span><span class="p">)</span>
      <span class="n">logits</span> <span class="o">=</span> <span class="n">last_position_logits</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span>
      <span class="n">probas</span> <span class="o">=</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>

      <span class="n">idx_next</span> <span class="o">=</span>
        <span class="n">probas</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="ss">keep_axis:</span> <span class="no">true</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">as_type</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">acc</span><span class="p">))</span>

      <span class="no">Nx</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">acc</span><span class="p">,</span> <span class="n">idx_next</span><span class="p">],</span> <span class="ss">axis:</span> <span class="mi">1</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">generate</span><span class="p">(</span>
          <span class="n">struct</span><span class="p">(),</span>
          <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="n">non_neg_integer</span><span class="p">(),</span>
          <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="n">number</span><span class="p">(),</span>
          <span class="n">pos_integer</span><span class="p">()</span> <span class="o">|</span> <span class="no">nil</span><span class="p">,</span>
          <span class="n">integer</span><span class="p">()</span> <span class="o">|</span> <span class="no">nil</span><span class="p">,</span>
          <span class="n">keyword</span><span class="p">()</span>
        <span class="p">)</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Generates new token ids with greedy or temperature-scaled sampling.

  This mirrors the book's `generate` function:

      idx_cond = idx[:, -context_size:]
      logits = model(idx_cond)
      logits = logits[:, -1, :]

      if top_k is not None:
          top_logits, _ = torch.topk(logits, top_k)
          min_val = top_logits[:, -1]
          logits = torch.where(logits &lt; min_val, -inf, logits)

      if temperature &gt; 0.0:
          logits = logits / temperature
          probs = torch.softmax(logits, dim=-1)
          idx_next = torch.multinomial(probs, num_samples=1)
      else:
          idx_next = torch.argmax(logits, dim=-1, keepdim=True)

  ## Arguments

    * `model` - a GPT-style model struct whose module exports `forward/2`.

    * `idx` - token ids shaped `{batch_size, seq_len}`.

    * `max_new_tokens` - maximum number of new tokens to append.

    * `context_size` - maximum number of latest tokens to pass to the model.

    * `temperature` - `0.0` uses deterministic argmax decoding. Values greater
      than `0.0` divide logits by the temperature and sample from the softmax
      distribution.

    * `top_k` - when set, keeps only the `top_k` highest logits per batch row
      before decoding.

    * `eos_id` - when set, generation stops before appending an EOS token once
      all batch rows predict that id.

  ## Options

    * `:seed` - deterministic sampling seed used when `temperature &gt; 0.0`.
      Defaults to `#{@default_sampling_seed}`, matching the chapter examples.

  ## Examples

      token_ids =
        LlmScratch.TextGeneration.generate(
          model,
          LlmScratch.TextUtils.text_to_token_ids("Every effort moves you", "code-davinci-002"),
          15,
          model.cfg.context_length,
          1.4,
          25
        )

  ## Algorithm

    Start with prompt token ids.
    Keep only the latest context_size tokens.
    Run the model.
    Take logits for the last position.
    Optionally keep only top-k logits.
    If temperature == 0, pick the highest logit.
    If temperature &gt; 0, softmax and sample.
    Stop if EOS is predicted.
  """</span>
  <span class="k">def</span> <span class="n">generate</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">idx</span><span class="p">,</span>
        <span class="n">max_new_tokens</span><span class="p">,</span>
        <span class="n">context_size</span><span class="p">,</span>
        <span class="n">temperature</span> <span class="p">\\</span> <span class="mf">0.0</span><span class="p">,</span>
        <span class="n">top_k</span> <span class="p">\\</span> <span class="no">nil</span><span class="p">,</span>
        <span class="n">eos_id</span> <span class="p">\\</span> <span class="no">nil</span><span class="p">,</span>
        <span class="n">opts</span> <span class="p">\\</span> <span class="p">[]</span>
      <span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">max_new_tokens</span><span class="p">)</span> <span class="ow">and</span> <span class="n">max_new_tokens</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">context_size</span><span class="p">)</span> <span class="ow">and</span>
             <span class="n">context_size</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">is_number</span><span class="p">(</span><span class="n">temperature</span><span class="p">)</span> <span class="ow">and</span> <span class="n">temperature</span> <span class="o">&gt;=</span> <span class="mf">0.0</span> <span class="ow">and</span>
             <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">validate_idx_shape!</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>
    <span class="n">validate_top_k!</span><span class="p">(</span><span class="n">top_k</span><span class="p">)</span>

    <span class="n">seed</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">,</span> <span class="nv">@default_sampling_seed</span><span class="p">)</span>
    <span class="n">key</span> <span class="o">=</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span>

    <span class="p">{</span><span class="n">_key</span><span class="p">,</span> <span class="n">generated_idx</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">reduce_while</span><span class="p">(</span><span class="mi">1</span><span class="o">..</span><span class="n">max_new_tokens</span><span class="o">//</span><span class="mi">1</span><span class="p">,</span> <span class="p">{</span><span class="n">key</span><span class="p">,</span> <span class="n">idx</span><span class="p">},</span> <span class="k">fn</span> <span class="n">_step</span><span class="p">,</span> <span class="p">{</span><span class="n">key</span><span class="p">,</span> <span class="n">acc</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="n">idx_cond</span> <span class="o">=</span> <span class="n">last_tokens</span><span class="p">(</span><span class="n">acc</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
        <span class="n">logits</span> <span class="o">=</span> <span class="n">model</span> <span class="o">|&gt;</span> <span class="n">forward!</span><span class="p">(</span><span class="n">idx_cond</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="n">last_position_logits</span><span class="p">()</span> <span class="o">|&gt;</span> <span class="n">mask_top_k</span><span class="p">(</span><span class="n">top_k</span><span class="p">)</span>

        <span class="p">{</span><span class="n">idx_next</span><span class="p">,</span> <span class="n">key</span><span class="p">}</span> <span class="o">=</span>
          <span class="n">next_token</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">temperature</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span>

        <span class="n">idx_next</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">as_type</span><span class="p">(</span><span class="n">idx_next</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">acc</span><span class="p">))</span>

        <span class="k">if</span> <span class="n">eos_reached?</span><span class="p">(</span><span class="n">idx_next</span><span class="p">,</span> <span class="n">eos_id</span><span class="p">)</span> <span class="k">do</span>
          <span class="p">{</span><span class="ss">:halt</span><span class="p">,</span> <span class="p">{</span><span class="n">key</span><span class="p">,</span> <span class="n">acc</span><span class="p">}}</span>
        <span class="k">else</span>
          <span class="p">{</span><span class="ss">:cont</span><span class="p">,</span> <span class="p">{</span><span class="n">key</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">acc</span><span class="p">,</span> <span class="n">idx_next</span><span class="p">],</span> <span class="ss">axis:</span> <span class="mi">1</span><span class="p">)}}</span>
        <span class="k">end</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="n">generated_idx</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">last_tokens</span><span class="p">(</span><span class="n">idx</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>
    <span class="n">length</span> <span class="o">=</span> <span class="n">min</span><span class="p">(</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
    <span class="n">start</span> <span class="o">=</span> <span class="n">seq_len</span> <span class="o">-</span> <span class="n">length</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">slice_along_axis</span><span class="p">(</span><span class="n">idx</span><span class="p">,</span> <span class="n">start</span><span class="p">,</span> <span class="n">length</span><span class="p">,</span> <span class="ss">axis:</span> <span class="mi">1</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">last_position_logits</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">_vocab_size</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="n">logits</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">slice_along_axis</span><span class="p">(</span><span class="n">seq_len</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="ss">axis:</span> <span class="mi">1</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="ss">axes:</span> <span class="p">[</span><span class="mi">1</span><span class="p">])</span>

      <span class="n">shape</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
              <span class="s2">"expected logits shape {batch_size, seq_len, vocab_size}, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">mask_top_k</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">logits</span>

  <span class="k">defp</span> <span class="n">mask_top_k</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">top_k</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">top_k</span> <span class="o">&gt;</span> <span class="n">vocab_size</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"top_k must be &lt;= vocab size </span><span class="si">#{</span><span class="n">vocab_size</span><span class="si">}</span><span class="s2">, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">top_k</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>

    <span class="p">{</span><span class="n">top_logits</span><span class="p">,</span> <span class="n">_top_pos</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">top_k</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="ss">k:</span> <span class="n">top_k</span><span class="p">)</span>

    <span class="n">min_top_logits</span> <span class="o">=</span>
      <span class="n">top_logits</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">slice_along_axis</span><span class="p">(</span><span class="n">top_k</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">logits</span><span class="p">))</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">select</span><span class="p">(</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">min_top_logits</span><span class="p">),</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="ss">:neg_infinity</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">logits</span><span class="p">)),</span>
      <span class="n">logits</span>
    <span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">next_token</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">temperature</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span> <span class="ow">when</span> <span class="n">temperature</span> <span class="o">&gt;</span> <span class="mf">0.0</span> <span class="k">do</span>
    <span class="n">probas</span> <span class="o">=</span>
      <span class="n">logits</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">divide</span><span class="p">(</span><span class="n">temperature</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>

    <span class="n">sample_from_batch</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">next_token</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">_temperature</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="ss">keep_axis:</span> <span class="no">true</span><span class="p">),</span> <span class="n">key</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">sample_from_batch</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">probas</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="n">probas</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>
    <span class="p">{</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">vocab_size</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">probas</span><span class="p">)</span>

    <span class="p">{</span><span class="n">sampled_ids</span><span class="p">,</span> <span class="n">key</span><span class="p">}</span> <span class="o">=</span>
      <span class="n">probas</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_list</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map_reduce</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="k">fn</span> <span class="n">row_probas</span><span class="p">,</span> <span class="n">key</span> <span class="o">-&gt;</span>
        <span class="p">{</span><span class="n">sampled_token_ids</span><span class="p">,</span> <span class="n">key</span><span class="p">}</span> <span class="o">=</span>
          <span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span>
            <span class="n">key</span><span class="p">,</span>
            <span class="no">Nx</span><span class="o">.</span><span class="n">iota</span><span class="p">({</span><span class="n">vocab_size</span><span class="p">}),</span>
            <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">row_probas</span><span class="p">),</span>
            <span class="ss">samples:</span> <span class="mi">1</span>
          <span class="p">)</span>

        <span class="p">{[</span><span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">(</span><span class="n">sampled_token_ids</span><span class="p">[</span><span class="mi">0</span><span class="p">])],</span> <span class="n">key</span><span class="p">}</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="p">{</span><span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">sampled_ids</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">reshape</span><span class="p">({</span><span class="n">batch_size</span><span class="p">,</span> <span class="mi">1</span><span class="p">}),</span> <span class="n">key</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">eos_reached?</span><span class="p">(</span><span class="n">_idx_next</span><span class="p">,</span> <span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">false</span>

  <span class="k">defp</span> <span class="n">eos_reached?</span><span class="p">(</span><span class="n">idx_next</span><span class="p">,</span> <span class="n">eos_id</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">idx_next</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">eos_id</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.==</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">forward!</span><span class="p">(%</span><span class="n">module</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">,</span> <span class="n">idx_cond</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">if</span> <span class="n">function_exported?</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="ss">:forward</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span> <span class="k">do</span>
      <span class="n">module</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">idx_cond</span><span class="p">)</span>
    <span class="k">else</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">module</span><span class="p">)</span><span class="si">}</span><span class="s2"> to export forward/2"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_idx_shape!</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">_seq_len</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="ss">:ok</span>

      <span class="n">shape</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected idx shape {batch_size, seq_len}, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_top_k!</span><span class="p">(</span><span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="ss">:ok</span>

  <span class="k">defp</span> <span class="n">validate_top_k!</span><span class="p">(</span><span class="n">top_k</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">top_k</span><span class="p">)</span> <span class="ow">and</span> <span class="n">top_k</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="ss">:ok</span>

  <span class="k">defp</span> <span class="n">validate_top_k!</span><span class="p">(</span><span class="n">top_k</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"top_k must be a positive integer or nil, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">top_k</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And a test that uses this new <code class="language-plaintext highlighter-rouge">generate</code> function:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="n">test</span> <span class="s2">"5.3.3 generate uses temperature scaling and top-k sampling"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="s2">"ch5_2_gpt_124m.nx"</span>
    <span class="n">assert</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">load!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>

    <span class="n">token_ids</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="s2">"Every effort moves you"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">),</span>
        <span class="mi">15</span><span class="p">,</span>
        <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="mf">1.4</span><span class="p">,</span>
        <span class="mi">25</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">decoded_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">(</span><span class="n">token_ids</span><span class="p">)</span> <span class="o">==</span> <span class="p">[</span>
             <span class="mi">6109</span><span class="p">,</span>
             <span class="mi">3626</span><span class="p">,</span>
             <span class="mi">6100</span><span class="p">,</span>
             <span class="mi">345</span><span class="p">,</span>
             <span class="mi">7109</span><span class="p">,</span>
             <span class="mi">284</span><span class="p">,</span>
             <span class="mi">766</span><span class="p">,</span>
             <span class="mi">9074</span><span class="p">,</span>
             <span class="mi">13</span><span class="p">,</span>
             <span class="mi">198</span><span class="p">,</span>
             <span class="mi">40</span><span class="p">,</span>
             <span class="mi">329</span><span class="p">,</span>
             <span class="mi">340</span><span class="p">,</span>
             <span class="mi">547</span><span class="p">,</span>
             <span class="mi">11</span><span class="p">,</span>
             <span class="mi">290</span><span class="p">,</span>
             <span class="mi">9074</span><span class="p">,</span>
             <span class="mi">13</span><span class="p">,</span>
             <span class="mi">843</span>
           <span class="p">]</span>

    <span class="n">assert</span> <span class="n">decoded_text</span> <span class="o">==</span> <span class="sd">"""
           Every effort moves youdr to see Mrs.
           I for it were, and Mrs. And\
           """</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>We have another exercise from Sebastian:</p>

<p>Exercise 5.2</p>

<p><em>Play around with different temperatures and top-k settings. Based on your observa- tions, can you think of applications where lower temperature and top-k settings are desired? Likewise, can you think of applications where higher temperature and top-k settings are preferred? (It’s recommended to also revisit this exercise at the end of the chapter after loading the pretrained weights from OpenAI.)</em></p>

<p>Lower temperature and lower top_k are best when you want reliable, conservative output.</p>

<p>Examples:</p>

<ul>
  <li>Code generation where syntax correctness matters.</li>
  <li>Summarization where the model should stay close to the source text.</li>
  <li>Extracting facts from documents.</li>
  <li>Customer support answers where consistency matters.</li>
  <li>Classification-style generation, like choosing labels.</li>
  <li>Autocomplete where the next token should be predictable.</li>
  <li>Tests or demos where reproducibility is important.</li>
</ul>

<p>Higher temperature and higher top_k are better when you want variety, exploration, or creative alternatives.</p>

<p>Examples:</p>

<ul>
  <li>Creative writing.</li>
  <li>Brainstorming product names or story ideas.</li>
  <li>Dialogue generation for fictional characters.</li>
  <li>Poetry, jokes, slogans, or game content.</li>
  <li>Generating multiple diverse completions for a human to choose from.</li>
</ul>

<p>And we will repeat this exercise when we load the publicly available GPT-2 model.</p>

<p>Exercise 5.3</p>

<p><em>What are the different combinations of settings for the generate function to force deterministic behavior, that is, disabling the random sampling such that it always pro- duces the same outputs similar to the generate_simple function?</em></p>

<p>To force deterministic behavior, we should set temperature = 0.0. Doing that forces the greedy algorithm to be used in next-token generation.</p>

<h2 id="54-loading-and-saving-model-weights">5.4 Loading and saving model weights</h2>

<p>We have already implemented saving and loading model weights. But Sebastian also saves optimizer parameters, which is handy when we continue the training process. So we are adding functions to save and load the training state:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">ModelCheckpoint</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Saves and loads Nx model containers.

  This is intended for trained models such as `LlmScratch.GPTModel`, whose
  tensors are serializable through `Nx.Container`.
  """</span>

  <span class="nv">@type</span> <span class="n">path</span> <span class="p">::</span> <span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Serializes `model` and writes it to `path`.

  Parent directories are created automatically. The model must implement
  `Nx.Container`; `LlmScratch.GPTModel` does through its container
  implementation in `LlmScratch.Training`.

  ## Example

      checkpoint_path = "ch5_2_gpt_124m.nx"
      LlmScratch.ModelCheckpoint.save!(trained_model, checkpoint_path)

  """</span>
  <span class="nv">@spec</span> <span class="n">save!</span><span class="p">(</span><span class="n">struct</span><span class="p">(),</span> <span class="n">path</span><span class="p">())</span> <span class="p">::</span> <span class="ss">:ok</span>
  <span class="k">def</span> <span class="n">save!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">write_serialized!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Reads a checkpoint from `path` and deserializes it back into a model struct.

  The checkpoint must have been written with `save!/2` or another compatible
  `Nx.serialize/1` call.

  ## Example

      trained_model = LlmScratch.ModelCheckpoint.load!("ch5_2_gpt_124m.nx")

  """</span>
  <span class="nv">@spec</span> <span class="n">load!</span><span class="p">(</span><span class="n">path</span><span class="p">())</span> <span class="p">::</span> <span class="n">struct</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">load!</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">path</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">deserialize</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Serializes a training checkpoint containing both model and optimizer state.

  This mirrors the common PyTorch checkpoint shape:

      torch.save({
          "model_state_dict": model.state_dict(),
          "optimizer_state_dict": optimizer.state_dict(),
      }, "model_and_optimizer.pth")

  The returned Elixir checkpoint uses atom keys:

      %{
        model_state_dict: model,
        optimizer_state_dict: optimizer
      }

  Both values must implement `Nx.Container`. `LlmScratch.GPTModel` and
  `LlmScratch.Training.AdamW` do.

  ## Example

      LlmScratch.ModelCheckpoint.save_training_state!(
        trained_model,
        optimizer,
        "model_and_optimizer.nx"
      )

  """</span>
  <span class="nv">@spec</span> <span class="n">save_training_state!</span><span class="p">(</span><span class="n">struct</span><span class="p">(),</span> <span class="n">struct</span><span class="p">(),</span> <span class="n">path</span><span class="p">())</span> <span class="p">::</span> <span class="ss">:ok</span>
  <span class="k">def</span> <span class="n">save_training_state!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optimizer</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_struct</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">write_serialized!</span><span class="p">(</span>
      <span class="p">%{</span>
        <span class="ss">model_state_dict:</span> <span class="n">model</span><span class="p">,</span>
        <span class="ss">optimizer_state_dict:</span> <span class="n">optimizer</span>
      <span class="p">},</span>
      <span class="n">path</span>
    <span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Reads a training checkpoint written by `save_training_state!/3`.

  Returns:

      %{
        model_state_dict: model,
        optimizer_state_dict: optimizer
      }

  """</span>
  <span class="nv">@spec</span> <span class="n">load_training_state!</span><span class="p">(</span><span class="n">path</span><span class="p">())</span> <span class="p">::</span> <span class="p">%{</span>
          <span class="ss">model_state_dict:</span> <span class="n">struct</span><span class="p">(),</span>
          <span class="ss">optimizer_state_dict:</span> <span class="n">struct</span><span class="p">()</span>
        <span class="p">}</span>
  <span class="k">def</span> <span class="n">load_training_state!</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">path</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">deserialize</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">write_serialized!</span><span class="p">(</span><span class="n">term</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">path</span>
    <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">dirname</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">mkdir_p!</span><span class="p">()</span>

    <span class="no">File</span><span class="o">.</span><span class="n">write!</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">serialize</span><span class="p">(</span><span class="n">term</span><span class="p">))</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>Exercise 5.4</p>

<p><em>After saving the weights, load the model and optimizer in a new Python session or Jupyter notebook file and continue pretraining it for one more epoch using the train_model_simple function.</em></p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="nv">@tag</span> <span class="ss">:train</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">180_000</span>
  <span class="n">test</span> <span class="s2">"5.4 load model and optimizer checkpoint and continue pretraining"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="n">device</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="no">System</span><span class="o">.</span><span class="n">tmp_dir!</span><span class="p">(),</span> <span class="s2">"ch5_4_model_and_optimizer.nx"</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">rm</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>
    <span class="n">raw_text</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">duplicate</span><span class="p">(</span><span class="s2">"Every effort moves you forward. "</span><span class="p">,</span> <span class="mi">120</span><span class="p">)</span>

    <span class="n">gpt_config</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">16</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">32</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">4</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">1</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.0</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="no">GptDatasetV1</span><span class="o">.</span><span class="n">create_dataloader_v1</span><span class="p">(</span>
        <span class="ss">raw_text:</span> <span class="n">raw_text</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span>
        <span class="ss">max_length:</span> <span class="n">gpt_config</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">stride:</span> <span class="n">gpt_config</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="mi">0</span>
      <span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="no">GptDatasetV1</span><span class="o">.</span><span class="n">create_dataloader_v1</span><span class="p">(</span>
        <span class="ss">raw_text:</span> <span class="n">raw_text</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span>
        <span class="ss">max_length:</span> <span class="n">gpt_config</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">stride:</span> <span class="n">gpt_config</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="mi">0</span>
      <span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>
    <span class="n">optimizer</span> <span class="o">=</span> <span class="no">Training</span><span class="o">.</span><span class="n">adamw</span><span class="p">(</span><span class="mf">0.0004</span><span class="p">,</span> <span class="ss">weight_decay:</span> <span class="mf">0.1</span><span class="p">)</span>

    <span class="p">{</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">train_losses</span><span class="p">,</span> <span class="n">val_losses</span><span class="p">,</span> <span class="n">tokens_seen</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Training</span><span class="o">.</span><span class="n">train_model_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">train_loader</span><span class="p">,</span>
        <span class="n">val_loader</span><span class="p">,</span>
        <span class="n">optimizer</span><span class="p">,</span>
        <span class="n">device</span><span class="p">,</span>
        <span class="mi">1</span><span class="p">,</span>
        <span class="mi">2</span><span class="p">,</span>
        <span class="mi">1</span><span class="p">,</span>
        <span class="s2">"Every effort moves you"</span><span class="p">,</span>
        <span class="n">tokenizer</span><span class="p">,</span>
        <span class="ss">return_optimizer:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">generate_samples:</span> <span class="no">false</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="p">%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">trained_model</span>
    <span class="n">assert</span> <span class="p">%</span><span class="no">Training</span><span class="o">.</span><span class="no">AdamW</span><span class="p">{</span><span class="ss">m:</span> <span class="n">m</span><span class="p">,</span> <span class="ss">v:</span> <span class="n">v</span><span class="p">}</span> <span class="o">=</span> <span class="n">trained_optimizer</span>
    <span class="n">assert</span> <span class="n">trained_optimizer</span><span class="o">.</span><span class="n">step</span> <span class="o">==</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">length</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">m</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">v</span><span class="p">)</span> <span class="o">==</span> <span class="n">length</span><span class="p">(</span><span class="n">m</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">train_losses</span> <span class="o">++</span> <span class="n">val_losses</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_float</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">tokens_seen</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_integer</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>

    <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">save_training_state!</span><span class="p">(</span><span class="n">trained_model</span><span class="p">,</span> <span class="n">trained_optimizer</span><span class="p">,</span> <span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="p">%{</span>
      <span class="ss">model_state_dict:</span> <span class="n">loaded_model</span><span class="p">,</span>
      <span class="ss">optimizer_state_dict:</span> <span class="n">loaded_optimizer</span>
    <span class="p">}</span> <span class="o">=</span> <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">load_training_state!</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span>

    <span class="n">assert</span> <span class="p">%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">loaded_model</span>
    <span class="n">assert</span> <span class="p">%</span><span class="no">Training</span><span class="o">.</span><span class="no">AdamW</span><span class="p">{}</span> <span class="o">=</span> <span class="n">loaded_optimizer</span>
    <span class="n">assert</span> <span class="n">loaded_optimizer</span><span class="o">.</span><span class="n">step</span> <span class="o">==</span> <span class="n">trained_optimizer</span><span class="o">.</span><span class="n">step</span>

    <span class="p">{</span><span class="n">_continued_model</span><span class="p">,</span> <span class="n">continued_optimizer</span><span class="p">,</span> <span class="n">continued_train_losses</span><span class="p">,</span> <span class="n">continued_val_losses</span><span class="p">,</span>
     <span class="n">continued_tokens_seen</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Training</span><span class="o">.</span><span class="n">train_model_simple</span><span class="p">(</span>
        <span class="n">loaded_model</span><span class="p">,</span>
        <span class="n">train_loader</span><span class="p">,</span>
        <span class="n">val_loader</span><span class="p">,</span>
        <span class="n">loaded_optimizer</span><span class="p">,</span>
        <span class="n">device</span><span class="p">,</span>
        <span class="mi">1</span><span class="p">,</span>
        <span class="mi">2</span><span class="p">,</span>
        <span class="mi">1</span><span class="p">,</span>
        <span class="s2">"Every effort moves you"</span><span class="p">,</span>
        <span class="n">tokenizer</span><span class="p">,</span>
        <span class="ss">return_optimizer:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">generate_samples:</span> <span class="no">false</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="n">continued_optimizer</span><span class="o">.</span><span class="n">step</span> <span class="o">==</span> <span class="n">trained_optimizer</span><span class="o">.</span><span class="n">step</span> <span class="o">+</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">length</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">continued_train_losses</span> <span class="o">++</span> <span class="n">continued_val_losses</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_float</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">assert</span> <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="n">continued_tokens_seen</span><span class="p">,</span> <span class="o">&amp;</span><span class="n">is_integer</span><span class="o">/</span><span class="mi">1</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="55-loading-pretrained-weights-from-openai">5.5 Loading pretrained weights from OpenAI</h2>

<p>Training our GPT-2 model on the short <em>The Verdict</em> story took 5 minutes on my MacBook Pro M3. Still, our validation loss was quite significant, so the model was not usable. OpenAI shared its GPT-2 model weights publicly. The weights are saved using TensorFlow, so we will also use Pythonx integration.</p>

<p>The 124M parameter model takes 900 MB, and the 1558M model takes 11 GB of disk space. Here is a module that downloads and loads weights into the GPTModel module:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">GPT2OpenAI</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Downloads and loads the public OpenAI GPT-2 checkpoints.

  The module is the bridge between OpenAI's original TensorFlow checkpoint
  layout and this repo's `%LlmScratch.GPTModel{}` structs. It handles three
  separate steps:

    * downloading the original checkpoint files
    * converting TensorFlow checkpoint tensors into `Nx.Tensor`s
    * assigning those tensors into the local GPT model implementation

  This mirrors the book's Python helper:

      settings, params = download_and_load_gpt2(
          model_size="124M",
          models_dir="gpt2"
      )

  The checkpoint files are downloaded from OpenAI's public GPT-2 storage and
  cached under `models_dir/model_size`. TensorFlow is used through `Pythonx` to
  read the original TensorFlow checkpoint and export its variables to a cached
  `model_params.npz` file in the same model directory.

  Supported checkpoint sizes are:

    * `"124M"`
    * `"355M"`
    * `"774M"`
    * `"1558M"`

  Python packages required by the checkpoint reader:

      pip install "tensorflow&gt;=2.15.0" "tqdm&gt;=4.66"

  At runtime, this Elixir module needs network access for the initial download
  and a Python environment with TensorFlow and NumPy for the checkpoint export.
  Subsequent loads reuse the downloaded files and cached `.npz` export when
  present.
  """</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span><span class="no">DummyLayerNorm</span><span class="p">,</span> <span class="no">EmbeddingNative</span><span class="p">,</span> <span class="no">GPTConfig</span><span class="p">,</span> <span class="no">GPTModel</span><span class="p">,</span> <span class="no">ModelCheckpoint</span><span class="p">}</span>

  <span class="nv">@allowed_sizes</span> <span class="p">[</span><span class="s2">"124M"</span><span class="p">,</span> <span class="s2">"355M"</span><span class="p">,</span> <span class="s2">"774M"</span><span class="p">,</span> <span class="s2">"1558M"</span><span class="p">]</span>
  <span class="nv">@filenames</span> <span class="p">[</span>
    <span class="s2">"checkpoint"</span><span class="p">,</span>
    <span class="s2">"encoder.json"</span><span class="p">,</span>
    <span class="s2">"hparams.json"</span><span class="p">,</span>
    <span class="s2">"model.ckpt.data-00000-of-00001"</span><span class="p">,</span>
    <span class="s2">"model.ckpt.index"</span><span class="p">,</span>
    <span class="s2">"model.ckpt.meta"</span><span class="p">,</span>
    <span class="s2">"vocab.bpe"</span>
  <span class="p">]</span>
  <span class="nv">@base_url</span> <span class="s2">"https://openaipublic.blob.core.windows.net/gpt-2/models"</span>
  <span class="nv">@backup_base_url</span> <span class="s2">"https://f001.backblazeb2.com/file/LLMs-from-scratch/gpt2"</span>

  <span class="nv">@type</span> <span class="n">settings</span> <span class="p">::</span> <span class="n">map</span><span class="p">()</span>
  <span class="nv">@type</span> <span class="n">params</span> <span class="p">::</span> <span class="n">map</span><span class="p">()</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Downloads GPT-2 files and loads TensorFlow checkpoint params.

  `model_size` must be one of `"124M"`, `"355M"`, `"774M"`, or `"1558M"`.

  Options:

    * `:models_dir` - directory used to cache downloaded checkpoints. Defaults
      to `"gpt2"`.

  Returns `{settings, params}`. `settings` is decoded from `hparams.json`.
  `params` follows the book's shape:

      %{
        "wte" =&gt; token_embedding,
        "wpe" =&gt; position_embedding,
        "g" =&gt; final_norm_scale,
        "b" =&gt; final_norm_shift,
        "blocks" =&gt; [
          %{
            "attn" =&gt; %{"c_attn" =&gt; %{"w" =&gt; ..., "b" =&gt; ...}, ...},
            "mlp" =&gt; ...,
            "ln_1" =&gt; ...,
            "ln_2" =&gt; ...
          }
        ]
      }

  The tensor values are `Nx.Tensor`s.

  Raises if the model size is unsupported, a checkpoint file cannot be
  downloaded, TensorFlow cannot export the checkpoint, or a required tensor is
  missing from the exported `.npz` file.
  """</span>
  <span class="nv">@spec</span> <span class="n">download_and_load_gpt2</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="p">{</span><span class="n">settings</span><span class="p">(),</span> <span class="n">params</span><span class="p">()}</span>
  <span class="k">def</span> <span class="n">download_and_load_gpt2</span><span class="p">(</span><span class="n">model_size</span> <span class="p">\\</span> <span class="s2">"124M"</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">models_dir</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:models_dir</span><span class="p">,</span> <span class="s2">"gpt2"</span><span class="p">)</span>
    <span class="n">model_dir</span> <span class="o">=</span> <span class="n">download_gpt2_files!</span><span class="p">(</span><span class="n">model_size</span><span class="p">,</span> <span class="n">models_dir</span><span class="p">)</span>
    <span class="n">settings</span> <span class="o">=</span> <span class="n">load_settings!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span>
    <span class="n">params</span> <span class="o">=</span> <span class="n">load_params_from_tf_checkpoint!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">,</span> <span class="n">settings</span><span class="p">)</span>

    <span class="p">{</span><span class="n">settings</span><span class="p">,</span> <span class="n">params</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Downloads GPT-2 files and returns the local model directory.

  If all expected files already exist and are non-empty, no download is
  attempted. Otherwise each checkpoint file is downloaded from OpenAI's public
  GPT-2 storage, with the book's Backblaze mirror used as a fallback.

  The returned directory is `Path.join(models_dir, model_size)`.
  """</span>
  <span class="nv">@spec</span> <span class="n">download_gpt2_files!</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">download_gpt2_files!</span><span class="p">(</span><span class="n">model_size</span><span class="p">,</span> <span class="n">models_dir</span> <span class="p">\\</span> <span class="s2">"gpt2"</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">models_dir</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">validate_model_size!</span><span class="p">(</span><span class="n">model_size</span><span class="p">)</span>

    <span class="n">model_dir</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">models_dir</span><span class="p">,</span> <span class="n">model_size</span><span class="p">)</span>
    <span class="no">File</span><span class="o">.</span><span class="n">mkdir_p!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span>

    <span class="k">unless</span> <span class="n">all_checkpoint_files_present?</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span> <span class="k">do</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">each</span><span class="p">(</span><span class="nv">@filenames</span><span class="p">,</span> <span class="k">fn</span> <span class="n">filename</span> <span class="o">-&gt;</span>
        <span class="n">destination</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">model_dir</span><span class="p">,</span> <span class="n">filename</span><span class="p">)</span>
        <span class="n">primary_url</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nv">@base_url</span><span class="p">,</span> <span class="n">model_size</span><span class="p">,</span> <span class="n">filename</span><span class="p">],</span> <span class="s2">"/"</span><span class="p">)</span>
        <span class="n">backup_url</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nv">@backup_base_url</span><span class="p">,</span> <span class="n">model_size</span><span class="p">,</span> <span class="n">filename</span><span class="p">],</span> <span class="s2">"/"</span><span class="p">)</span>

        <span class="n">download_file!</span><span class="p">(</span><span class="n">primary_url</span><span class="p">,</span> <span class="n">destination</span><span class="p">,</span> <span class="n">backup_url</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>
    <span class="k">end</span>

    <span class="n">model_dir</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Builds a `%LlmScratch.GPTModel{}` with the public GPT-2 weights assigned.

  This is the highest-level loading helper when you want a ready-to-run model:
  it downloads missing checkpoint files, exports the TensorFlow checkpoint,
  creates the matching `%LlmScratch.GPTConfig{}`, initializes a local model, and
  assigns the downloaded weights.

  Options:

    * `:models_dir` - checkpoint cache directory. Defaults to `"gpt2"`.
    * `:seed` - seed used for the temporary model initialization before
      checkpoint weights are assigned. Defaults to `123`.

  The returned model uses `norm_eps: 1.0e-5`, matching GPT-2's layer norm
  epsilon.
  """</span>
  <span class="nv">@spec</span> <span class="n">load_model</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">load_model</span><span class="p">(</span><span class="n">model_size</span> <span class="p">\\</span> <span class="s2">"124M"</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">settings</span><span class="p">,</span> <span class="n">params</span><span class="p">}</span> <span class="o">=</span> <span class="n">download_and_load_gpt2</span><span class="p">(</span><span class="n">model_size</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="n">settings</span>
    <span class="o">|&gt;</span> <span class="n">config_from_settings</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">seed:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">,</span> <span class="mi">123</span><span class="p">),</span> <span class="ss">norm_eps:</span> <span class="mf">1.0e-5</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">load_weights_into_gpt</span><span class="p">(</span><span class="n">params</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Downloads, loads, assigns, and saves public GPT-2 weights as an Nx checkpoint.

  Returns the loaded `%LlmScratch.GPTModel{}`.

  This is useful when you want to pay the TensorFlow checkpoint conversion cost
  once and later load the model through `LlmScratch.ModelCheckpoint.load!/1`.
  Parent directories for `checkpoint_path` are created by
  `LlmScratch.ModelCheckpoint.save!/2`.
  """</span>
  <span class="nv">@spec</span> <span class="n">save_model!</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">save_model!</span><span class="p">(</span><span class="n">model_size</span> <span class="p">\\</span> <span class="s2">"124M"</span><span class="p">,</span> <span class="n">checkpoint_path</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span>
      <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">checkpoint_path</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">model</span> <span class="o">=</span> <span class="n">load_model</span><span class="p">(</span><span class="n">model_size</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="no">ModelCheckpoint</span><span class="o">.</span><span class="n">save!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">checkpoint_path</span><span class="p">)</span>
    <span class="n">model</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Validates a GPT-2 source tensor against a target tensor and returns the source.

  `left` is the tensor already present in the local model. `right` is the
  tensor loaded from the OpenAI checkpoint. The function raises when their
  shapes differ, then casts `right` to `left`'s type before returning it.

  This is the Nx equivalent of the book's PyTorch helper:

      def assign(left, right):
          if left.shape != right.shape:
              raise ValueError(...)
          return torch.nn.Parameter(torch.tensor(right))

  Elixir data is immutable and this repo stores parameters directly as
  `Nx.Tensor`s, so the function returns `right` converted to the target tensor's
  type after validating that the shapes match.
  """</span>
  <span class="nv">@spec</span> <span class="n">assign</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">assign</span><span class="p">(%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">left</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">right</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">left_shape</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">left</span><span class="p">)</span>
    <span class="n">right_shape</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">right</span><span class="p">)</span>

    <span class="k">unless</span> <span class="n">left_shape</span> <span class="o">==</span> <span class="n">right_shape</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
            <span class="s2">"shape mismatch. Left: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">left_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">, Right: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">right_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">as_type</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">left</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Loads decoded OpenAI GPT-2 params into this repo's `%LlmScratch.GPTModel{}`.

  This mirrors the book's `load_weights_into_gpt(gpt, params)` function:

    * token and positional embeddings are copied from `"wte"` and `"wpe"`
    * Q/K/V weights and biases are split from `"attn"["c_attn"]`
    * attention output, feed-forward, and layer norm parameters are assigned
    * the output head is tied to token embeddings

  One layout difference from the PyTorch code is intentional: this repo's dense
  kernels are stored as `{in_dim, out_dim}`, matching the OpenAI TensorFlow
  checkpoint matrices. PyTorch `nn.Linear.weight` is `{out_dim, in_dim}`, so the
  Python book code uses `.T` where this Nx implementation does not.
  """</span>
  <span class="nv">@spec</span> <span class="n">load_weights_into_gpt</span><span class="p">(</span><span class="no">GPTModel</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">params</span><span class="p">())</span> <span class="p">::</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">load_weights_into_gpt</span><span class="p">(%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">,</span> <span class="n">params</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">params</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">do_assign_params!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">params</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Backward-compatible alias for `load_weights_into_gpt/2`.

  Prefer `load_weights_into_gpt/2` in new code.
  """</span>
  <span class="nv">@spec</span> <span class="n">assign_params!</span><span class="p">(</span><span class="no">GPTModel</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">params</span><span class="p">())</span> <span class="p">::</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">assign_params!</span><span class="p">(%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">,</span> <span class="n">params</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">params</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">load_weights_into_gpt</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">params</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">do_assign_params!</span><span class="p">(%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">,</span> <span class="n">params</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">params</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">blocks</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="s2">"blocks"</span><span class="p">)</span>

    <span class="k">unless</span> <span class="n">length</span><span class="p">(</span><span class="n">blocks</span><span class="p">)</span> <span class="o">==</span> <span class="n">length</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">)</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
            <span class="s2">"expected </span><span class="si">#{</span><span class="n">length</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">)</span><span class="si">}</span><span class="s2"> GPT-2 blocks, got </span><span class="si">#{</span><span class="n">length</span><span class="p">(</span><span class="n">blocks</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>

    <span class="n">token_embedding</span> <span class="o">=</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="s2">"wte"</span><span class="p">)</span>

    <span class="p">%{</span>
      <span class="n">model</span>
      <span class="o">|</span> <span class="ss">tok_emb:</span> <span class="n">put_embedding_weight!</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">tok_emb</span><span class="p">,</span> <span class="n">token_embedding</span><span class="p">),</span>
        <span class="ss">pos_emb:</span> <span class="n">put_embedding_weight!</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span><span class="p">,</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="s2">"wpe"</span><span class="p">)),</span>
        <span class="ss">final_norm:</span>
          <span class="n">put_layer_norm!</span><span class="p">(</span>
            <span class="n">model</span><span class="o">.</span><span class="n">final_norm</span><span class="p">,</span>
            <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="s2">"g"</span><span class="p">),</span>
            <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="s2">"b"</span><span class="p">)</span>
          <span class="p">),</span>
        <span class="ss">out_head:</span> <span class="n">put_output_head!</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="p">,</span> <span class="n">token_embedding</span><span class="p">),</span>
        <span class="ss">trf_blocks:</span>
          <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">,</span> <span class="n">blocks</span><span class="p">)</span>
          <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">block</span><span class="p">,</span> <span class="n">block_params</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="n">put_block!</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">block_params</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Converts GPT-2 `hparams.json` settings into this repo's `%GPTConfig{}`.

  The OpenAI settings determine the model size from the triple
  `{n_embd, n_layer, n_head}`. The returned config keeps the checkpoint's
  vocabulary size and context length, and uses GPT-2's `qkv_bias: true`
  setting.
  """</span>
  <span class="nv">@spec</span> <span class="n">config_from_settings</span><span class="p">(</span><span class="n">settings</span><span class="p">())</span> <span class="p">::</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">config_from_settings</span><span class="p">(</span><span class="n">settings</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">settings</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span>
      <span class="no">GPTConfig</span><span class="o">.</span><span class="n">openai_gpt2</span><span class="p">(</span><span class="n">model_size_from_settings!</span><span class="p">(</span><span class="n">settings</span><span class="p">))</span>
      <span class="o">|</span> <span class="ss">vocab_size:</span> <span class="n">fetch_setting!</span><span class="p">(</span><span class="n">settings</span><span class="p">,</span> <span class="s2">"n_vocab"</span><span class="p">),</span>
        <span class="ss">context_length:</span> <span class="n">fetch_setting!</span><span class="p">(</span><span class="n">settings</span><span class="p">,</span> <span class="s2">"n_ctx"</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Loads `hparams.json` from a downloaded GPT-2 model directory.

  `model_dir` should be the directory returned by `download_gpt2_files!/2`.
  """</span>
  <span class="nv">@spec</span> <span class="n">load_settings!</span><span class="p">(</span><span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">settings</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">load_settings!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">model_dir</span>
    <span class="o">|&gt;</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">"hparams.json"</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Jason</span><span class="o">.</span><span class="n">decode!</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Loads GPT-2 params from a downloaded TensorFlow checkpoint directory.

  This exports the TensorFlow checkpoint to `model_params.npz` if that cache
  file is missing, then reads the expected GPT-2 variables from the `.npz` file
  into the nested params map consumed by `load_weights_into_gpt/2`.
  """</span>
  <span class="nv">@spec</span> <span class="n">load_params_from_tf_checkpoint!</span><span class="p">(</span><span class="no">Path</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">settings</span><span class="p">())</span> <span class="p">::</span> <span class="n">params</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">load_params_from_tf_checkpoint!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">,</span> <span class="n">settings</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_map</span><span class="p">(</span><span class="n">settings</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">npz_path</span> <span class="o">=</span> <span class="n">export_tf_checkpoint_to_npz!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span>

    <span class="p">%{</span>
      <span class="s2">"blocks"</span> <span class="o">=&gt;</span>
        <span class="mi">0</span><span class="o">..</span><span class="p">(</span><span class="n">fetch_setting!</span><span class="p">(</span><span class="n">settings</span><span class="p">,</span> <span class="s2">"n_layer"</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="n">load_block_params!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="nv">&amp;1</span><span class="p">)),</span>
      <span class="s2">"wte"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"model/wte"</span><span class="p">),</span>
      <span class="s2">"wpe"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"model/wpe"</span><span class="p">),</span>
      <span class="s2">"g"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"model/ln_f/g"</span><span class="p">),</span>
      <span class="s2">"b"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"model/ln_f/b"</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">load_block_params!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="n">layer_idx</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">prefix</span> <span class="o">=</span> <span class="s2">"model/h</span><span class="si">#{</span><span class="n">layer_idx</span><span class="si">}</span><span class="s2">"</span>

    <span class="p">%{</span>
      <span class="s2">"attn"</span> <span class="o">=&gt;</span> <span class="p">%{</span>
        <span class="s2">"c_attn"</span> <span class="o">=&gt;</span> <span class="p">%{</span>
          <span class="s2">"w"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/attn/c_attn/w"</span><span class="p">),</span>
          <span class="s2">"b"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/attn/c_attn/b"</span><span class="p">)</span>
        <span class="p">},</span>
        <span class="s2">"c_proj"</span> <span class="o">=&gt;</span> <span class="p">%{</span>
          <span class="s2">"w"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/attn/c_proj/w"</span><span class="p">),</span>
          <span class="s2">"b"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/attn/c_proj/b"</span><span class="p">)</span>
        <span class="p">}</span>
      <span class="p">},</span>
      <span class="s2">"mlp"</span> <span class="o">=&gt;</span> <span class="p">%{</span>
        <span class="s2">"c_fc"</span> <span class="o">=&gt;</span> <span class="p">%{</span>
          <span class="s2">"w"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/mlp/c_fc/w"</span><span class="p">),</span>
          <span class="s2">"b"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/mlp/c_fc/b"</span><span class="p">)</span>
        <span class="p">},</span>
        <span class="s2">"c_proj"</span> <span class="o">=&gt;</span> <span class="p">%{</span>
          <span class="s2">"w"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/mlp/c_proj/w"</span><span class="p">),</span>
          <span class="s2">"b"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/mlp/c_proj/b"</span><span class="p">)</span>
        <span class="p">}</span>
      <span class="p">},</span>
      <span class="s2">"ln_1"</span> <span class="o">=&gt;</span> <span class="p">%{</span>
        <span class="s2">"g"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/ln_1/g"</span><span class="p">),</span>
        <span class="s2">"b"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/ln_1/b"</span><span class="p">)</span>
      <span class="p">},</span>
      <span class="s2">"ln_2"</span> <span class="o">=&gt;</span> <span class="p">%{</span>
        <span class="s2">"g"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/ln_2/g"</span><span class="p">),</span>
        <span class="s2">"b"</span> <span class="o">=&gt;</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="s2">"</span><span class="si">#{</span><span class="n">prefix</span><span class="si">}</span><span class="s2">/ln_2/b"</span><span class="p">)</span>
      <span class="p">}</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">load_tensor!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">,</span> <span class="n">name</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">python_code</span> <span class="o">=</span> <span class="sd">"""
    import numpy as np

    if isinstance(npz_path, bytes):
        npz_path = npz_path.decode()

    if isinstance(name, bytes):
        name = name.decode()

    with np.load(npz_path) as data:
        array = data[name].astype(np.float32)

    (list(array.shape), array.tobytes(order="C"))
    """</span>

    <span class="p">{</span><span class="n">result</span><span class="p">,</span> <span class="n">_globals</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">Pythonx</span><span class="o">.</span><span class="n">eval</span><span class="p">(</span><span class="n">python_code</span><span class="p">,</span> <span class="p">%{</span>
        <span class="s2">"npz_path"</span> <span class="o">=&gt;</span> <span class="n">npz_path</span><span class="p">,</span>
        <span class="s2">"name"</span> <span class="o">=&gt;</span> <span class="n">name</span>
      <span class="p">})</span>

    <span class="p">{</span><span class="n">shape</span><span class="p">,</span> <span class="n">bytes</span><span class="p">}</span> <span class="o">=</span> <span class="no">Pythonx</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="n">result</span><span class="p">)</span>

    <span class="n">bytes</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">from_binary</span><span class="p">({</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="no">List</span><span class="o">.</span><span class="n">to_tuple</span><span class="p">(</span><span class="n">shape</span><span class="p">))</span>
  <span class="k">rescue</span>
    <span class="n">error</span> <span class="ow">in</span> <span class="p">[</span><span class="no">Pythonx</span><span class="o">.</span><span class="no">Error</span><span class="p">]</span> <span class="o">-&gt;</span>
      <span class="k">raise</span> <span class="sd">"""
      Failed to load GPT-2 variable #{inspect(name)} from #{inspect(npz_path)}.

      Make sure NumPy is available to Pythonx.

      Original error:

      #{Exception.message(error)}
      """</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">export_tf_checkpoint_to_npz!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">npz_path</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">model_dir</span><span class="p">,</span> <span class="s2">"model_params.npz"</span><span class="p">)</span>

    <span class="k">if</span> <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">npz_path</span><span class="p">)</span> <span class="ow">and</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">npz_path</span><span class="p">)</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
      <span class="n">npz_path</span>
    <span class="k">else</span>
      <span class="n">ckpt_path</span> <span class="o">=</span> <span class="n">latest_checkpoint!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span>
      <span class="n">python</span> <span class="o">=</span> <span class="n">python_executable!</span><span class="p">()</span>

      <span class="n">script</span> <span class="o">=</span> <span class="sd">"""
      import sys
      import numpy as np
      import tensorflow as tf

      ckpt_path = sys.argv[1]
      npz_path = sys.argv[2]

      params = {}
      for name, _shape in tf.train.list_variables(ckpt_path):
          if name.startswith("model/"):
              params[name] = np.squeeze(tf.train.load_variable(ckpt_path, name)).astype(np.float32)

      np.savez_compressed(npz_path, **params)
      """</span>

      <span class="k">case</span> <span class="no">System</span><span class="o">.</span><span class="n">cmd</span><span class="p">(</span><span class="n">python</span><span class="p">,</span> <span class="p">[</span><span class="s2">"-c"</span><span class="p">,</span> <span class="n">script</span><span class="p">,</span> <span class="n">ckpt_path</span><span class="p">,</span> <span class="n">npz_path</span><span class="p">],</span>
             <span class="ss">stderr_to_stdout:</span> <span class="no">true</span><span class="p">,</span>
             <span class="ss">env:</span> <span class="p">[{</span><span class="s2">"TF_CPP_MIN_LOG_LEVEL"</span><span class="p">,</span> <span class="s2">"2"</span><span class="p">}]</span>
           <span class="p">)</span> <span class="k">do</span>
        <span class="p">{</span><span class="n">_output</span><span class="p">,</span> <span class="mi">0</span><span class="p">}</span> <span class="o">-&gt;</span>
          <span class="n">npz_path</span>

        <span class="p">{</span><span class="n">output</span><span class="p">,</span> <span class="n">status</span><span class="p">}</span> <span class="o">-&gt;</span>
          <span class="no">File</span><span class="o">.</span><span class="n">rm</span><span class="p">(</span><span class="n">npz_path</span><span class="p">)</span>

          <span class="k">raise</span> <span class="sd">"""
          Failed to export TensorFlow checkpoint #{inspect(ckpt_path)} to #{inspect(npz_path)}.

          Make sure TensorFlow is available in the Python environment:

              pip install "tensorflow&gt;=2.15.0" "tqdm&gt;=4.66"

          Python exited with status #{status}.

          #{output}
          """</span>
      <span class="k">end</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">latest_checkpoint!</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">model_dir</span><span class="p">,</span> <span class="s2">"checkpoint"</span><span class="p">)</span>

    <span class="n">checkpoint_file</span> <span class="o">=</span>
      <span class="n">checkpoint_path</span>
      <span class="o">|&gt;</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="p">,</span> <span class="ss">trim:</span> <span class="no">true</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="o">&amp;</span><span class="no">String</span><span class="o">.</span><span class="n">starts_with?</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="s2">"model_checkpoint_path:"</span><span class="p">))</span>

    <span class="k">case</span> <span class="n">checkpoint_file</span> <span class="k">do</span>
      <span class="no">nil</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"no TensorFlow checkpoint found in </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>

      <span class="n">line</span> <span class="o">-&gt;</span>
        <span class="p">[</span><span class="n">_key</span><span class="p">,</span> <span class="n">quoted_path</span><span class="p">]</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">line</span><span class="p">,</span> <span class="s2">":"</span><span class="p">,</span> <span class="ss">parts:</span> <span class="mi">2</span><span class="p">)</span>

        <span class="n">ckpt_path</span> <span class="o">=</span>
          <span class="n">quoted_path</span>
          <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">trim</span><span class="p">()</span>
          <span class="o">|&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">trim</span><span class="p">(</span><span class="s2">"</span><span class="se">\"</span><span class="s2">"</span><span class="p">)</span>

        <span class="k">if</span> <span class="no">Path</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">ckpt_path</span><span class="p">)</span> <span class="o">==</span> <span class="ss">:absolute</span> <span class="k">do</span>
          <span class="n">ckpt_path</span>
        <span class="k">else</span>
          <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">model_dir</span><span class="p">,</span> <span class="n">ckpt_path</span><span class="p">)</span>
        <span class="k">end</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">python_executable!</span> <span class="k">do</span>
    <span class="n">pythonx_python</span> <span class="o">=</span>
      <span class="ss">:pythonx</span>
      <span class="o">|&gt;</span> <span class="no">Application</span><span class="o">.</span><span class="n">app_dir</span><span class="p">(</span><span class="s2">"priv/uv/project/.venv/bin/python"</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="n">to_string</span><span class="p">()</span>

    <span class="k">cond</span> <span class="k">do</span>
      <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">pythonx_python</span><span class="p">)</span> <span class="o">-&gt;</span>
        <span class="n">pythonx_python</span>

      <span class="n">python</span> <span class="o">=</span> <span class="no">System</span><span class="o">.</span><span class="n">find_executable</span><span class="p">(</span><span class="s2">"python3"</span><span class="p">)</span> <span class="o">-&gt;</span>
        <span class="n">python</span>

      <span class="n">python</span> <span class="o">=</span> <span class="no">System</span><span class="o">.</span><span class="n">find_executable</span><span class="p">(</span><span class="s2">"python"</span><span class="p">)</span> <span class="o">-&gt;</span>
        <span class="n">python</span>

      <span class="no">true</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="s2">"could not find a Python executable for TensorFlow checkpoint export"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_block!</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">block_params</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">q_w</span><span class="p">,</span> <span class="n">k_w</span><span class="p">,</span> <span class="n">v_w</span><span class="p">}</span> <span class="o">=</span> <span class="n">split_qkv</span><span class="p">(</span><span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">block_params</span><span class="p">[</span><span class="s2">"attn"</span><span class="p">][</span><span class="s2">"c_attn"</span><span class="p">],</span> <span class="s2">"w"</span><span class="p">))</span>
    <span class="p">{</span><span class="n">q_b</span><span class="p">,</span> <span class="n">k_b</span><span class="p">,</span> <span class="n">v_b</span><span class="p">}</span> <span class="o">=</span> <span class="n">split_qkv</span><span class="p">(</span><span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">block_params</span><span class="p">[</span><span class="s2">"attn"</span><span class="p">][</span><span class="s2">"c_attn"</span><span class="p">],</span> <span class="s2">"b"</span><span class="p">))</span>

    <span class="p">%{</span>
      <span class="n">block</span>
      <span class="o">|</span> <span class="ss">norm1:</span>
          <span class="n">put_layer_norm!</span><span class="p">(</span>
            <span class="n">block</span><span class="o">.</span><span class="n">norm1</span><span class="p">,</span>
            <span class="n">block_params</span><span class="p">[</span><span class="s2">"ln_1"</span><span class="p">][</span><span class="s2">"g"</span><span class="p">],</span>
            <span class="n">block_params</span><span class="p">[</span><span class="s2">"ln_1"</span><span class="p">][</span><span class="s2">"b"</span><span class="p">]</span>
          <span class="p">),</span>
        <span class="ss">norm2:</span>
          <span class="n">put_layer_norm!</span><span class="p">(</span>
            <span class="n">block</span><span class="o">.</span><span class="n">norm2</span><span class="p">,</span>
            <span class="n">block_params</span><span class="p">[</span><span class="s2">"ln_2"</span><span class="p">][</span><span class="s2">"g"</span><span class="p">],</span>
            <span class="n">block_params</span><span class="p">[</span><span class="s2">"ln_2"</span><span class="p">][</span><span class="s2">"b"</span><span class="p">]</span>
          <span class="p">),</span>
        <span class="ss">att:</span> <span class="p">%{</span>
          <span class="n">block</span><span class="o">.</span><span class="n">att</span>
          <span class="o">|</span> <span class="ss">w_q:</span> <span class="n">put_dense!</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">w_q</span><span class="p">,</span> <span class="n">q_w</span><span class="p">,</span> <span class="n">q_b</span><span class="p">),</span>
            <span class="ss">w_k:</span> <span class="n">put_dense!</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">w_k</span><span class="p">,</span> <span class="n">k_w</span><span class="p">,</span> <span class="n">k_b</span><span class="p">),</span>
            <span class="ss">w_v:</span> <span class="n">put_dense!</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">w_v</span><span class="p">,</span> <span class="n">v_w</span><span class="p">,</span> <span class="n">v_b</span><span class="p">),</span>
            <span class="ss">out_proj:</span>
              <span class="n">put_dense!</span><span class="p">(</span>
                <span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">out_proj</span><span class="p">,</span>
                <span class="n">block_params</span><span class="p">[</span><span class="s2">"attn"</span><span class="p">][</span><span class="s2">"c_proj"</span><span class="p">][</span><span class="s2">"w"</span><span class="p">],</span>
                <span class="n">block_params</span><span class="p">[</span><span class="s2">"attn"</span><span class="p">][</span><span class="s2">"c_proj"</span><span class="p">][</span><span class="s2">"b"</span><span class="p">]</span>
              <span class="p">)</span>
        <span class="p">},</span>
        <span class="ss">ff:</span> <span class="p">%{</span>
          <span class="n">block</span><span class="o">.</span><span class="n">ff</span>
          <span class="o">|</span> <span class="ss">layers:</span> <span class="p">%{</span>
              <span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="o">.</span><span class="n">layers</span>
              <span class="o">|</span> <span class="ss">first:</span>
                  <span class="n">put_dense!</span><span class="p">(</span>
                    <span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="p">,</span>
                    <span class="n">block_params</span><span class="p">[</span><span class="s2">"mlp"</span><span class="p">][</span><span class="s2">"c_fc"</span><span class="p">][</span><span class="s2">"w"</span><span class="p">],</span>
                    <span class="n">block_params</span><span class="p">[</span><span class="s2">"mlp"</span><span class="p">][</span><span class="s2">"c_fc"</span><span class="p">][</span><span class="s2">"b"</span><span class="p">]</span>
                  <span class="p">),</span>
                <span class="ss">second:</span>
                  <span class="n">put_dense!</span><span class="p">(</span>
                    <span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="p">,</span>
                    <span class="n">block_params</span><span class="p">[</span><span class="s2">"mlp"</span><span class="p">][</span><span class="s2">"c_proj"</span><span class="p">][</span><span class="s2">"w"</span><span class="p">],</span>
                    <span class="n">block_params</span><span class="p">[</span><span class="s2">"mlp"</span><span class="p">][</span><span class="s2">"c_proj"</span><span class="p">][</span><span class="s2">"b"</span><span class="p">]</span>
                  <span class="p">)</span>
            <span class="p">}</span>
        <span class="p">}</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_embedding_weight!</span><span class="p">(%</span><span class="no">EmbeddingNative</span><span class="p">{}</span> <span class="o">=</span> <span class="n">embedding</span><span class="p">,</span> <span class="n">weight</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span><span class="n">embedding</span> <span class="o">|</span> <span class="ss">weight:</span> <span class="n">assign</span><span class="p">(</span><span class="n">embedding</span><span class="o">.</span><span class="n">weight</span><span class="p">,</span> <span class="n">weight</span><span class="p">)}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_layer_norm!</span><span class="p">(%</span><span class="no">DummyLayerNorm</span><span class="p">{}</span> <span class="o">=</span> <span class="n">layer_norm</span><span class="p">,</span> <span class="n">scale</span><span class="p">,</span> <span class="n">shift</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span><span class="n">layer_norm</span> <span class="o">|</span> <span class="ss">scale:</span> <span class="n">assign</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">scale</span><span class="p">,</span> <span class="n">scale</span><span class="p">),</span> <span class="ss">shift:</span> <span class="n">assign</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">shift</span><span class="p">,</span> <span class="n">shift</span><span class="p">)}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_dense!</span><span class="p">(</span><span class="n">dense</span><span class="p">,</span> <span class="n">kernel</span><span class="p">,</span> <span class="n">bias</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span><span class="n">dense</span> <span class="o">|</span> <span class="ss">kernel:</span> <span class="n">assign</span><span class="p">(</span><span class="n">dense</span><span class="o">.</span><span class="n">kernel</span><span class="p">,</span> <span class="n">kernel</span><span class="p">),</span> <span class="ss">bias:</span> <span class="n">assign</span><span class="p">(</span><span class="n">dense</span><span class="o">.</span><span class="n">bias</span><span class="p">,</span> <span class="n">bias</span><span class="p">)}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">put_output_head!</span><span class="p">(</span><span class="n">out_head</span><span class="p">,</span> <span class="n">token_embedding</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">kernel</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="n">token_embedding</span><span class="p">)</span>
    <span class="p">%{</span><span class="n">out_head</span> <span class="o">|</span> <span class="ss">kernel:</span> <span class="n">assign</span><span class="p">(</span><span class="n">out_head</span><span class="o">.</span><span class="n">kernel</span><span class="p">,</span> <span class="n">kernel</span><span class="p">)}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">split_qkv</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">shape</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>
    <span class="n">last_dim</span> <span class="o">=</span> <span class="n">shape</span> <span class="o">|&gt;</span> <span class="no">Tuple</span><span class="o">.</span><span class="n">to_list</span><span class="p">()</span> <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">last</span><span class="p">()</span>

    <span class="k">unless</span> <span class="n">rem</span><span class="p">(</span><span class="n">last_dim</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
            <span class="s2">"expected QKV tensor last dim to be divisible by 3, got </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>

    <span class="n">width</span> <span class="o">=</span> <span class="n">div</span><span class="p">(</span><span class="n">last_dim</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>

    <span class="p">{</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">slice_along_axis</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">width</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">slice_along_axis</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">width</span><span class="p">,</span> <span class="n">width</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">slice_along_axis</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">width</span><span class="p">,</span> <span class="n">width</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">download_file!</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="n">destination</span><span class="p">,</span> <span class="n">backup_url</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">if</span> <span class="n">file_current?</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="n">destination</span><span class="p">)</span> <span class="k">do</span>
      <span class="ss">:ok</span>
    <span class="k">else</span>
      <span class="n">url</span>
      <span class="o">|&gt;</span> <span class="n">do_download_file</span><span class="p">(</span><span class="n">destination</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="k">case</span> <span class="k">do</span>
        <span class="ss">:ok</span> <span class="o">-&gt;</span> <span class="ss">:ok</span>
        <span class="p">{</span><span class="ss">:error</span><span class="p">,</span> <span class="n">primary_error</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="n">download_backup_file!</span><span class="p">(</span><span class="n">backup_url</span><span class="p">,</span> <span class="n">destination</span><span class="p">,</span> <span class="n">primary_error</span><span class="p">)</span>
      <span class="k">end</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">download_backup_file!</span><span class="p">(</span><span class="n">backup_url</span><span class="p">,</span> <span class="n">destination</span><span class="p">,</span> <span class="n">primary_error</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="n">do_download_file</span><span class="p">(</span><span class="n">backup_url</span><span class="p">,</span> <span class="n">destination</span><span class="p">)</span> <span class="k">do</span>
      <span class="ss">:ok</span> <span class="o">-&gt;</span>
        <span class="ss">:ok</span>

      <span class="p">{</span><span class="ss">:error</span><span class="p">,</span> <span class="n">backup_error</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="sd">"""
        Failed to download GPT-2 checkpoint file #{Path.basename(destination)}.

        Primary error: #{inspect(primary_error)}
        Backup error: #{inspect(backup_error)}
        """</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">do_download_file</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="n">destination</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">tmp_destination</span> <span class="o">=</span> <span class="n">destination</span> <span class="o">&lt;&gt;</span> <span class="s2">".download"</span>
    <span class="no">File</span><span class="o">.</span><span class="n">rm</span><span class="p">(</span><span class="n">tmp_destination</span><span class="p">)</span>

    <span class="n">response</span> <span class="o">=</span>
      <span class="no">Req</span><span class="o">.</span><span class="n">get!</span><span class="p">(</span>
        <span class="n">url</span><span class="p">,</span>
        <span class="ss">into:</span> <span class="no">File</span><span class="o">.</span><span class="n">stream!</span><span class="p">(</span><span class="n">tmp_destination</span><span class="p">),</span>
        <span class="ss">receive_timeout:</span> <span class="mi">120_000</span><span class="p">,</span>
        <span class="ss">retry:</span> <span class="no">false</span>
      <span class="p">)</span>

    <span class="k">if</span> <span class="n">response</span><span class="o">.</span><span class="n">status</span> <span class="ow">in</span> <span class="mi">200</span><span class="o">..</span><span class="mi">299</span> <span class="k">do</span>
      <span class="no">File</span><span class="o">.</span><span class="n">rename!</span><span class="p">(</span><span class="n">tmp_destination</span><span class="p">,</span> <span class="n">destination</span><span class="p">)</span>
      <span class="ss">:ok</span>
    <span class="k">else</span>
      <span class="no">File</span><span class="o">.</span><span class="n">rm</span><span class="p">(</span><span class="n">tmp_destination</span><span class="p">)</span>
      <span class="p">{</span><span class="ss">:error</span><span class="p">,</span> <span class="p">{</span><span class="ss">:http_status</span><span class="p">,</span> <span class="n">response</span><span class="o">.</span><span class="n">status</span><span class="p">,</span> <span class="n">url</span><span class="p">}}</span>
    <span class="k">end</span>
  <span class="k">rescue</span>
    <span class="n">error</span> <span class="o">-&gt;</span>
      <span class="no">File</span><span class="o">.</span><span class="n">rm</span><span class="p">(</span><span class="n">destination</span> <span class="o">&lt;&gt;</span> <span class="s2">".download"</span><span class="p">)</span>
      <span class="p">{</span><span class="ss">:error</span><span class="p">,</span> <span class="n">error</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">file_current?</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="n">destination</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">destination</span><span class="p">)</span> <span class="ow">and</span> <span class="n">remote_content_length</span><span class="p">(</span><span class="n">url</span><span class="p">)</span> <span class="o">==</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">destination</span><span class="p">)</span><span class="o">.</span><span class="n">size</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">all_checkpoint_files_present?</span><span class="p">(</span><span class="n">model_dir</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Enum</span><span class="o">.</span><span class="n">all?</span><span class="p">(</span><span class="nv">@filenames</span><span class="p">,</span> <span class="k">fn</span> <span class="n">filename</span> <span class="o">-&gt;</span>
      <span class="n">path</span> <span class="o">=</span> <span class="no">Path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">model_dir</span><span class="p">,</span> <span class="n">filename</span><span class="p">)</span>
      <span class="no">File</span><span class="o">.</span><span class="n">exists?</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="ow">and</span> <span class="no">File</span><span class="o">.</span><span class="n">stat!</span><span class="p">(</span><span class="n">path</span><span class="p">)</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">remote_content_length</span><span class="p">(</span><span class="n">url</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">response</span> <span class="o">=</span> <span class="no">Req</span><span class="o">.</span><span class="n">head!</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="ss">receive_timeout:</span> <span class="mi">30_000</span><span class="p">,</span> <span class="ss">retry:</span> <span class="no">false</span><span class="p">)</span>

    <span class="n">response</span>
    <span class="o">|&gt;</span> <span class="no">Req</span><span class="o">.</span><span class="no">Response</span><span class="o">.</span><span class="n">get_header</span><span class="p">(</span><span class="s2">"content-length"</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">first</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="k">case</span> <span class="k">do</span>
      <span class="no">nil</span> <span class="o">-&gt;</span> <span class="ss">:unknown</span>
      <span class="n">size</span> <span class="o">-&gt;</span> <span class="no">String</span><span class="o">.</span><span class="n">to_integer</span><span class="p">(</span><span class="n">size</span><span class="p">)</span>
    <span class="k">end</span>
  <span class="k">rescue</span>
    <span class="n">_error</span> <span class="o">-&gt;</span> <span class="ss">:unknown</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_model_size!</span><span class="p">(</span><span class="n">model_size</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">unless</span> <span class="n">model_size</span> <span class="ow">in</span> <span class="nv">@allowed_sizes</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
            <span class="s2">"model_size must be one of </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="nv">@allowed_sizes</span><span class="p">)</span><span class="si">}</span><span class="s2">, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">model_size</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">fetch_setting!</span><span class="p">(</span><span class="n">settings</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Map</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">settings</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">model_size_from_settings!</span><span class="p">(</span><span class="n">settings</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="p">{</span><span class="n">fetch_setting!</span><span class="p">(</span><span class="n">settings</span><span class="p">,</span> <span class="s2">"n_embd"</span><span class="p">),</span> <span class="n">fetch_setting!</span><span class="p">(</span><span class="n">settings</span><span class="p">,</span> <span class="s2">"n_layer"</span><span class="p">),</span>
          <span class="n">fetch_setting!</span><span class="p">(</span><span class="n">settings</span><span class="p">,</span> <span class="s2">"n_head"</span><span class="p">)}</span> <span class="k">do</span>
      <span class="p">{</span><span class="mi">768</span><span class="p">,</span> <span class="mi">12</span><span class="p">,</span> <span class="mi">12</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="s2">"124M"</span>
      <span class="p">{</span><span class="mi">1024</span><span class="p">,</span> <span class="mi">24</span><span class="p">,</span> <span class="mi">16</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="s2">"355M"</span>
      <span class="p">{</span><span class="mi">1280</span><span class="p">,</span> <span class="mi">36</span><span class="p">,</span> <span class="mi">20</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="s2">"774M"</span>
      <span class="p">{</span><span class="mi">1600</span><span class="p">,</span> <span class="mi">48</span><span class="p">,</span> <span class="mi">25</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="s2">"1558M"</span>
      <span class="n">other</span> <span class="o">-&gt;</span> <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"unsupported GPT-2 settings dimensions: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">other</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>We also added a new function in <code class="language-plaintext highlighter-rouge">LlmScratch.GPTConfig</code> that helps us initialize GPTModel from OpenAI parameters:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nv">@openai_gpt2_model_configs</span> <span class="p">%{</span>
    <span class="s2">"gpt2-small (124M)"</span> <span class="o">=&gt;</span> <span class="p">%{</span><span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span> <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span> <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">},</span>
    <span class="s2">"gpt2-medium (355M)"</span> <span class="o">=&gt;</span> <span class="p">%{</span><span class="ss">emb_dim:</span> <span class="mi">1024</span><span class="p">,</span> <span class="ss">n_layers:</span> <span class="mi">24</span><span class="p">,</span> <span class="ss">n_heads:</span> <span class="mi">16</span><span class="p">},</span>
    <span class="s2">"gpt2-large (774M)"</span> <span class="o">=&gt;</span> <span class="p">%{</span><span class="ss">emb_dim:</span> <span class="mi">1280</span><span class="p">,</span> <span class="ss">n_layers:</span> <span class="mi">36</span><span class="p">,</span> <span class="ss">n_heads:</span> <span class="mi">20</span><span class="p">},</span>
    <span class="s2">"gpt2-xl (1558M)"</span> <span class="o">=&gt;</span> <span class="p">%{</span><span class="ss">emb_dim:</span> <span class="mi">1600</span><span class="p">,</span> <span class="ss">n_layers:</span> <span class="mi">48</span><span class="p">,</span> <span class="ss">n_heads:</span> <span class="mi">25</span><span class="p">},</span>
    <span class="s2">"124M"</span> <span class="o">=&gt;</span> <span class="p">%{</span><span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span> <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span> <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">},</span>
    <span class="s2">"355M"</span> <span class="o">=&gt;</span> <span class="p">%{</span><span class="ss">emb_dim:</span> <span class="mi">1024</span><span class="p">,</span> <span class="ss">n_layers:</span> <span class="mi">24</span><span class="p">,</span> <span class="ss">n_heads:</span> <span class="mi">16</span><span class="p">},</span>
    <span class="s2">"774M"</span> <span class="o">=&gt;</span> <span class="p">%{</span><span class="ss">emb_dim:</span> <span class="mi">1280</span><span class="p">,</span> <span class="ss">n_layers:</span> <span class="mi">36</span><span class="p">,</span> <span class="ss">n_heads:</span> <span class="mi">20</span><span class="p">},</span>
    <span class="s2">"1558M"</span> <span class="o">=&gt;</span> <span class="p">%{</span><span class="ss">emb_dim:</span> <span class="mi">1600</span><span class="p">,</span> <span class="ss">n_layers:</span> <span class="mi">48</span><span class="p">,</span> <span class="ss">n_heads:</span> <span class="mi">25</span><span class="p">}</span>
  <span class="p">}</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns a GPT-2 config matching OpenAI's released checkpoint sizes.

  This mirrors the book's update pattern:

      NEW_CONFIG = GPT_CONFIG_124M.copy()
      NEW_CONFIG.update(model_configs[model_name])
      NEW_CONFIG.update({"context_length": 1024})
      NEW_CONFIG.update({"qkv_bias": True})

  Accepted names are the descriptive labels from the book and the short model
  sizes used by the downloader:

    * `"gpt2-small (124M)"` or `"124M"`
    * `"gpt2-medium (355M)"` or `"355M"`
    * `"gpt2-large (774M)"` or `"774M"`
    * `"gpt2-xl (1558M)"` or `"1558M"`
  """</span>
  <span class="nv">@spec</span> <span class="n">openai_gpt2</span><span class="p">(</span><span class="no">String</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">openai_gpt2</span><span class="p">(</span><span class="n">model_name</span> <span class="p">\\</span> <span class="s2">"gpt2-small (124M)"</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_binary</span><span class="p">(</span><span class="n">model_name</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Map</span><span class="o">.</span><span class="n">fetch</span><span class="p">(</span><span class="nv">@openai_gpt2_model_configs</span><span class="p">,</span> <span class="n">model_name</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">model_config</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
          <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
          <span class="ss">emb_dim:</span> <span class="n">model_config</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span>
          <span class="ss">n_heads:</span> <span class="n">model_config</span><span class="o">.</span><span class="n">n_heads</span><span class="p">,</span>
          <span class="ss">n_layers:</span> <span class="n">model_config</span><span class="o">.</span><span class="n">n_layers</span><span class="p">,</span>
          <span class="ss">drop_rate:</span> <span class="mf">0.0</span><span class="p">,</span>
          <span class="ss">qkv_bias:</span> <span class="no">true</span>
        <span class="p">}</span>

      <span class="ss">:error</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
              <span class="s2">"unknown OpenAI GPT-2 model </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">model_name</span><span class="p">)</span><span class="si">}</span><span class="s2">; expected one of </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="no">Map</span><span class="o">.</span><span class="n">keys</span><span class="p">(</span><span class="nv">@openai_gpt2_model_configs</span><span class="p">))</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>And a test that loads GPT-2 params and predicts new tokens:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
  <span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">900_000</span>
  <span class="n">test</span> <span class="s2">"5.5 downloads and loads public OpenAI GPT-2 124M settings and params"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="p">{</span><span class="n">settings</span><span class="p">,</span> <span class="n">params</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">download_and_load_gpt2</span><span class="p">(</span>
        <span class="s2">"124M"</span><span class="p">,</span>
        <span class="ss">models_dir:</span> <span class="s2">"gpt2"</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="n">settings</span> <span class="o">==</span> <span class="p">%{</span>
             <span class="s2">"n_vocab"</span> <span class="o">=&gt;</span> <span class="mi">50_257</span><span class="p">,</span>
             <span class="s2">"n_ctx"</span> <span class="o">=&gt;</span> <span class="mi">1024</span><span class="p">,</span>
             <span class="s2">"n_embd"</span> <span class="o">=&gt;</span> <span class="mi">768</span><span class="p">,</span>
             <span class="s2">"n_head"</span> <span class="o">=&gt;</span> <span class="mi">12</span><span class="p">,</span>
             <span class="s2">"n_layer"</span> <span class="o">=&gt;</span> <span class="mi">12</span>
           <span class="p">}</span>

    <span class="n">assert</span> <span class="no">Map</span><span class="o">.</span><span class="n">keys</span><span class="p">(</span><span class="n">params</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">sort</span><span class="p">()</span> <span class="o">==</span> <span class="p">[</span><span class="s2">"b"</span><span class="p">,</span> <span class="s2">"blocks"</span><span class="p">,</span> <span class="s2">"g"</span><span class="p">,</span> <span class="s2">"wpe"</span><span class="p">,</span> <span class="s2">"wte"</span><span class="p">]</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">params</span><span class="p">[</span><span class="s2">"wte"</span><span class="p">])</span> <span class="o">==</span> <span class="p">{</span><span class="mi">50257</span><span class="p">,</span> <span class="mi">768</span><span class="p">}</span>

    <span class="n">model</span> <span class="o">=</span>
      <span class="n">settings</span>
      <span class="o">|&gt;</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">config_from_settings</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="ss">seed:</span> <span class="mi">123</span><span class="p">,</span> <span class="ss">norm_eps:</span> <span class="mf">1.0e-5</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_weights_into_gpt</span><span class="p">(</span><span class="n">params</span><span class="p">)</span>

    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>

    <span class="n">token_ids</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="s2">"Every effort moves you"</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">),</span>
        <span class="mi">25</span><span class="p">,</span>
        <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="mf">1.5</span><span class="p">,</span>
        <span class="mi">50</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>

    <span class="n">decoded_text</span> <span class="o">=</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">token_ids</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">decoded_text</span> <span class="o">==</span>
             <span class="s2">"Every effort moves you closer to victory!</span><span class="se">\"</span><span class="s2"> he told the story of his friend. But if all else failed, he was determined to defeat, his"</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Finally, we got tokens that make some sense! And we did not burn a single dollar on model training.</p>

<p>Exercise 5.5</p>

<p><em>Calculate the training and validation set losses of the GPTModel with the pretrained weights from OpenAI on the “The Verdict” dataset.</em></p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">900_000</span>
  <span class="n">test</span> <span class="s2">"5.5 exercise calculates The Verdict losses with OpenAI GPT-2 124M weights"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="s2">"124M"</span><span class="p">,</span> <span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>

    <span class="n">file_content</span> <span class="o">=</span> <span class="no">File</span><span class="o">.</span><span class="n">read!</span><span class="p">(</span><span class="s2">"the-verdict.txt"</span><span class="p">)</span>
    <span class="n">train_ratio</span> <span class="o">=</span> <span class="mf">0.90</span>
    <span class="n">split_idx</span> <span class="o">=</span> <span class="n">trunc</span><span class="p">(</span><span class="n">train_ratio</span> <span class="o">*</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">file_content</span><span class="p">))</span>
    <span class="n">train_data</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">file_content</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">split_idx</span><span class="p">)</span>
    <span class="n">val_data</span> <span class="o">=</span> <span class="no">String</span><span class="o">.</span><span class="n">slice</span><span class="p">(</span><span class="n">file_content</span><span class="p">,</span> <span class="n">split_idx</span><span class="p">,</span> <span class="no">String</span><span class="o">.</span><span class="n">length</span><span class="p">(</span><span class="n">file_content</span><span class="p">)</span> <span class="o">-</span> <span class="n">split_idx</span><span class="p">)</span>

    <span class="n">train_loader</span> <span class="o">=</span>
      <span class="no">GptDatasetV1</span><span class="o">.</span><span class="n">create_dataloader_v1</span><span class="p">(</span>
        <span class="ss">raw_text:</span> <span class="n">train_data</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span>
        <span class="ss">max_length:</span> <span class="mi">256</span><span class="p">,</span>
        <span class="ss">stride:</span> <span class="mi">256</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">true</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="mi">0</span>
      <span class="p">)</span>

    <span class="n">val_loader</span> <span class="o">=</span>
      <span class="no">GptDatasetV1</span><span class="o">.</span><span class="n">create_dataloader_v1</span><span class="p">(</span>
        <span class="ss">raw_text:</span> <span class="n">val_data</span><span class="p">,</span>
        <span class="ss">batch_size:</span> <span class="mi">2</span><span class="p">,</span>
        <span class="ss">max_length:</span> <span class="mi">256</span><span class="p">,</span>
        <span class="ss">stride:</span> <span class="mi">256</span><span class="p">,</span>
        <span class="ss">drop_last:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">shuffle:</span> <span class="no">false</span><span class="p">,</span>
        <span class="ss">num_workers:</span> <span class="mi">0</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">==</span> <span class="mi">9</span>
    <span class="n">assert</span> <span class="n">val_loader</span><span class="o">.</span><span class="n">length</span> <span class="o">==</span> <span class="mi">1</span>

    <span class="n">train_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">val_loss</span> <span class="o">=</span> <span class="no">LossUtils</span><span class="o">.</span><span class="n">calc_loss_loader</span><span class="p">(</span><span class="n">val_loader</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>

    <span class="n">assert_in_delta</span> <span class="n">train_loss</span><span class="p">,</span> <span class="mf">3.7547634177737765</span><span class="p">,</span> <span class="mf">1.0e-5</span>
    <span class="n">assert_in_delta</span> <span class="n">val_loss</span><span class="p">,</span> <span class="mf">3.5596354007720947</span><span class="p">,</span> <span class="mf">1.0e-5</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>We can see that validation loss is almost the same as training loss and much better than with our model trained on <em>The Verdict</em> story.</p>

<p>Exercise 5.6</p>

<p><em>Experiment with GPT-2 models of different sizes—for example, the largest 1,558 mil- lion parameter model—and compare the generated text to the 124 million model.</em></p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="nv">@tag</span> <span class="ss">:download</span>
  <span class="nv">@tag</span> <span class="ss">timeout:</span> <span class="mi">3_600_000</span>
  <span class="n">test</span> <span class="s2">"5.6 exercise compares generated text from GPT-2 124M and 1558M"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">prompt</span> <span class="o">=</span> <span class="s2">"Every effort moves you"</span>
    <span class="n">tokenizer</span> <span class="o">=</span> <span class="s2">"code-davinci-002"</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">GPT2OpenAI</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="s2">"1558M"</span><span class="p">,</span> <span class="ss">models_dir:</span> <span class="s2">"gpt2"</span><span class="p">)</span>

    <span class="n">text</span> <span class="o">=</span>
      <span class="n">model</span>
      <span class="o">|&gt;</span> <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate</span><span class="p">(</span>
        <span class="no">TextUtils</span><span class="o">.</span><span class="n">text_to_token_ids</span><span class="p">(</span><span class="n">prompt</span><span class="p">,</span> <span class="n">tokenizer</span><span class="p">),</span>
        <span class="mi">25</span><span class="p">,</span>
        <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
        <span class="mf">1.5</span><span class="p">,</span>
        <span class="mi">50</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">backend_transfer</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">BinaryBackend</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">TextUtils</span><span class="o">.</span><span class="n">token_ids_to_text</span><span class="p">(</span><span class="n">tokenizer</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">text</span> <span class="o">==</span> <span class="s2">"Every effort moves you farther and farther closer to it. I don't need to try all their tricks.</span><span class="se">\n\n\n</span><span class="s2">(END OF TRANSCRIPT"</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>The predicted text is different from the text generated by the 124M model.</p>

<h2 id="conclusion">Conclusion</h2>

<p>We managed to train our model, and we saw that this requires significant computational power. Luckily, we can use an existing pretrained model that we can fine-tune. This is our next chapter.</p>]]></content><author><name></name></author><category term="llm-from-scratch" /><category term="gpt" /><category term="tutorials" /><category term="llm" /><category term="gpt" /><category term="workbook" /><summary type="html"><![CDATA[A practical Elixir walkthrough of GPT pretraining: calculating generation loss, training on The Verdict, controlling decoding with temperature and top-k sampling, saving checkpoints, and loading OpenAI GPT-2 weights.]]></summary></entry><entry><title type="html">Estimating the population of Zabok using public data</title><link href="https://karlosmid.com/2026/05/estimating-the-number-of-population-for-zabok-using-public-data/" rel="alternate" type="text/html" title="Estimating the population of Zabok using public data" /><published>2026-05-14T09:00:00+00:00</published><updated>2026-05-14T09:00:00+00:00</updated><id>https://karlosmid.com/2026/05/estimating-the-number-of-population-for-zabok-using-public-data</id><content type="html" xml:base="https://karlosmid.com/2026/05/estimating-the-number-of-population-for-zabok-using-public-data/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>The last public population census in Croatia was done in 2021. With the latest development in Zabok, I made a simple estimate of the current population of Zabok.</p>

<h2 id="public-data">Public data</h2>

<p>I am using public data from the Croatian Bureau of Statistics and the Ministry of Justice, Public Administration and Digital Transformation:</p>
<ul>
  <li><a href="https://podaci.dzs.hr/2025/hr/97601">statistics on the natural increase of population</a></li>
  <li><a href="https://mpudt.gov.hr/UserDocsImages/dokumenti/Lokalni%20izbori%202025/Predstavni%C4%8Dka%20tijela%20-%20broj%20bira%C4%8Da%20po%20op%C4%87inama%20i%20gradovima.pdf">number of registered voters on May 9, 2025, for the 2025 Croatian local elections</a></li>
  <li><a href="https://podaci.dzs.hr/media/rqybclnx/popis_2021-stanovnistvo_po_naseljima.xlsx">Croatian 2021 population census</a></li>
</ul>

<p>In 2021, Zabok had a population of 8656. The number of voters on May 9, 2025 was 7563. This estimate assumes that the voter register is a good proxy for the adult resident population. The estimated percentage of minors is 17.3%, which gives us an estimated Zabok population on May 9, 2025 of 7563 / 0.827 = 9145.</p>

<p>The difference between the 2021 population and this estimate is 489. Zabok is losing 47 residents per year based on births minus deaths. From the end of 2021 to May 9, 2025, this gives an estimated natural decrease of about 158 residents.</p>

<p>So, the estimated population change from the end of 2021 to May 9, 2025 is +489 residents, or 489 / 8656 = 0.056, which is 5.6%. To reach that result despite the estimated natural decrease, Zabok would need an estimated gain of about 647 residents from other sources, such as migration.</p>]]></content><author><name></name></author><category term="statistics" /><category term="testing" /><category term="statistics" /><category term="testing" /><summary type="html"><![CDATA[The last public population census in Croatia was done in 2021. With the latest development in Zabok, I made a simple estimate of the current population of Zabok.]]></summary></entry><entry><title type="html">Recap on implementing GPT model from scratch</title><link href="https://karlosmid.com/2026/05/recap-on-implementing-GPT-model-from-scratch/" rel="alternate" type="text/html" title="Recap on implementing GPT model from scratch" /><published>2026-05-08T09:00:00+00:00</published><updated>2026-05-08T09:00:00+00:00</updated><id>https://karlosmid.com/2026/05/recap-on-implementing-GPT-model-from-scratch</id><content type="html" xml:base="https://karlosmid.com/2026/05/recap-on-implementing-GPT-model-from-scratch/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>I am recaping blog post <a href="/implementing-a-GPT-model-from-scratch">Implementing a GPT model from scratch</a> by answering workbook questions Build llm from scratch chapter 4, and reading Giles blog posts on this chapter:</p>
<ul>
  <li><a href="https://www.gilesthomas.com/2025/05/llm-from-scratch-15-from-context-vectors-to-logits">part15</a></li>
  <li><a href="https://www.gilesthomas.com/2025/07/llm-from-scratch-16-layer-normalisation">part16</a></li>
  <li><a href="https://www.gilesthomas.com/2025/08/llm-from-scratch-17-the-feed-forward-network">part17</a></li>
  <li><a href="https://www.gilesthomas.com/2025/08/llm-from-scratch-18-residuals-shortcut-connections-and-the-talmud">part18</a></li>
  <li><a href="https://www.gilesthomas.com/2025/08/llm-from-scratch-19-wrapping-up-chapter-4">part19</a></li>
</ul>

<h2 id="giles-blog-posts">Giles’ blog posts</h2>

<p>I got lost during Giles’ part 15 on the first read. The second read made things clearer. He goes deep into the GPTModel architecture and how data flows from tokens to output prediction. This is much deeper than the book content. He explains the weight-tying trick with examples. For a given token, the logits tensor gives us probabilities for the next token across the entire vocabulary.</p>

<p>Part 16 is about layer normalization. Giles goes deeper than the book in explaining why we want token embedding vectors to have a mean of 0 and a variance of 1. He explains the problem of vanishing and exploding gradients. A vanishing gradient means that we did not move our calculation for the next token in any direction. An exploding gradient is the opposite: one gradient makes all the movement in its direction.</p>

<p>In part 17, Giles hunts for the why behind the feed-forward block. It seems that the feed-forward layer is doing additional thinking, and he uses the diamond-in-a-cup riddle to show that. It is a very interesting example.</p>

<p>Part 18 is about shortcut connections, which tackle the problem of vanishing gradients. In shortcut connections, we add the input and output of a layer. Giles goes deeper into what that actually means, since we are changing the network. He also uses the Talmud, a central Jewish text, to provide an example of shortcut connections. Again, Sebastian’s book is about how, not why, so Giles’ posts are useful side reading.</p>

<p>Part 19 is about three posts that explain how LLMs work, based on Giles’ understanding of the material covered up to chapter 5 of the book. One part is about math for LLMs, with a very good explanation of which math we need for LLMs. I strongly recommend reading all three posts to deepen your knowledge of how LLMs work.</p>

<h2 id="book-github-resources">Book github resources</h2>

<p><a href="https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04">Here</a> is the link to book github ch04 that contains juniper notebooks and code used in book. There is also a screencast of Sebastian live coding for chapter 4.</p>

<h2 id="workbook-questions">Workbook questions</h2>

<p>1 What is the purpose of the context_length parameter in the GPT_CONFIG_124M dictionary?</p>

<p>A. It specifies the number of transformer blocks in the model.<br />
B. It indicates the count of attention heads in the multi-head attention mechanism.<br />
C. It represents the embedding size, transforming each token into a vector.<br />
D. It denotes the maximum number of input tokens the model can handle via the positional embeddings.</p>

<p>My answer: D<br />
Correct answer: D</p>

<p>2 The primary purpose of layer normalization in a GPT model is to adjust the activations of a neural network layer to have a mean of 0 and a <strong>__</strong><strong>__</strong> of 1.</p>

<p>A. variance<br />
B. gradient<br />
C. weight</p>

<p>My answer: A<br />
Correct answer: A</p>

<p>3 What is the main advantage of layer normalization over batch normalization in LLMs?</p>

<p>A. Layer normalization is more suitable for handling sequential data than batch normalization.<br />
B. Layer normalization normalizes each input independently of the batch size<br />
C. Layer normalization is computationally more efficient than batch normalization.<br />
D. Layer normalization is more effective at preventing overfitting than batch normalization.</p>

<p>My answer: B<br />
Correct answer: B</p>

<p>4 Which activation function is commonly used in LLMs like GPT-2, offering a smoother alternative to ReLU?</p>

<p>A. ReLU (Rectified Linear Unit)<br />
B. GELU (Gaussian Error Linear Unit)<br />
C. Sigmoid<br />
D. Tanh</p>

<p>My answer: B<br />
Correct answer: B</p>

<p>5 What is the primary purpose of shortcut connections in a deep neural network?</p>

<p>A. To reduce the number of parameters in the model.<br />
B. To preserve the flow of gradients during the backward pass in training.<br />
C. To prevent overfitting by dropping out neurons during training.<br />
D. To increase the computational efficiency of the model.</p>

<p>My answer: B<br />
Correct answer: B</p>

<p>6 What are the main components of a transformer block in a GPT model?</p>

<p>A. Convolutional layers, pooling layers, and maxout activations.<br />
B. Linear layers, ReLU activations, and batch normalization.<br />
C. Multi-head attention, layer normalization, dropout, feed forward layers, and GELU activations.<br />
D. Token embedding, positional embedding, and dropout.</p>

<p>My answer: C<br />
Correct answer: C</p>

<p>7 What is the primary reason why the GPT model generates gibberish when it is not trained?</p>

<p>A. The model is not using the correct activation function.<br />
B. The model is not using the correct tokenizer.<br />
C. The model has not learned the relationships between words and patterns in language.<br />
D. The model is not using the correct dropout rate.</p>

<p>My answer: C<br />
Correct answer: C</p>

<p>Workbook has eight answers, but there are only seven question. I assume that answer number 6 is duplicated.</p>

<h3 id="41-coding-an-llm-architecture">4.1 Coding an LLM architecture</h3>

<p>1 What is the main purpose of a GPT model, and how does it achieve this?</p>

<p>My answer: The main purpose of a GPT model is to generate a next token based on current input. To achieve this, a GPT model first tokenizes input text, then it creates token embedings and positional embedings and adds those two embeddings. Then data enters transformer block. First module in transformer block is layer normalization, then multihead attention where we have shortcut circuit between multihead attention input and output. After that fast forward block that consists of linear block, gelu activation and another linear block. Fast forward output has shortcircut with fast forward input. This is transformer block and we have 12 of those. After we are done with transformer blocks, their output goes to output layer.</p>

<p>Correct answer: A GPT model is designed to generate new text one word at a time. It achieves this by using a large deep neural network architecture that learns patterns from a massive dataset of text and then uses these patterns to predict the next word in a sequence.</p>

<p>2 Label this diagram:</p>

<p><img src="/assets/2026/05/gpt-model-diagram-labels.png" alt="GPT model diagram labels" /></p>

<p>My answer: 1 =&gt; Fast forward, 2 =&gt; layer normalization, 3 =&gt; tokenizer<br />
Correct answer: 1 =&gt; output layers, 2 =&gt; embedding layers, 3 =&gt; tokenized text</p>

<p>3 What are the key components of a GPT model, and how do they contribute to the model’s functionality?</p>

<p>My answer: Key components of a GPT model are tokenizer, transformer, linear normalization and fast forward. Tokenizer creates from input text embeding vectors where input is transformed into weights. Transformer has multihead attention that enables GPT model to know relations between tokens of input text. Layer normalization helps GPT model to resolve problem of diminishing or exploding gradients and fast forward model does adittional thinking based on multihead attention.</p>

<p>Correct answer: The key components of a GPT model include embedding layers, transformer blocks, and an output layer. Embedding layers convert words into numerical representations, transformer blocks process these representations to capture relationships between words, and the output layer predicts the probability of each word in the vocabulary.</p>

<p>4 Explain the concept of parameters in the context of LLMs like GPT.</p>

<p>My answer: If we see GPT model as a function, then parameters are attributes of that function. Based on values of those attributes and input text, we get the next token.</p>

<p>Correct answer: Parameters in LLMs refer to the trainable weights of the model. These weights are adjusted during training to minimize a loss function, allowing the model to learn from the training data and improve its ability to generate text.</p>

<p>5 What are the main differences between GPT-2 and GPT-3, and why is GPT-2 a better choice for learning LLM implementation?</p>

<p>My answer: Main difference is number of parameters. GPT-2 has less parameters so it is easier option for pretraning and usage that GPT-3</p>

<p>Correct answer: GPT-3 is a larger model with more parameters and trained on more data than GPT-2. However, GPT-2 is more suitable for learning LLM implementation because its pretrained weights are publicly available, and it can be run on a single laptop computer, unlike GPT-3, which requires a GPU cluster.</p>

<p>6 Describe the purpose of the GPT_CONFIG_124M dictionary and explain the meaning of each key-value pair.</p>

<p>My answer: This is configuration for GPT model GPT-2. It contains following parameters: vocab_size is vocabulary size of tiktoken tokenizer, emb_size is the dimension of embedding tokens, context_length is the size of positional embedding vector or how long input GPT supports, n_heads is the number of heads in multihead attention, n_layers is the number transformer blocks, drop rate is the value of optional dropout rate and qkv_bias determines do our query, key and value weights also contain along the kernel, bias.</p>

<p>Correct answer: The GPT_CONFIG_124M dictionary defines the configuration of the small GPT-2 model. It includes parameters like vocabulary size, context length, embedding dimension, number of attention heads, number of layers, dropout rate, and query-key-value bias, which determine the model’s architecture and behavior.</p>

<p>Finally, one correct answer.</p>

<p>7 What is the role of the DummyGPTModel class in the code, and how does it contribute to the overall implementation of the GPT model?</p>

<p>My answer: DummGPTModel has structure as the final GPTModel with difference that is using not implemented layer normalization and FastForward. Purpose is to shows use the overal structure of GPTModel.</p>

<p>Correct answer: The DummyGPTModel class provides a placeholder architecture for the GPT model, outlining the main components and their order. It serves as a starting point for building the full GPT model by defining the data flow and providing a framework for implementing the individual components.</p>

<p>Correct answer.</p>

<h3 id="42-normalizing-activations-with-layer-normalization">4.2 Normalizing activations with layer normalization</h3>

<p>1 What is the main purpose of layer normalization in neural networks, and how does it contribute to improved training dynamics?</p>

<p>My answer: The main purpose of layer normalization in neural networks is prevent the problem of exploding or vanishing gradients. It contributes to training dinamics by making the mean value of gradients to be zero, and to make variance of gradients equal to one.</p>

<p>Correct answer: Layer normalization aims to stabilize and accelerate neural network training by adjusting the activations of a layer to have a mean of 0 and a variance of 1. This normalization process helps prevent vanishing or exploding gradients, ensuring consistent and reliable training.</p>

<p>2 Describe the typical placement of layer normalization within GPT-2 and modern transformer architectures.</p>

<p>My answer: Layer normalization is the first layer in transformer block.</p>

<p>Correct answer: Layer normalization is commonly applied both before and after the multi- head attention module in GPT-2 and modern transformer architectures. It is also applied before the final output layer.</p>

<p>3 Pieces of the code have been removed from three places in this listing. Which of these terms have been removed and where should they go?</p>

<p>A mean B variance C norm_x D ones E scale F zeros</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">DummyLayerNorm</span> <span class="k">do</span>

  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">layer_norm</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">validate_last_axis!</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">layer_norm</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span>

    <span class="n">mean</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">__1__</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="ss">axes:</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="ss">keep_axes:</span> <span class="no">true</span><span class="p">)</span>
    <span class="n">var</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">__2__</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="ss">axes:</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="ss">keep_axes:</span> <span class="no">true</span><span class="p">)</span>
    <span class="n">norm_x</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">divide</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">subtract</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">mean</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">var</span><span class="p">,</span> <span class="n">layer_norm</span><span class="o">.</span><span class="n">eps</span><span class="p">)))</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">scale</span><span class="p">,</span> <span class="n">__3__</span><span class="p">),</span> <span class="n">layer_norm</span><span class="o">.</span><span class="n">shift</span><span class="p">)</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; mean, 2 =&gt; variance, 3 =&gt; norm_x</p>

<p>Correct answer: as my answer as I adopted this question to Elixir</p>

<p>4 Explain the difference between biased and unbiased variance calculation, and why is the biased approach preferred in the context of LLMs?</p>

<p>My answer: In biased approach, we add a small number close to zero to dividors so we do not have division with zero error.</p>

<p>Correct answer: Biased variance calculation divides by the number of inputs <em>n</em>, while unbiased variance uses <em>n</em> – <em>1</em> in the denominator to correct for bias.  In LLMs, where the embedding dimension <em>n</em> is large, the difference between these approaches is negligible, and the biased method is preferred for compatibility with GPT-2’s normalization layers and TensorFlow’s default behavior.</p>

<p>5 What are the key differences between layer normalization and batch normalization, and why layer normalization often favored in LLMs?</p>

<p>My answer: to be honest I totally missed batch normaliztion concept from the book, never heard of it.</p>

<p>Correct answer: Layer normalization normalizes across the feature dimension, while batch normalization normalizes across the batch dimension. Layer normalization offers greater flexibility and stability, especially in scenarios with varying batch sizes or resource constraints, making it suitable for LLMs.</p>

<p>6 Match the term on the left to its description on the right:</p>

<table>
  <thead>
    <tr>
      <th>Term</th>
      <th>Definition</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>Transformer Block</td>
      <td>The trainable weights of a neural network model, which are adjusted during the training process to minimize a specific loss function.</td>
    </tr>
    <tr>
      <td>Layer Normalization</td>
      <td>The output of a neural network model before applying the softmax function, representing the unnormalized probabilities of each possible output class.</td>
    </tr>
    <tr>
      <td>Parameters</td>
      <td>The core building block of a GPT model, consisting of a masked multi-head attention module and a feedforward neural network, which are applied sequentially to the input data.</td>
    </tr>
    <tr>
      <td>Logits</td>
      <td>A technique used to normalize the output of each layer in a neural network, ensuring that the data has a mean of zero and a standard deviation of one, which helps to improve the stability and performance of the model.</td>
    </tr>
  </tbody>
</table>

<p>My answer:  Transformer block =&gt; 3, layer normalization =&gt; 4, parameters =&gt; 1, logits =&gt; 2</p>

<p>Correct answer: mine answer is correct</p>

<h3 id="43-implementing-a-feed-forward-network-with-gelu-activations">4.3 Implementing a feed forward network with GELU activations</h3>

<p>1 What is GELU activation function and how does it differ from the ReLU activation function?</p>

<p>My answer: GELU activation function helps deep neural network to do better pattern recognition. GELU function values bellow zero normalizes using bell curve, while ReLU function is linear, all values bellow zero become zero.</p>

<p>Correct answer: The GELU activation function is a smooth, nonlinear function that approximates ReLU but with a non-zero gradient for almost all negative values. Unlike ReLU, which outputs zero for any negative input, GELU allows for a small, non-zero output for negative values.</p>

<p>2 Explain the purpose and structure of the FeedForward module in the context of LLM.</p>

<p>My answer: Purpose of feed forward module is to learn additional patterns from input data. It consists of two linear layers, we first transform input data to space that is 4 time bigger than original dimension, then use GELU activation function, and again we transform back to original dimension in second layer.</p>

<p>Correct answer: The FeedForward module is a small neural network consisting of two linear layers and a GELU activation function. It expands the embedding dimension into a higher-dimensional space, applies a nonlinear transformation, and then contracts back to the original dimension, allowing for richer representation learning.</p>

<p>3 Piece of code have been removed from two places in this listing. Which of these terms have been removed and where should they go?</p>

<p>A RELU B linear_defn C Embedding D GELU E Sequential</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">FeedForward</span> <span class="k">do</span>

  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">feed_forward</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">validate_last_axis!</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">feed_forward</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span>

    <span class="n">x</span>
    <span class="o">|&gt;</span> <span class="n">__1__</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">GELU</span><span class="o">.</span><span class="n">forward</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">__2__</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="p">)</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; linear_defn, 2 =&gt; linear_defn</p>

<p>Correct answer: mine is correct as I come up with thie Elixir example.</p>

<p>4 How does the FeedForward module contribute to the model’s ability to learn and generalize data?</p>

<p>My answer: do not have a clue about this question</p>

<p>Correct answer: The FeedForward module enhances the model’s ability to learn and generalize data by exploring a richer representation space through the expansion and contraction of the embedding dimension. This allows the model to capture more complex relationships within the data.</p>

<p>I do not member explanation of this why either in the book or in the screencast for chapter 4.</p>

<p>5 What is the significance of the FeedForward module having the same input and output dimensions?</p>

<p>My answer: do not know</p>

<p>Correct answer: The uniformity in input and output dimensions simplifies the architecture by enabling the stacking of multiple layers without the need to adjust dimensions between them, making the model more scalable.</p>

<h3 id="44-adding-shortcut-connections">4.4 Adding shortcut connections</h3>

<p>1 What is the vanishing gradient problem and how does it affect training deep neural networks?</p>

<p>My answer: Vanishing gradient problem is when we go backward in the deep neural network, gradient values become almost zero. This is not good for training, as network does not learn anything from gradients that does not point in any direction.</p>

<p>Correct answer:  The vanishing gradient problem occurs when gradients become progressively smaller as they propagate backward through the layers of a deep neural network, making it difficult to effectively train earlier layers.  This can hinder the learning process and prevent the model from achieving optimal performance.</p>

<p>2 Explain the concept of shortcut connections and how they address the vanishing gradient problem.</p>

<p>My answer: In shortcut connections concept, we add to the output of layer, input to that layer. This helps that we do not have gradient values close to zero.</p>

<p>Correct answer: Shortcut connections, also known as skip connections, create alternative paths for gradients to flow through the network by bypassing certain layers. This helps preserve the flow of gradients during the backward pass, mitigating the vanishing gradient problem and enabling more effective training of deeper networks.</p>

<p>3 How are shortcut connections implemented in this code listing:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">ExampleDeepNeuralNetwork</span> <span class="k">do</span>
  <span class="n">defnp</span> <span class="n">residual_linear_gelu</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">layer</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">x</span> <span class="o">+</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">layer</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">defnp</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">layer</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">x</span>
    <span class="o">|&gt;</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">linear_defn</span><span class="p">(</span><span class="n">layer</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">GELU</span><span class="o">.</span><span class="n">forward_defn</span><span class="p">()</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>My answer: In this code we are adding to output of linear_gelu function, input x to that function.</p>

<p>Correct answer: mine as I addopted this question to Elixir</p>

<p>4 What is the purpose of the ExampleDeepNeuralNetwork.weight_gradient_means and how does it demonstrate the effectiveness of shortcut connections?</p>

<p>My answer: this function calculates gradient means for weigts of each layer. It demonstrates how usage of short connections helps to resolve problem of diminishing gradients as we move backwards through the layers.</p>

<p>Correct answer: My as I addpoted this question to Elixir</p>

<p>5 This diagram shows (on the left) a deep neural network consisting of five layers and one with shortcut connections (on the right). Why are the values of the gradients bigger in the network on the right?</p>

<p><img src="/assets/2026/05/shortcut-connections-gradient-diagram.png" alt="Shortcut connections gradient diagram" /></p>

<p>My answer: Network on the right is using shortcut connections where inputs to each layer (except the layer 5) are added to layer output.</p>

<p>Correct answer: The bigger gradient values on the right are because of the shortcut connections -these stop the gradient values from getting vanishingly small through the layers.</p>

<h3 id="45-connecting-attention-and-linear-layers-in-a-transformer-block">4.5 Connecting attention and linear layers in a transformer block</h3>

<p>1 What are the key components of a transformer block and how do they contribute to the processing of input sequences?</p>

<p>My answer: The key components of a transformer block are: layer normalization, multi head attention, fast forward layer (with two linear layers and one GELU layer between them) and dropout that is optional. We have two layer normalization. First one is processing token embedings, and second is processing output of multihead attention,  so we have for each gradient mean value 0 and variance of 1. Multihead attention block adds dependencies between embedded tokens, while fast forwad also adds dependencies, but in a much larger dimension space.</p>

<p>Correct answer: A transformer block consists of multi-head attention, feed forward layers, layer normalization, and dropout. The multi-head attention analyzes relationships between elements in the input sequence, while the feed forward network modifies the data individually at each position. Layer normalization ensures consistent scaling, and dropout prevents overfitting.</p>

<p>2 Explain the concept of Pre-LayerNorm and its significance in the transformer block architecture.</p>

<p>My answer: In pre layer normalization we process token embeddings so each gradient has mean value zero, and variance of one.</p>

<p>Correct answer: Pre-LayerNorm refers to applying layer normalization before the multi- head attention and feed forward layers. This approach, unlike Post- LayerNorm, has been shown to improve training dynamics and lead to better performance in transformer models.</p>

<p>3 Describe the role of shortcut connections in the transformer block and their impact on gradient flow.</p>

<p>My answer: The role of shortcut connections is to add input to first layer normalization to output of multi head attention and to add input to second layer normalization to output of fast forward layer. Doing that if output values are close to zero, output gets short connection to input, so total output is not zero. This helps transformer block to have backward flow of gradient values that are not close to zero value.</p>

<p>Correct answer: Shortcut connections add the input of the block to its output, allowing gradients to flow more easily through the network during training. This helps to prevent vanishing gradients and improves the learning of deep models.</p>

<p>4 How does transformer block preserve the input dimensions in its output?</p>

<p>My answer: do not know</p>

<p>Correct answer: The transformer block maintains the input dimensions by applying operations that do not alter the shape of the input sequence. This allows for a one-to-one correspondence between input and output vectors, enabling its use in various sequence-to-sequence tasks.</p>

<p>5 Explain how the transformer block integrates contextual information from the entire input sequence into its output.</p>

<p>My answer: do not know</p>

<p>Correct answer: While the transformer block preserves the physical dimensions of the input sequence, the content of each output vector is re-encoded to incorporate contextual information from across the entire input sequence.  This allows the model to capture complex relationships between elements in the sequence.</p>

<h3 id="46-coding-the-gpt-model">4.6 Coding the GPT model</h3>

<p>1 What is the purpose of the GPTModel module and how does it relate to the TransformerBlock module?</p>

<p>My answer: GPTModel module represents the wole GPT architecture. It connects all blocks so we get for an input sequence a prediction of next token at the time. GPTModel module relates to TransformerBlock module so that is contains a list of TransformerBlock modules.</p>

<p>Correct answer: The GPTModel class assembles the complete GPT architecture, utilizing the previously defined TransformerBlock class as a building block. It combines token and positional embeddings, applies multiple TransformerBlock layers, and finally projects the output into the vocabulary space to predict the next token.</p>

<p>2 Explain the role of LayerNorm layer in the GPTModel architecture.</p>

<p>My answer: LayerNorm layer role is to make gradient values more stable by making sure that gradient values have a mean zero and variance one.</p>

<p>Correct answer: The LayerNorm layer is applied after the transformer blocks to normalize the output, ensuring that the data has a consistent scale and distribution.  This helps to stabilize the learning process and improve the model’s performance.</p>

<p>3 What is weight tying and how does it affect the number of parameters in a GPT model?</p>

<p>My answer: Weight tying concept is when we reuse Wq, Wk and Wv parameters from the embedding layer in the output layer. This reduces the total number of parameters needed in a GPT model.</p>

<p>Correct answer: Weight tying is a technique where the weights of the token embedding layer are reused in the output layer. This reduces the number of trainable parameters, leading to a smaller model footprint and potentially faster training.</p>

<p>4 Descripe the process of converting the output of the GPTModel into text.</p>

<p>My answer: Output of GPTModel is a matrix that contains logits, each row represents one input token, and each row has number of columns equal to number of vocabulary size. logit values are relations of each input token and all vocabulay tokens, depending on token position in input. We first must convert logits values to probability values using softmax function. Then for the last token in current input, we find for that row column index that has the highest probability value. We lookup that index in tokenizer vocabulary and this is our next token.</p>

<p>Correct answer: The output of the GPTModel is a tensor of shape [batch_size, num_tokens, vocab_size], representing the logits for each token in the vocabulary. To convert these logits into text, we need to apply a softmax function to obtain probabilities, then select the token with the highest probability for each position in the sequence.</p>

<p>5 Pieces of the code have been removed from three places in this listing. Which of these terms have been removed and where should they go?</p>

<p>A EmbeddingNative B validate_dropout! C DummyLayerNorm D init_out_head</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">GPTModel</span> <span class="k">do</span>
    <span class="k">def</span> <span class="n">new</span><span class="p">(%</span><span class="no">GPTConfig</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">))</span>
    <span class="n">norm_eps</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:norm_eps</span><span class="p">,</span> <span class="mf">1.0e-5</span><span class="p">)</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">cfg:</span> <span class="n">cfg</span><span class="p">,</span>
      <span class="ss">tok_emb:</span> <span class="n">__1__</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span><span class="p">),</span>
      <span class="ss">pos_emb:</span> <span class="n">__2__</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
      <span class="ss">drop_emb:</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">),</span>
      <span class="ss">trf_blocks:</span> <span class="n">transformer_blocks</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span><span class="p">,</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">final_norm:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">eps:</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">out_head:</span> <span class="n">__3__</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span> <span class="o">+</span> <span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>My answer: 1 =&gt; A, 2 =&gt; A, 3 =&gt; D</p>

<p>Correct answer: yes because I come up with this elixir example</p>

<p>6 How does the number of transformer blocks affect the complexity and performance of a GPT model?</p>

<p>My answer: Greater the number of transformer block means that we need more memory as the number of parameter increases, and we need more processing power as the number of needed calculations also increases.</p>

<p>Correct answer: Increasing the number of transformer blocks in a GPT model generally leads to a larger model with more parameters, requiring more computational resources. However, it can also improve the model’s ability to capture long-range dependencies in the input text, potentially leading to better performance on tasks like text generation.</p>

<p>7 Match the terms on the left to the description on the right:</p>

<table>
  <thead>
    <tr>
      <th> </th>
      <th> </th>
      <th> </th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>Multi-Head Attention</td>
      <td> </td>
      <td>a technique that normalizes the activations of a layer to have a mean of zero and a standard deviation of one, improving the stability and performance of the model.</td>
    </tr>
    <tr>
      <td>Layer Normalization</td>
      <td> </td>
      <td>allows the gradient to flow directly from the input to the output of a layer, preventing the vanishing gradient problem and enabling the training of deeper models.</td>
    </tr>
    <tr>
      <td>Shortcut Connection</td>
      <td> </td>
      <td>allows the model to attend to different parts of the input sequence simultaneously, capturing complex relationships between words and phrases.</td>
    </tr>
  </tbody>
</table>

<p>My answer: multi-head attention =&gt; row 3, layer normalization =&gt; row 1, shortcut connection =&gt; row 2</p>

<p>Correct answer: my answer</p>

<h3 id="47-generating-text">4.7 Generating text</h3>

<p>1 Explain the process by which a GPT model generates text, starting from its output tensors.</p>

<p>My answer: We call values in output matrix logits. In order to determine next token, we first identify row in logits matrix that belongs to last token in current input. Then we normalize probability values for that row using softmax. Last step is to identify index of probability with greatest value. We lookup tokenizer vocabulary with that index, and token that we get back is the next token.</p>

<p>Correct answer: The GPT model converts its output tensors into text by decoding them, selecting tokens based on a probability distribution derived from the softmax function, and then converting these tokens back into human- readable text</p>

<p>2 Describe the role of the softmax function in the text generation process.</p>

<p>My answer: The role of the softmax function is to calculate the next token probability value based on value of logits.</p>

<p>Correct answer: The softmax function transforms the output logits into a probability distribution, where each value represents the likelihood of a particular token being the next in the sequence. This allows the model to select the most probable token for generation.</p>

<p>3 What is the purpose of the generate_text_simple function, and how does it work?</p>

<p>My answer: The generate_text_simple function calculates next token based on the attributes: output logits and no of tokens that we want to generate and input tokens. It first takes first token, and predicts its next token. In output logits it identifies row that belongs to input token, transforms logits values into probabilities with softmax function, predicts next token by identifying largest probability value and looking up index of that probability in tokenizer vocabulary. That next token is added to input token, and we repeat the process for this new input where the last token is now token that we just have predicted.</p>

<p>Correct answer:  The generate_text_simple function implements a simple generative loop for a language model. It iteratively predicts the next token based on the current context, appends it to the input sequence, and repeats this process until a specified number of new tokens are generated.</p>

<p>4 Why is the softmax step technically redundant in the generate_text_simple function?</p>

<p>My answer: It is redundant because softmax normalizes logits values to interval [0,1] keeping biggest logit value also biggest in interval [0,1]</p>

<p>Correct answer: The softmax function is monotonic, meaning it preserves the order of its inputs. Therefore, applying torch.argmax directly to the logits tensor would yield the same result as applying it to the softmax output, as the position of the highest value remains unchanged</p>

<p>5 What is the significance of the greedy decoding approach in text generation?</p>

<p>My answer: The significance is that we are looing for biggest logit value.</p>

<p>Correct answer: Greedy decoding refers to the strategy of always selecting the most likely token at each step. While this approach can be efficient, it can also lead to repetitive or predictable text, as the model always chooses the most obvious continuation.</p>

<p>6 Why does the GPT model generate gibberish when it hasn’t been trained?</p>

<p>My answer: Because we still have random values for all model parameters. By training the model, those random values are updated with values that represnt token relationships with identified patterns in training data set.</p>

<p>Correct answer: The model generates incoherent text because it has not yet learned the relationships between words and their contexts. Training is crucial for the model to develop the ability to generate meaningful and coherent text.</p>

<p>7 Put these steps to implement the GPTModel architecture into the correct order:</p>

<p>A Initialize token and positional embeddings, dropout, and a linear output layer.<br />
B Implement the forward pass, combining embeddings, transformer blocks, layer normalization, and the output layer.<br />
C Create a sequential container for TransformerBlock instances.<br />
D Create a GPTModel class inheriting from nn.Module.</p>

<p>1 =&gt; D, 2 =&gt; A, 3 =&gt; B, 4 =&gt; C</p>

<p>Correct answer: 1 =&gt; D, 2 = A, 3 =&gt; C, 4 =&gt; B</p>]]></content><author><name></name></author><category term="llm-from-scratch" /><category term="gpt" /><category term="tutorials" /><category term="llm" /><category term="gpt" /><category term="workbook" /></entry><entry><title type="html">Implementing a GPT model from scratch</title><link href="https://karlosmid.com/2026/05/implementing-a-GPT-model-from-scratch/" rel="alternate" type="text/html" title="Implementing a GPT model from scratch" /><published>2026-05-02T09:00:00+00:00</published><updated>2026-05-02T09:00:00+00:00</updated><id>https://karlosmid.com/2026/05/implementing-a-GPT-model-from-scratch</id><content type="html" xml:base="https://karlosmid.com/2026/05/implementing-a-GPT-model-from-scratch/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>This post implements a GPT-style model from scratch in Elixir/Nx, building on the <a href="/coding-attention-mechanism-in-elixir">attention mechanism post</a>. We start with the GPT-2 configuration and parameter counts, then add layer normalization, GELU feed-forward layers, shortcut connections, transformer blocks, the full GPTModel, and a simple greedy text generation loop.</p>

<p>You can find the source code in my <a href="https://github.com/karlosmid/llm-from-scratch">GitHub repository</a>.</p>

<h2 id="coding-an-llm-architecture">Coding an LLM architecture</h2>

<p>To finish building an LLM, we need to create an LLM architecture. We already know that we are building a GPT-style LLM. To finish the GPT model, we need to add several transformer blocks to the embedding layer and masked multi-head attention.<br />
Let’s talk about the scale of our future GPT model. So far, we have used a dimension of three, just to make it possible to fit examples on the book pages. We are scaling up to 124 million parameters, the size of a small GPT-2 model. What are those parameters? Parameters are trainable weights (remember query, key, and value weights). You can think of them as model variables. An LLM is all about crazy scales. Imagine that you need to develop a function with 124 million parameters. Well, this is basically an LLM. Remember, we first create initial weight matrices using randomization, and in the training phase, we tune those values.<br />
If we have a tensor of size {2048, 2048}, which is a matrix of 2048 rows and columns, we call this tensor a layer with 2048x2048 = 4,196,304 parameters.<br />
Now we are ready to understand the configuration parameters of the GPT-2 model, since we mentioned all these parameters in previous blog posts:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">LlmScratch</span><span class="o">.</span><span class="no">GPTConfig</span><span class="p">{</span>
  <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
  <span class="ss">context_length:</span> <span class="mi">1_024</span><span class="p">,</span>
  <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
  <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
  <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
  <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
  <span class="ss">qkv_bias:</span> <span class="no">false</span>
<span class="p">}</span>

<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">GPTConfig</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  GPT model configuration.

  This struct mirrors the config keys used by the Python examples.

  ## Fields

    * `:vocab_size` - number of tokens in the tokenizer vocabulary. This
      determines the number of rows in the token embedding table and the final
      logits dimension produced by the output head. For GPT-2 this is commonly
      `50_257`.

    * `:context_length` - maximum number of token positions the model can
      process in one forward pass. This determines the size of the positional
      embedding table. Inputs longer than this cannot be represented by the
      configured positional embeddings.

    * `:emb_dim` - width of each token and positional embedding vector. This is
      also the model's hidden size: tensors flowing through transformer blocks
      have shape `{batch_size, seq_len, emb_dim}`.

    * `:n_heads` - number of attention heads used by real transformer blocks.
      The dummy GPT model stores this value for config compatibility, but its
      placeholder transformer blocks do not use it yet.

    * `:n_layers` - number of transformer blocks in the model. In
      `LlmScratch.DummyGPTModel`, this controls how many dummy identity blocks
      are created.

    * `:drop_rate` - dropout probability used after adding token and positional
      embeddings. A value of `0.1` means ten percent of activations are dropped
      during training mode. Dropout is skipped in inference mode.

    * `:qkv_bias` - whether query, key, and value projections include bias
      terms in real attention layers. The dummy GPT model keeps this field for
      compatibility with the full GPT config shape, but does not use it because
      its transformer blocks are placeholders.
  """</span>

  <span class="nv">@enforce_keys</span> <span class="p">[</span>
    <span class="ss">:vocab_size</span><span class="p">,</span>
    <span class="ss">:context_length</span><span class="p">,</span>
    <span class="ss">:emb_dim</span><span class="p">,</span>
    <span class="ss">:n_heads</span><span class="p">,</span>
    <span class="ss">:n_layers</span><span class="p">,</span>
    <span class="ss">:drop_rate</span><span class="p">,</span>
    <span class="ss">:qkv_bias</span>
  <span class="p">]</span>

  <span class="k">defstruct</span> <span class="p">[</span>
    <span class="ss">:vocab_size</span><span class="p">,</span>
    <span class="ss">:context_length</span><span class="p">,</span>
    <span class="ss">:emb_dim</span><span class="p">,</span>
    <span class="ss">:n_heads</span><span class="p">,</span>
    <span class="ss">:n_layers</span><span class="p">,</span>
    <span class="ss">:drop_rate</span><span class="p">,</span>
    <span class="ss">:qkv_bias</span>
  <span class="p">]</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">vocab_size:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">context_length:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">emb_dim:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">n_heads:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">n_layers:</span> <span class="n">non_neg_integer</span><span class="p">(),</span>
          <span class="ss">drop_rate:</span> <span class="n">float</span><span class="p">(),</span>
          <span class="ss">qkv_bias:</span> <span class="n">boolean</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">estimated_parameter_count</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Estimates GPT model parameters directly from the configuration.

  This mirrors the architecture used by `LlmScratch.GPTModel` without
  allocating tensors:

    * token embeddings: `vocab_size * emb_dim`
    * positional embeddings: `context_length * emb_dim`
    * untied output projection: `emb_dim * vocab_size`
    * each transformer block: attention projections, feed-forward projections,
      and two layer norms
    * final layer norm: scale and shift, each shaped `{emb_dim}`

  Use `LlmScratch.GPTModel.total_parameters/1` when you already have a
  concrete model and want to count its actual tensors.
  """</span>
  <span class="k">def</span> <span class="n">estimated_parameter_count</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">token_embedding_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span>
    <span class="n">positional_embedding_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span>
    <span class="n">output_projection_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span>
    <span class="n">block_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">*</span> <span class="n">estimated_transformer_block_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
    <span class="n">final_norm_params</span> <span class="o">=</span> <span class="n">estimated_layer_norm_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span>

    <span class="n">token_embedding_params</span> <span class="o">+</span> <span class="n">positional_embedding_params</span> <span class="o">+</span> <span class="n">output_projection_params</span> <span class="o">+</span> <span class="n">block_params</span> <span class="o">+</span>
      <span class="n">final_norm_params</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">estimated_transformer_block_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">estimated_attention_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">estimated_feed_forward_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span> <span class="o">+</span>
      <span class="mi">2</span> <span class="o">*</span> <span class="n">estimated_layer_norm_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">estimated_attention_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">qkv_kernel_params</span> <span class="o">=</span> <span class="mi">3</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span>
    <span class="n">qkv_bias_params</span> <span class="o">=</span> <span class="k">if</span> <span class="n">cfg</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="mi">3</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="k">else</span><span class="p">:</span> <span class="mi">0</span>
    <span class="n">output_projection_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">+</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span>

    <span class="n">qkv_kernel_params</span> <span class="o">+</span> <span class="n">qkv_bias_params</span> <span class="o">+</span> <span class="n">output_projection_params</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">estimated_feed_forward_parameter_count</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">expanded_dim</span> <span class="o">=</span> <span class="mi">4</span> <span class="o">*</span> <span class="n">emb_dim</span>
    <span class="n">first_projection_params</span> <span class="o">=</span> <span class="n">emb_dim</span> <span class="o">*</span> <span class="n">expanded_dim</span> <span class="o">+</span> <span class="n">expanded_dim</span>
    <span class="n">second_projection_params</span> <span class="o">=</span> <span class="n">expanded_dim</span> <span class="o">*</span> <span class="n">emb_dim</span> <span class="o">+</span> <span class="n">emb_dim</span>

    <span class="n">first_projection_params</span> <span class="o">+</span> <span class="n">second_projection_params</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">estimated_layer_norm_parameter_count</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">emb_dim</span>
<span class="k">end</span>
</code></pre></div></div>

<p><code class="language-plaintext highlighter-rouge">n_layers</code> is new, and we will learn about it in this blog post. <code class="language-plaintext highlighter-rouge">context_length</code> is the number of tokens in our input tensor.<br />
Time for some Elixir code. Let’s create a GPT backbone module. I added documentation for the <code class="language-plaintext highlighter-rouge">new</code> and <code class="language-plaintext highlighter-rouge">forward</code> functions. For now, the <code class="language-plaintext highlighter-rouge">Transformer</code> and <code class="language-plaintext highlighter-rouge">Layer</code> modules just return the input unchanged in their <code class="language-plaintext highlighter-rouge">forward</code> functions.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">DummyGPTModel</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Minimal GPT-shaped model implemented with Nx tensors.

  This mirrors the Python `DummyGPTModel` from the book:

    * token embedding
    * positional embedding
    * embedding dropout
    * a list of placeholder transformer blocks
    * placeholder layer norm
    * vocabulary projection without bias

  The transformer blocks and layer norm are intentionally no-ops.
  """</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span>
    <span class="no">DummyLayerNorm</span><span class="p">,</span>
    <span class="no">DummyTransformerBlock</span><span class="p">,</span>
    <span class="no">EmbeddingNative</span><span class="p">,</span>
    <span class="no">GPTConfig</span><span class="p">,</span>
    <span class="no">SelfAttentionV2</span>
  <span class="p">}</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:cfg</span><span class="p">,</span> <span class="ss">:tok_emb</span><span class="p">,</span> <span class="ss">:pos_emb</span><span class="p">,</span> <span class="ss">:drop_emb</span><span class="p">,</span> <span class="ss">:trf_blocks</span><span class="p">,</span> <span class="ss">:final_norm</span><span class="p">,</span> <span class="ss">:out_head</span><span class="p">]</span>

  <span class="nv">@type</span> <span class="n">linear_no_bias</span> <span class="p">::</span> <span class="p">%{</span><span class="ss">kernel:</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()}</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">cfg:</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">tok_emb:</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">pos_emb:</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">drop_emb:</span> <span class="n">float</span><span class="p">(),</span>
          <span class="ss">trf_blocks:</span> <span class="p">[</span><span class="no">DummyTransformerBlock</span><span class="o">.</span><span class="n">t</span><span class="p">()],</span>
          <span class="ss">final_norm:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">out_head:</span> <span class="n">linear_no_bias</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Creates a dummy GPT model.

  `cfg` is a `LlmScratch.GPTConfig` struct.

  ## Options

    * `:seed` - deterministic seed for native Nx/Axon initialization.

  ## Returns

  A `%LlmScratch.DummyGPTModel{}` struct with these fields:

    * `:cfg` - the `LlmScratch.GPTConfig` used to build the model. It records
      the vocabulary size, context length, embedding dimension, number of
      layers, dropout rate, and compatibility fields from the GPT config.

    * `:tok_emb` - token embedding table with stored weight shape
      `{vocab_size, emb_dim}`.

    * `:pos_emb` - positional embedding table with stored weight shape
      `{context_length, emb_dim}`.

    * `:drop_emb` - dropout rate from `cfg.drop_rate`.

    * `:trf_blocks` - list of dummy transformer blocks. Its length is
      `cfg.n_layers`. In this dummy model each block returns its input
      unchanged.

    * `:final_norm` - placeholder final layer normalization module. It currently
      returns its input unchanged.

    * `:out_head` - final vocabulary projection without bias. It contains a
      kernel shaped `{emb_dim, vocab_size}`.
  """</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(%</span><span class="no">GPTConfig</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">))</span>

    <span class="n">tok_emb</span> <span class="o">=</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span><span class="p">)</span>
    <span class="n">pos_emb</span> <span class="o">=</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">out_head</span> <span class="o">=</span> <span class="n">init_out_head</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span><span class="p">)</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">cfg:</span> <span class="n">cfg</span><span class="p">,</span>
      <span class="ss">tok_emb:</span> <span class="n">tok_emb</span><span class="p">,</span>
      <span class="ss">pos_emb:</span> <span class="n">pos_emb</span><span class="p">,</span>
      <span class="ss">drop_emb:</span> <span class="n">cfg</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">,</span>
      <span class="ss">trf_blocks:</span> <span class="n">for</span><span class="p">(</span><span class="n">_</span> <span class="o">&lt;-</span> <span class="mi">1</span><span class="o">..</span><span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="no">DummyTransformerBlock</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">)),</span>
      <span class="ss">final_norm:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">),</span>
      <span class="ss">out_head:</span> <span class="n">out_head</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Runs a forward pass over token ids shaped `{batch_size, seq_len}`.

  ## Steps

    * Read `seq_len` from the input tensor shape.
    * Look up token embeddings for every token id in `in_idx`, producing a
      tensor shaped `{batch_size, seq_len, emb_dim}`.
    * Build positional indices `0..seq_len-1`.
    * Look up positional embeddings for those positions, producing a tensor
      shaped `{seq_len, emb_dim}`.
    * Add token embeddings and positional embeddings. Nx broadcasts the
      positional embeddings across the batch dimension.
    * Apply embedding dropout when `mode: :train`; leave the tensor unchanged
      in inference mode.
    * Pass the tensor through each dummy transformer block. These placeholder
      blocks currently return their input unchanged.
    * Apply the placeholder final layer norm. It currently returns its input
      unchanged.
    * Project the final hidden states through the output head kernel.

  Options:

    * `:mode` - `:inference` (default) or `:train`.
    * `:key` - optional `Nx.Random` key for dropout in train mode.

  ## Returns

  A logits tensor shaped `{batch_size, seq_len, vocab_size}`.

  For each token position, the last dimension contains one score per vocabulary
  token.
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">in_idx</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">in_idx</span><span class="p">)</span>
    <span class="n">tok_embeds</span> <span class="o">=</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">tok_emb</span><span class="p">,</span> <span class="n">in_idx</span><span class="p">)</span>

    <span class="n">pos_indices</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">iota</span><span class="p">({</span><span class="n">seq_len</span><span class="p">},</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>
    <span class="n">pos_embeds</span> <span class="o">=</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span><span class="p">,</span> <span class="n">pos_indices</span><span class="p">)</span>

    <span class="n">x</span> <span class="o">=</span>
      <span class="n">tok_embeds</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">pos_embeds</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="n">maybe_dropout</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">drop_emb</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="n">x</span> <span class="o">=</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="o">&amp;</span><span class="no">DummyTransformerBlock</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="nv">&amp;2</span><span class="p">))</span>
    <span class="n">x</span> <span class="o">=</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">final_norm</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

    <span class="c1"># {batch_size, seq_len, emb_dim} dot {emb_dim, vocab_size} -&gt;</span>
    <span class="c1"># {batch_size, seq_len, vocab_size}</span>
    <span class="n">linear</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">call</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">call</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">in_idx</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">in_idx</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">linear</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="p">%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">dot</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="p">[</span><span class="mi">2</span><span class="p">],</span> <span class="n">kernel</span><span class="p">,</span> <span class="p">[</span><span class="mi">0</span><span class="p">])</span>

  <span class="k">defp</span> <span class="n">maybe_dropout</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">drop_rate</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span> <span class="ow">when</span> <span class="n">drop_rate</span> <span class="o">&gt;</span> <span class="mf">0.0</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:mode</span><span class="p">,</span> <span class="ss">:inference</span><span class="p">)</span> <span class="k">do</span>
      <span class="ss">:train</span> <span class="o">-&gt;</span>
        <span class="n">key</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:key</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="no">System</span><span class="o">.</span><span class="n">unique_integer</span><span class="p">([</span><span class="ss">:positive</span><span class="p">])))</span>

        <span class="p">%</span><span class="no">Axon</span><span class="o">.</span><span class="no">StatefulOutput</span><span class="p">{</span><span class="ss">output:</span> <span class="n">dropped</span><span class="p">}</span> <span class="o">=</span>
          <span class="no">Axon</span><span class="o">.</span><span class="no">Layers</span><span class="o">.</span><span class="n">dropout</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="ss">rate:</span> <span class="n">drop_rate</span><span class="p">,</span> <span class="ss">mode:</span> <span class="ss">:train</span><span class="p">)</span>

        <span class="n">dropped</span>

      <span class="ss">:inference</span> <span class="o">-&gt;</span>
        <span class="n">x</span>

      <span class="n">mode</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"mode must be :train or :inference, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">mode</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">maybe_dropout</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">_drop_rate</span><span class="p">,</span> <span class="n">_opts</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">x</span>

  <span class="nv">@spec</span> <span class="n">init_out_head</span><span class="p">(</span><span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">integer</span><span class="p">())</span> <span class="p">::</span> <span class="n">linear_no_bias</span><span class="p">()</span>
  <span class="k">defp</span> <span class="n">init_out_head</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">SelfAttentionV2</span><span class="o">.</span><span class="n">init_dense_weights</span><span class="p">(</span>
        <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span>
        <span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span>
        <span class="n">seed</span><span class="p">,</span>
        <span class="no">false</span><span class="p">,</span>
        <span class="s2">"out_head"</span>
      <span class="p">)</span>

    <span class="p">%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">System</span><span class="o">.</span><span class="n">unique_integer</span><span class="p">([</span><span class="ss">:positive</span><span class="p">])</span>
  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">seed</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">seed</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"seed must be an integer or nil, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>
<span class="k">end</span>

<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">DummyTransformerBlock</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Placeholder transformer block.

  It mirrors the book's dummy PyTorch block and returns its input unchanged.
  """</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:cfg</span><span class="p">]</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span><span class="ss">cfg:</span> <span class="n">map</span><span class="p">()}</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="n">map</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">cfg</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span><span class="ss">cfg:</span> <span class="n">cfg</span><span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{},</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">x</span>
<span class="k">end</span>

<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">DummyLayerNorm</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Placeholder layer normalization.

  It stores the intended normalized shape and epsilon, but returns inputs
  unchanged just like the Python dummy implementation.
  """</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:normalized_shape</span><span class="p">,</span> <span class="ss">:eps</span><span class="p">]</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">normalized_shape:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">eps:</span> <span class="n">float</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="n">pos_integer</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">normalized_shape</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span>
      <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">normalized_shape</span><span class="p">)</span> <span class="ow">and</span> <span class="n">normalized_shape</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="n">eps</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:eps</span><span class="p">,</span> <span class="mf">1.0e-5</span><span class="p">)</span>
    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span><span class="ss">normalized_shape:</span> <span class="n">normalized_shape</span><span class="p">,</span> <span class="ss">eps:</span> <span class="n">eps</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{},</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">x</span>
<span class="k">end</span>
</code></pre></div></div>

<p>Here is a test that uses this dummy GPT module to produce some logits. Since we have big tensors, we call EXLA to the rescue!</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="n">test</span> <span class="s2">"dummy GPT model returns logits for GPT-2 tokenized batch"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">txt1</span> <span class="o">=</span> <span class="s2">"Every effort moves you"</span>
    <span class="n">txt2</span> <span class="o">=</span> <span class="s2">"Every day holds a"</span>

    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">tokens1</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">,</span> <span class="n">txt1</span><span class="p">)</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">tokens2</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">,</span> <span class="n">txt2</span><span class="p">)</span>

    <span class="n">batch</span> <span class="o">=</span>
      <span class="p">[</span><span class="n">tokens1</span><span class="p">,</span> <span class="n">tokens2</span><span class="p">]</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="o">&amp;</span><span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}))</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">stack</span><span class="p">()</span>

    <span class="n">expected_batch</span> <span class="o">=</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span>
        <span class="p">[</span>
          <span class="p">[</span><span class="mi">6109</span><span class="p">,</span> <span class="mi">3626</span><span class="p">,</span> <span class="mi">6100</span><span class="p">,</span> <span class="mi">345</span><span class="p">],</span>
          <span class="p">[</span><span class="mi">6109</span><span class="p">,</span> <span class="mi">1110</span><span class="p">,</span> <span class="mi">6622</span><span class="p">,</span> <span class="mi">257</span><span class="p">]</span>
        <span class="p">],</span>
        <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
      <span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">expected_batch</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all</span><span class="p">()</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span> <span class="o">==</span> <span class="mi">1</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">DummyGPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>
    <span class="n">logits</span> <span class="o">=</span> <span class="no">DummyGPTModel</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">batch</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">50_257</span><span class="p">}</span>

    <span class="c1"># We only check the first 5 dimensions of logits, as their full dimension of 50_257 is too large.</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">logits</span><span class="p">[[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="o">..</span><span class="mi">4</span><span class="p">]],</span> <span class="p">[</span>
      <span class="mf">0.11659280955791473</span><span class="p">,</span>
      <span class="o">-</span><span class="mf">0.01700635254383087</span><span class="p">,</span>
      <span class="mf">0.1895323395729065</span><span class="p">,</span>
      <span class="o">-</span><span class="mf">0.0104571133852005</span><span class="p">,</span>
      <span class="o">-</span><span class="mf">0.09112129360437393</span>
    <span class="p">])</span>

    <span class="n">assert_close</span><span class="p">(</span><span class="n">logits</span><span class="p">[[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="o">..</span><span class="mi">4</span><span class="p">]],</span> <span class="p">[</span>
      <span class="mf">0.0482957549393177</span><span class="p">,</span>
      <span class="o">-</span><span class="mf">0.3872452676296234</span><span class="p">,</span>
      <span class="mf">0.2935338020324707</span><span class="p">,</span>
      <span class="mf">0.048269566148519516</span><span class="p">,</span>
      <span class="mf">0.058328352868556976</span>
    <span class="p">])</span>

    <span class="n">assert_close</span><span class="p">(</span><span class="n">logits</span><span class="p">[[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">0</span><span class="o">..</span><span class="mi">4</span><span class="p">]],</span> <span class="p">[</span>
      <span class="o">-</span><span class="mf">0.32439061999320984</span><span class="p">,</span>
      <span class="mf">0.3353412449359894</span><span class="p">,</span>
      <span class="mf">0.25242942571640015</span><span class="p">,</span>
      <span class="mf">0.3892766535282135</span><span class="p">,</span>
      <span class="o">-</span><span class="mf">0.49838411808013916</span>
    <span class="p">])</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">assert_close</span><span class="p">(</span><span class="n">actual</span><span class="p">,</span> <span class="n">expected</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">expected</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">expected</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all_close</span><span class="p">(</span><span class="n">actual</span><span class="p">,</span> <span class="n">expected</span><span class="p">,</span> <span class="ss">atol:</span> <span class="mf">1.0e-6</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span> <span class="o">==</span> <span class="mi">1</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="normalizing-activations-with-layer-normalization">Normalizing activations with layer normalization</h2>

<p>For a neural network to effectively predict the next token from the input tokens, its set of weight parameters should be adjusted during training so the loss function is minimized. Without normalizing activations, the gradients of the weight parameters could easily vanish or explode, which makes minimizing the function harder.</p>

<p>What is a gradient? Let’s dive into some Elixir code. <code class="language-plaintext highlighter-rouge">defn</code> is an <code class="language-plaintext highlighter-rouge">Nx</code> macro that tells <code class="language-plaintext highlighter-rouge">Nx</code> to run this function in <code class="language-plaintext highlighter-rouge">EXLA</code> if it is available:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">SimpleGradient</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Small Nx gradient example matching the scalar PyTorch example from the book.

  Nx tensors do not store gradients in a mutable `.grad` field. The `backward/4`
  function returns the loss and gradients as values.
  """</span>

  <span class="kn">import</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Computes binary cross entropy loss for a one-input logistic unit.

  ## Parameters

    * `w1` - trainable weight tensor.
    * `b` - trainable bias tensor.
    * `x1` - input feature tensor.
    * `y` - target label tensor.

  ## Returns

  A scalar tensor with the mean binary cross entropy between `sigmoid(x1 * w1 + b)`
  and `y`.
  """</span>
  <span class="n">defn</span> <span class="n">loss</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">z</span> <span class="o">=</span> <span class="n">x1</span> <span class="o">*</span> <span class="n">w1</span> <span class="o">+</span> <span class="n">b</span>
    <span class="n">a</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">sigmoid</span><span class="p">(</span><span class="n">z</span><span class="p">)</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="o">-</span><span class="p">(</span><span class="n">y</span> <span class="o">*</span> <span class="no">Nx</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">a</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">y</span><span class="p">)</span> <span class="o">*</span> <span class="no">Nx</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">a</span><span class="p">)))</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Computes mean squared error loss.

  This mirrors `nn.MSELoss()` in the PyTorch examples. It is defined with
  `Nx.Defn`, so it can be called from other differentiable functions that use
  `grad/2` or `value_and_grad/2`.

  ## Parameters

    * `output` - predicted values.
    * `target` - expected values with a shape compatible with `output`.

  ## Returns

  A scalar tensor with the mean of `(output - target)^2`.

  ## Example

      iex&gt; output = Nx.tensor([[1.0, 2.0]])
      iex&gt; target = Nx.tensor([[0.0, 1.0]])
      iex&gt; LlmScratch.SimpleGradient.mse_loss(output, target) |&gt; Nx.to_number()
      1.0
  """</span>
  <span class="n">defn</span> <span class="n">mse_loss</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">output</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">subtract</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Computes the gradient of `loss/4` with respect to `w1`.

  ## Parameters

    * `w1` - trainable weight tensor to differentiate.
    * `b` - trainable bias tensor used to compute the loss.
    * `x1` - input feature tensor used to compute the loss.
    * `y` - target label tensor used to compute the loss.

  ## Returns

  A tensor with the same shape as `w1`, containing `dL/dw1`.
  """</span>
  <span class="n">defn</span> <span class="n">grad_w1</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">grad</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="k">fn</span> <span class="n">w1</span> <span class="o">-&gt;</span>
      <span class="n">loss</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Computes the gradient of `loss/4` with respect to `b`.

  ## Parameters

    * `w1` - trainable weight tensor used to compute the loss.
    * `b` - trainable bias tensor to differentiate.
    * `x1` - input feature tensor used to compute the loss.
    * `y` - target label tensor used to compute the loss.

  ## Returns

  A tensor with the same shape as `b`, containing `dL/db`.
  """</span>
  <span class="n">defn</span> <span class="n">grad_b</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">grad</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="k">fn</span> <span class="n">b</span> <span class="o">-&gt;</span>
      <span class="n">loss</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Computes loss and gradients for the trainable parameters `{w1, b}`.

  This is the Nx counterpart to calling `loss.backward()` in the PyTorch example.
  Instead of mutating `w1.grad` and `b.grad`, it returns the gradients.

  ## Parameters

    * `w1` - trainable weight tensor.
    * `b` - trainable bias tensor.
    * `x1` - input feature tensor used to compute the loss.
    * `y` - target label tensor used to compute the loss.

  ## Returns

  `{loss, {grad_w1, grad_b}}`, where `loss` is a scalar tensor and the gradient
  tensors match the shapes of `w1` and `b`.
  """</span>
  <span class="n">defn</span> <span class="n">backward</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">value_and_grad</span><span class="p">({</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">},</span> <span class="k">fn</span> <span class="p">{</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">}</span> <span class="o">-&gt;</span>
      <span class="n">loss</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>This is a lot of math, but it gives you a feeling for it. Sebastian said that a deeper knowledge of gradients is not mandatory to successfully build a GPT LLM.</p>

<p><img src="/assets/loss-gradient-curve.png" alt="Loss curve for the simple gradient example" /></p>

<p>This graph shows that the loss function goes to zero as <code class="language-plaintext highlighter-rouge">w1</code> values get larger. That means the prediction of a <code class="language-plaintext highlighter-rouge">y</code> value of 1 gets better.</p>

<p><img src="/assets/gradient-curve.png" alt="Gradient curve for the simple gradient example" /></p>

<p>The same is true for the gradient, which goes to zero as <code class="language-plaintext highlighter-rouge">w1</code> gets larger and we get a better prediction of 1.</p>

<p>Here is a test with the book example:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
<span class="n">test</span> <span class="s2">"simple gradient example matches book data"</span> <span class="k">do</span>
    <span class="n">y</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">1.0</span><span class="p">])</span>
    <span class="n">x1</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">1.1</span><span class="p">])</span>
    <span class="n">w1</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">2.2</span><span class="p">])</span>
    <span class="n">b</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">0.0</span><span class="p">])</span>

    <span class="n">grad_l_w1</span> <span class="o">=</span> <span class="no">SimpleGradient</span><span class="o">.</span><span class="n">grad_w1</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
    <span class="n">grad_l_b</span> <span class="o">=</span> <span class="no">SimpleGradient</span><span class="o">.</span><span class="n">grad_b</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>

    <span class="p">{</span><span class="n">loss</span><span class="p">,</span> <span class="p">{</span><span class="n">w1_grad</span><span class="p">,</span> <span class="n">b_grad</span><span class="p">}}</span> <span class="o">=</span> <span class="no">SimpleGradient</span><span class="o">.</span><span class="n">backward</span><span class="p">(</span><span class="n">w1</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>

    <span class="n">assert_close</span><span class="p">(</span><span class="n">loss</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.08518788</span><span class="p">))</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">w1_grad</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.0898</span><span class="p">],</span> <span class="ss">atol:</span> <span class="mf">1.0e-4</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">b_grad</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.0817</span><span class="p">],</span> <span class="ss">atol:</span> <span class="mf">1.0e-4</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">grad_l_w1</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.0898</span><span class="p">],</span> <span class="ss">atol:</span> <span class="mf">1.0e-4</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">grad_l_b</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.0817</span><span class="p">],</span> <span class="ss">atol:</span> <span class="mf">1.0e-4</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Now we can implement the layer normalization module:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">DummyLayerNorm</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Layer normalization implemented with Nx tensors.

  Layer normalization normalizes each input vector over its last dimension. For
  GPT-style hidden states shaped `{batch_size, seq_len, emb_dim}`, this means
  every token embedding is normalized independently across its `emb_dim`
  features.

  This mirrors the book's PyTorch implementation:

      mean = x.mean(dim=-1, keepdim=True)
      var = x.var(dim=-1, keepdim=True, unbiased=False)
      norm_x = (x - mean) / torch.sqrt(var + eps)
      scale * norm_x + shift

  The Nx equivalents are:

      mean = Nx.mean(x, axes: [-1], keep_axes: true)
      var = Nx.variance(x, axes: [-1], keep_axes: true)

  `Nx.variance/2` computes population variance, which matches
  `torch.var(..., unbiased=False)`. The `scale` and `shift` tensors correspond
  to PyTorch's trainable parameters initialized with `torch.ones(emb_dim)` and
  `torch.zeros(emb_dim)`.
  """</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:emb_dim</span><span class="p">,</span> <span class="ss">:eps</span><span class="p">,</span> <span class="ss">:scale</span><span class="p">,</span> <span class="ss">:shift</span><span class="p">]</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">emb_dim:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">eps:</span> <span class="n">float</span><span class="p">(),</span>
          <span class="ss">scale:</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">shift:</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="n">pos_integer</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Creates a layer normalization module for vectors with `emb_dim` features.

  ## Options

    * `:eps` - small value added to the variance before the square root.
      Defaults to `1.0e-5`, matching the PyTorch example.

  ## Returns

  A `%LlmScratch.DummyLayerNorm{}` struct containing:

    * `:emb_dim` - expected size of the input tensor's last dimension.
    * `:eps` - numerical stability value.
    * `:scale` - tensor of ones shaped `{emb_dim}`.
    * `:shift` - tensor of zeros shaped `{emb_dim}`.

  ## Example

      iex&gt; ln = LlmScratch.DummyLayerNorm.new(5)
      iex&gt; Nx.shape(ln.scale)
      {5}
      iex&gt; Nx.shape(ln.shift)
      {5}
  """</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span> <span class="ow">and</span> <span class="n">emb_dim</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="n">eps</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:eps</span><span class="p">,</span> <span class="mf">1.0e-5</span><span class="p">)</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">emb_dim:</span> <span class="n">emb_dim</span><span class="p">,</span>
      <span class="ss">eps:</span> <span class="n">eps</span><span class="p">,</span>
      <span class="ss">scale:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="mf">1.0</span><span class="p">,</span> <span class="p">{</span><span class="n">emb_dim</span><span class="p">})</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">as_type</span><span class="p">({</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">}),</span>
      <span class="ss">shift:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="p">{</span><span class="n">emb_dim</span><span class="p">})</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">as_type</span><span class="p">({</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Applies layer normalization over the last axis of `x`.

  The input tensor's last dimension must equal the layer's `emb_dim`. Any number
  of leading dimensions is supported, for example `{batch_size, emb_dim}` or
  `{batch_size, seq_len, emb_dim}`.

  ## Example

      {batch_example, _key} =
        Nx.Random.normal(Nx.Random.key(123), 0.0, 1.0, shape: {2, 5})

      ln = LlmScratch.DummyLayerNorm.new(5)
      out_ln = LlmScratch.DummyLayerNorm.forward(ln, batch_example)

      Nx.mean(out_ln, axes: [-1], keep_axes: true)
      Nx.variance(out_ln, axes: [-1], keep_axes: true)

  The resulting mean is close to zero and the population variance is close to
  one for each row. Because `eps` is included in the denominator, variance can
  be slightly below `1.0` before rounding.

  Raises `ArgumentError` if `x` does not have `emb_dim` features on its last
  axis.
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">layer_norm</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">validate_last_axis!</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">layer_norm</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span>

    <span class="n">mean</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="ss">axes:</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="ss">keep_axes:</span> <span class="no">true</span><span class="p">)</span>
    <span class="n">var</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">variance</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="ss">axes:</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="ss">keep_axes:</span> <span class="no">true</span><span class="p">)</span>
    <span class="n">norm_x</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">divide</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">subtract</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">mean</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">var</span><span class="p">,</span> <span class="n">layer_norm</span><span class="o">.</span><span class="n">eps</span><span class="p">)))</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">scale</span><span class="p">,</span> <span class="n">norm_x</span><span class="p">),</span> <span class="n">layer_norm</span><span class="o">.</span><span class="n">shift</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_last_axis!</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">last_axis</span> <span class="o">=</span>
      <span class="n">x</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Tuple</span><span class="o">.</span><span class="n">to_list</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">last</span><span class="p">()</span>

    <span class="k">unless</span> <span class="n">last_axis</span> <span class="o">==</span> <span class="n">emb_dim</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
            <span class="s2">"expected input last dimension to be </span><span class="si">#{</span><span class="n">emb_dim</span><span class="si">}</span><span class="s2">, got </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">last_axis</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>Here is a test that exercises this module:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"layer norm normalizes over the last dimension"</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">batch_example</span><span class="p">,</span> <span class="n">_key</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">normal</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">123</span><span class="p">),</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="ss">shape:</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">5</span><span class="p">})</span>
    <span class="n">layer_norm</span> <span class="o">=</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="mi">5</span><span class="p">)</span>
    <span class="n">out_ln</span> <span class="o">=</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">layer_norm</span><span class="p">,</span> <span class="n">batch_example</span><span class="p">)</span>
    <span class="n">mean</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">out_ln</span><span class="p">,</span> <span class="ss">axes:</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="ss">keep_axes:</span> <span class="no">true</span><span class="p">)</span>
    <span class="n">var</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">variance</span><span class="p">(</span><span class="n">out_ln</span><span class="p">,</span> <span class="ss">axes:</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="ss">keep_axes:</span> <span class="no">true</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">scale</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">5</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">shift</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">5</span><span class="p">}</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">scale</span><span class="p">,</span> <span class="p">[</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">])</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">shift</span><span class="p">,</span> <span class="p">[</span><span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">])</span>

    <span class="n">assert_close</span><span class="p">(</span><span class="n">mean</span><span class="p">,</span> <span class="p">[[</span><span class="mf">0.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">0.0</span><span class="p">]],</span> <span class="ss">atol:</span> <span class="mf">1.0e-6</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">var</span><span class="p">,</span> <span class="p">[[</span><span class="mf">1.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">1.0</span><span class="p">]],</span> <span class="ss">atol:</span> <span class="mf">1.0e-4</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>The point of layer normalization is that the <code class="language-plaintext highlighter-rouge">out_ln</code> means are close to zero, and the variances are close to 1.</p>

<h2 id="implementing-a-feed-forward-network-with-gelu-activations">Implementing a feed-forward network with GELU activations</h2>

<p>Prepare for a scary word: GELU, or Gaussian Error Linear Unit. Remember the Gaussian distribution, which many natural values follow?</p>

<p>For example, if the average height in a group is <code class="language-plaintext highlighter-rouge">180 cm</code>, and most people are close to that value, the values often look like a bell curve:</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>frequency
   20 |                    *
   15 |               *         *
   10 |
    8 |          *                   *
    5 |
    3 |     *                             *
    1 |*                                      *
      +-----------------------------------------
       160  165  170  175  180  185  190  195  200
                              height
</code></pre></div></div>

<p>Most values are near the mean, <code class="language-plaintext highlighter-rouge">180 cm</code>. Values far away from the mean, like <code class="language-plaintext highlighter-rouge">160 cm</code> or <code class="language-plaintext highlighter-rouge">200 cm</code>, happen less often.</p>

<p>That is the intuition behind the Gaussian distribution: a value near the center is common, and a value far from the center is less common. GELU uses this idea when it decides how much of an input value should pass through the activation function.<br />
Remember that in order to have effective LLM parameters, their values should be between 0 and 1. For that purpose, we first used the ReLU function, which replaced negative values with zero. But that was too harsh for LLM parameters, so GELU was introduced.</p>

<p><img src="/assets/gelu.png" alt="GELU activation function" /></p>

<p>Here we compare GELU and ReLU. We can see that neurons that receive small negative values still contribute to the learning process.</p>

<p>Here is a GELU module and an appropriate test:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">GELU</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Gaussian Error Linear Unit activation implemented with Nx tensors.

  GELU is a smooth nonlinear activation commonly used in transformer blocks.
  Instead of zeroing all negative inputs like ReLU, it smoothly gates values
  based on their magnitude. GPT-style models typically use the tanh-based
  approximation rather than the exact error-function version.

  This module mirrors the PyTorch implementation:

      0.5 * x * (1 + torch.tanh(
        torch.sqrt(torch.tensor(2.0 / torch.pi)) *
        (x + 0.044715 * torch.pow(x, 3))
      ))

  The equivalent Nx formula is:

      0.5 * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x^3)))

  The activation is stateless, so the struct only exists to match the
  module-style shape used by other examples in this project.
  """</span>

  <span class="kn">import</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span>

  <span class="k">defstruct</span> <span class="p">[]</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{}</span>

  <span class="nv">@sqrt_2_over_pi</span> <span class="ss">:math</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="mf">2.0</span> <span class="o">/</span> <span class="ss">:math</span><span class="o">.</span><span class="n">pi</span><span class="p">())</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">()</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Creates a stateless GELU activation module.

  ## Example

      iex&gt; gelu = LlmScratch.GELU.new()
      iex&gt; LlmScratch.GELU.forward(gelu, Nx.tensor([0.0])) |&gt; Nx.to_flat_list()
      [0.0]
  """</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{}</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Applies GELU to `x` using a `%LlmScratch.GELU{}` module struct.

  This form is useful when composing modules that all expose a `forward/2`
  function. The input can have any shape; GELU is applied element-wise and the
  output keeps the same shape.

  ## Example

      iex&gt; gelu = LlmScratch.GELU.new()
      iex&gt; x = Nx.tensor([-1.0, 0.0, 1.0])
      iex&gt; LlmScratch.GELU.forward(gelu, x) |&gt; Nx.round(6) |&gt; Nx.to_flat_list()
      [-0.158808, 0.0, 0.841192]
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{},</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">forward</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Applies GELU to `x` directly.

  Since GELU has no parameters, this convenience form produces the same result
  as `forward(LlmScratch.GELU.new(), x)`.

  ## Example

      iex&gt; x = Nx.tensor([[-3.0, 0.0, 3.0]])
      iex&gt; LlmScratch.GELU.forward(x) |&gt; Nx.round(6) |&gt; Nx.to_flat_list()
      [-0.003637, 0.0, 2.996363]
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">forward_defn</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Applies GELU inside `Nx.Defn` code.

  Use this function from other `defn`/`defnp` functions when the activation
  needs to participate in automatic differentiation.

  ## Example

      iex&gt; x = Nx.tensor([[-3.0, 0.0, 3.0]])
      iex&gt; LlmScratch.GELU.forward_defn(x) |&gt; Nx.round(6) |&gt; Nx.to_flat_list()
      [-0.003637, 0.0, 2.996363]
  """</span>
  <span class="n">defn</span> <span class="n">forward_defn</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">x_cubed</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
    <span class="n">inner</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="nv">@sqrt_2_over_pi</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="mf">0.044715</span><span class="p">,</span> <span class="n">x_cubed</span><span class="p">)))</span>

    <span class="n">x</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="mf">0.5</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="mf">1.0</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">inner</span><span class="p">)))</span>
  <span class="k">end</span>
<span class="k">end</span>

</code></pre></div></div>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"gelu applies the GPT approximate activation"</span> <span class="k">do</span>
    <span class="n">gelu</span> <span class="o">=</span> <span class="no">GELU</span><span class="o">.</span><span class="n">new</span><span class="p">()</span>
    <span class="n">input</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([[</span><span class="o">-</span><span class="mf">3.0</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="mf">3.0</span><span class="p">]],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>

    <span class="n">output</span> <span class="o">=</span> <span class="no">GELU</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">gelu</span><span class="p">,</span> <span class="n">input</span><span class="p">)</span>

    <span class="n">assert_close</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="p">[[</span><span class="o">-</span><span class="mf">0.00363743</span><span class="p">,</span> <span class="o">-</span><span class="mf">0.158808</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.841192</span><span class="p">,</span> <span class="mf">2.9963627</span><span class="p">]],</span> <span class="ss">atol:</span> <span class="mf">1.0e-6</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">GELU</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">input</span><span class="p">),</span> <span class="n">output</span><span class="p">,</span> <span class="ss">atol:</span> <span class="mf">1.0e-6</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>And now we are ready to implement a small feed-forward neural network with two linear layers, which is part of our GPT model architecture. The most important property of this neural network is that we first project the input embedding dimension to a first-layer embedding dimension that is bigger by a factor of four. Then we scale the output of the second layer back to the original input dimension:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">FeedForward</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Position-wise feed-forward network used inside GPT transformer blocks.

  This mirrors the book's PyTorch module:

      nn.Sequential(
        nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
        GELU(),
        nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"])
      )

  The first dense layer expands the embedding dimension by a factor of four,
  `LlmScratch.GELU` applies the nonlinearity element-wise, and the second dense
  layer projects the activations back to the original embedding dimension.

  The input can have any number of leading dimensions as long as its final axis
  is `emb_dim`. For GPT hidden states shaped `{batch_size, seq_len, emb_dim}`,
  the output has the same shape.
  """</span>

  <span class="kn">import</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span><span class="no">GELU</span><span class="p">,</span> <span class="no">GPTConfig</span><span class="p">,</span> <span class="no">SelfAttentionV2</span><span class="p">}</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:emb_dim</span><span class="p">,</span> <span class="ss">:layers</span><span class="p">]</span>

  <span class="nv">@type</span> <span class="n">dense_weights</span> <span class="p">::</span> <span class="p">%{</span><span class="ss">kernel:</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="ss">bias:</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()}</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">emb_dim:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">layers:</span> <span class="p">%{</span>
            <span class="ss">first:</span> <span class="n">dense_weights</span><span class="p">(),</span>
            <span class="ss">gelu:</span> <span class="no">GELU</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
            <span class="ss">second:</span> <span class="n">dense_weights</span><span class="p">()</span>
          <span class="p">}</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">()</span> <span class="o">|</span> <span class="n">map</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Creates a feed-forward network from a GPT config.

  `cfg` can be a `%LlmScratch.GPTConfig{}` or a map with either `:emb_dim` or
  `"emb_dim"`.

  ## Options

    * `:seed` - deterministic seed for dense layer initialization.

  ## Example

      iex&gt; cfg = %{emb_dim: 4}
      iex&gt; ff = LlmScratch.FeedForward.new(cfg, seed: 123)
      iex&gt; Nx.shape(ff.layers.first.kernel)
      {4, 16}
      iex&gt; Nx.shape(ff.layers.second.kernel)
      {16, 4}
  """</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">emb_dim</span> <span class="o">=</span> <span class="n">emb_dim_from_cfg!</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">))</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">emb_dim:</span> <span class="n">emb_dim</span><span class="p">,</span>
      <span class="ss">layers:</span> <span class="p">%{</span>
        <span class="ss">first:</span> <span class="no">SelfAttentionV2</span><span class="o">.</span><span class="n">init_dense_weights</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">,</span> <span class="mi">4</span> <span class="o">*</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">seed</span><span class="p">,</span> <span class="no">true</span><span class="p">,</span> <span class="s2">"ff_first"</span><span class="p">),</span>
        <span class="ss">gelu:</span> <span class="no">GELU</span><span class="o">.</span><span class="n">new</span><span class="p">(),</span>
        <span class="ss">second:</span>
          <span class="no">SelfAttentionV2</span><span class="o">.</span><span class="n">init_dense_weights</span><span class="p">(</span><span class="mi">4</span> <span class="o">*</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">,</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span> <span class="no">true</span><span class="p">,</span> <span class="s2">"ff_second"</span><span class="p">)</span>
      <span class="p">}</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Runs the feed-forward network over `x`.

  GELU and both dense layers are applied over the last dimension. The output
  preserves all leading dimensions and restores the last dimension to `emb_dim`.

  Raises `ArgumentError` if the final axis of `x` is not the configured
  embedding dimension.

  ## Example

      iex&gt; ff = LlmScratch.FeedForward.new(%{emb_dim: 4}, seed: 123)
      iex&gt; x = Nx.broadcast(1.0, {2, 3, 4})
      iex&gt; LlmScratch.FeedForward.forward(ff, x) |&gt; Nx.shape()
      {2, 3, 4}
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">feed_forward</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">validate_last_axis!</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">feed_forward</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span>

    <span class="n">x</span>
    <span class="o">|&gt;</span> <span class="n">linear_defn</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">GELU</span><span class="o">.</span><span class="n">forward</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="n">linear_defn</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">call</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="k">def</span> <span class="n">call</span><span class="p">(</span><span class="n">feed_forward</span><span class="p">,</span> <span class="n">x</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">forward</span><span class="p">(</span><span class="n">feed_forward</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Applies a dense layer over the last axis of `x`.

  The dense weights are maps with `:kernel` shaped `{d_in, d_out}` and `:bias`
  shaped `{d_out}`. The function is defined with `Nx.Defn`, so it can be reused
  inside differentiable code such as `Nx.Defn.value_and_grad/2`.

  ## Example

      iex&gt; layer = %{kernel: Nx.broadcast(1.0, {3, 2}), bias: Nx.tensor([0.1, -0.1])}
      iex&gt; x = Nx.tensor([[1.0, 2.0, 3.0]])
      iex&gt; LlmScratch.FeedForward.linear_defn(x, layer) |&gt; Nx.to_flat_list()
      [6.1, 5.9]
  """</span>
  <span class="n">defn</span> <span class="n">linear_defn</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="p">%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">,</span> <span class="ss">bias:</span> <span class="n">bias</span><span class="p">})</span> <span class="k">do</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">dot</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">kernel</span><span class="p">,</span> <span class="p">[</span><span class="mi">0</span><span class="p">]),</span> <span class="n">bias</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">emb_dim_from_cfg!</span><span class="p">(%</span><span class="no">GPTConfig</span><span class="p">{</span><span class="ss">emb_dim:</span> <span class="n">emb_dim</span><span class="p">}),</span> <span class="k">do</span><span class="p">:</span> <span class="n">validate_emb_dim!</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">emb_dim_from_cfg!</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">cfg</span>
    <span class="o">|&gt;</span> <span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="ss">:emb_dim</span><span class="p">,</span> <span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="s2">"emb_dim"</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="n">validate_emb_dim!</span><span class="p">()</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_emb_dim!</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span> <span class="ow">and</span> <span class="n">emb_dim</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="n">emb_dim</span>

  <span class="k">defp</span> <span class="n">validate_emb_dim!</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
          <span class="s2">"expected cfg to contain a positive integer emb_dim, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_last_axis!</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">emb_dim</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">last_axis</span> <span class="o">=</span>
      <span class="n">x</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Tuple</span><span class="o">.</span><span class="n">to_list</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">last</span><span class="p">()</span>

    <span class="k">unless</span> <span class="n">last_axis</span> <span class="o">==</span> <span class="n">emb_dim</span> <span class="k">do</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
            <span class="s2">"expected input last dimension to be </span><span class="si">#{</span><span class="n">emb_dim</span><span class="si">}</span><span class="s2">, got </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">last_axis</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">System</span><span class="o">.</span><span class="n">unique_integer</span><span class="p">([</span><span class="ss">:positive</span><span class="p">])</span>
  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">seed</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">seed</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"seed must be an integer or nil, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And here is a test where we exercise the feed-forward network:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"feed forward expands through gelu and projects back to embedding dimension"</span> <span class="k">do</span>
    <span class="n">cfg</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">4</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">2</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">1</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">feed_forward</span> <span class="o">=</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>
    <span class="n">x</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([[[</span><span class="mf">0.1</span><span class="p">,</span> <span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.3</span><span class="p">,</span> <span class="mf">0.4</span><span class="p">],</span> <span class="p">[</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">1.1</span><span class="p">,</span> <span class="mf">1.2</span><span class="p">,</span> <span class="mf">1.3</span><span class="p">]]],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>

    <span class="n">y</span> <span class="o">=</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">feed_forward</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="o">.</span><span class="n">kernel</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">4</span><span class="p">,</span> <span class="mi">16</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="o">.</span><span class="n">bias</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">16</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="o">.</span><span class="n">kernel</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">16</span><span class="p">,</span> <span class="mi">4</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="o">.</span><span class="n">bias</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">4</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">y</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">4</span><span class="p">}</span>

    <span class="n">assert_close</span><span class="p">(</span><span class="no">FeedForward</span><span class="o">.</span><span class="n">call</span><span class="p">(</span><span class="n">feed_forward</span><span class="p">,</span> <span class="n">x</span><span class="p">),</span> <span class="n">y</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">test</span> <span class="s2">"feed forward preserves GPT-124M hidden state shape with EXLA"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="p">{</span><span class="n">x</span><span class="p">,</span> <span class="n">_key</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">123</span><span class="p">),</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="ss">shape:</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">768</span><span class="p">})</span>
    <span class="n">feed_forward</span> <span class="o">=</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="n">out</span> <span class="o">=</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">feed_forward</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">out</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">768</span><span class="p">}</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="adding-shortcut-connections">Adding Shortcut Connections</h2>

<p>We need shortcut connections between neural network layers to address the problem of vanishing gradients. Gradient values are used to update LLM parameters. Vanishing gradients occur when gradient values approach zero as we move backward through the layers of the LLM, making the training of earlier layers less effective.</p>

<p>To overcome this, we use the shortcut connection technique. Here is an ASCII graph of this concept:</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>x  -&gt; [Layer 1] -&gt; + -&gt; x1 -&gt; [Layer 2] -&gt; + -&gt; x2 -&gt; [Layer 3] -&gt; + -&gt; x3
|                 ^   |                 ^   |                 ^
+-----------------+   +-----------------+   +-----------------+


x1 = Layer1(x)  + x  
x2 = Layer2(x1) + x1  
x3 = Layer3(x2) + x2  
</code></pre></div></div>

<p>For each layer, we add the layer input to the layer output, and the result becomes the input for the next layer.</p>

<p>Here is a well-documented module for an example neural network with five layers, with and without shortcuts. Notice that with shortcuts, gradient values do not vanish toward zero, and all gradients in the network with shortcuts have larger values than gradients in the network without shortcuts.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">ExampleDeepNeuralNetwork</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Small deep network used to demonstrate gradient flow with and without shortcut
  connections.

  This module mirrors the book's PyTorch example:

      nn.ModuleList([
        nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]), GELU()),
        nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]), GELU()),
        nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]), GELU()),
        nn.Sequential(nn.Linear(layer_sizes[3], layer_sizes[4]), GELU()),
        nn.Sequential(nn.Linear(layer_sizes[4], layer_sizes[5]), GELU())
      ])

  The network is intentionally small and fixed to five dense-plus-GELU layers
  because the surrounding chapter example focuses on comparing gradient flow,
  not on building a general model container.

  Unlike PyTorch, Nx tensors do not accumulate gradients in mutable `.grad`
  fields. `backward/3` returns `{loss, gradients}` directly, where `gradients`
  has the same nested structure as the model layers.

  Shortcut behavior follows the example's shape check. For `use_shortcut: true`,
  the first four layers in `layer_sizes = [3, 3, 3, 3, 3, 1]` add their output
  back to the input because their input and output shapes match. The final
  `3 -&gt; 1` layer cannot use a shortcut and is applied normally.
  """</span>

  <span class="kn">import</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Defn</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span><span class="no">FeedForward</span><span class="p">,</span> <span class="no">GELU</span><span class="p">,</span> <span class="no">SelfAttentionV2</span><span class="p">,</span> <span class="no">SimpleGradient</span><span class="p">}</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:layers</span><span class="p">,</span> <span class="ss">:use_shortcut</span><span class="p">]</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Builds a five-layer dense network.

  ## Parameters

    * `layer_sizes` - six positive integers describing the input and output
      width of each dense layer. For example, `[3, 3, 3, 3, 3, 1]` creates the
      layers `3 -&gt; 3`, `3 -&gt; 3`, `3 -&gt; 3`, `3 -&gt; 3`, and `3 -&gt; 1`.

    * `opts` - keyword options:
      `:use_shortcut` is required and controls residual connections.
      `:seed` is optional and defaults to `123`.

  Dense kernels and biases are initialized through Axon so their layout and
  initialization behavior match the rest of this project.

  ## Example

      iex&gt; model = LlmScratch.ExampleDeepNeuralNetwork.new([3, 3, 3, 3, 3, 1],
      ...&gt;   use_shortcut: false,
      ...&gt;   seed: 123
      ...&gt; )
      iex&gt; tuple_size(model.layers)
      5
      iex&gt; Nx.shape(elem(model.layers, 0).kernel)
      {3, 3}
      iex&gt; Nx.shape(elem(model.layers, 4).kernel)
      {3, 1}
  """</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">layer_sizes</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">use_shortcut</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">fetch!</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:use_shortcut</span><span class="p">)</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">,</span> <span class="mi">123</span><span class="p">)</span>

    <span class="n">layers</span> <span class="o">=</span>
      <span class="n">layer_sizes</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">chunk_every</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="ss">:discard</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">with_index</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="p">{[</span><span class="n">d_in</span><span class="p">,</span> <span class="n">d_out</span><span class="p">],</span> <span class="n">idx</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="no">SelfAttentionV2</span><span class="o">.</span><span class="n">init_dense_weights</span><span class="p">(</span>
          <span class="n">d_in</span><span class="p">,</span>
          <span class="n">d_out</span><span class="p">,</span>
          <span class="n">seed</span> <span class="o">+</span> <span class="n">idx</span><span class="p">,</span>
          <span class="no">true</span><span class="p">,</span>
          <span class="s2">"example_layer_</span><span class="si">#{</span><span class="n">idx</span><span class="si">}</span><span class="s2">"</span>
        <span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">List</span><span class="o">.</span><span class="n">to_tuple</span><span class="p">()</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span><span class="ss">layers:</span> <span class="n">layers</span><span class="p">,</span> <span class="ss">use_shortcut:</span> <span class="n">use_shortcut</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Runs the network on `x`.

  `x` is expected to be a rank-2 tensor whose second axis matches the first
  layer input width. The returned tensor has the output width of the final
  layer.

  When `use_shortcut` is enabled, residual additions are applied for layers
  whose dense-plus-GELU output shape matches the incoming `x` shape.

  ## Example

      iex&gt; model = LlmScratch.ExampleDeepNeuralNetwork.new([3, 3, 3, 3, 3, 1],
      ...&gt;   use_shortcut: true,
      ...&gt;   seed: 123
      ...&gt; )
      iex&gt; x = Nx.tensor([[1.0, 0.0, -1.0]], type: {:f, 32})
      iex&gt; LlmScratch.ExampleDeepNeuralNetwork.forward(model, x) |&gt; Nx.shape()
      {1, 1}
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{</span><span class="ss">layers:</span> <span class="n">layers</span><span class="p">,</span> <span class="ss">use_shortcut:</span> <span class="no">false</span><span class="p">},</span> <span class="n">x</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">forward_without_shortcut</span><span class="p">(</span><span class="n">layers</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{</span><span class="ss">layers:</span> <span class="n">layers</span><span class="p">,</span> <span class="ss">use_shortcut:</span> <span class="no">true</span><span class="p">},</span> <span class="n">x</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">forward_with_shortcut</span><span class="p">(</span><span class="n">layers</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Computes mean-squared-error loss and gradients for the model layers.

  This is the Nx equivalent of a PyTorch forward pass followed by
  `loss.backward()`. Instead of mutating parameter gradients, it returns them:

      {loss, gradients} = backward(model, x, target)

  `loss` is a scalar tensor. `gradients` is a tuple with one map per layer. Each
  layer gradient map contains `:kernel` and `:bias` tensors matching the
  corresponding layer parameter shapes.
  """</span>
  <span class="k">def</span> <span class="n">backward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{</span><span class="ss">layers:</span> <span class="n">layers</span><span class="p">,</span> <span class="ss">use_shortcut:</span> <span class="no">false</span><span class="p">},</span> <span class="n">x</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">backward_without_shortcut</span><span class="p">(</span><span class="n">layers</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

  <span class="k">def</span> <span class="n">backward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{</span><span class="ss">layers:</span> <span class="n">layers</span><span class="p">,</span> <span class="ss">use_shortcut:</span> <span class="no">true</span><span class="p">},</span> <span class="n">x</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">backward_with_shortcut</span><span class="p">(</span><span class="n">layers</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

  <span class="nv">@doc</span> <span class="sd">"""
  Returns the absolute mean of each layer's kernel gradient.

  This mirrors the Python example's:

      param.grad.abs().mean().item()

  Bias gradients are intentionally ignored because the chapter example prints
  only gradients for parameters whose names include `"weight"`.
  """</span>
  <span class="k">def</span> <span class="n">weight_gradient_means</span><span class="p">(</span><span class="n">gradients</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">gradients</span>
    <span class="o">|&gt;</span> <span class="no">Tuple</span><span class="o">.</span><span class="n">to_list</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">map</span><span class="p">(</span><span class="k">fn</span> <span class="p">%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}</span> <span class="o">-&gt;</span>
      <span class="n">kernel</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">abs</span><span class="p">()</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">defn</span> <span class="n">forward_without_shortcut</span><span class="p">({</span><span class="n">layer_0</span><span class="p">,</span> <span class="n">layer_1</span><span class="p">,</span> <span class="n">layer_2</span><span class="p">,</span> <span class="n">layer_3</span><span class="p">,</span> <span class="n">layer_4</span><span class="p">},</span> <span class="n">x</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">x</span>
    <span class="o">|&gt;</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">layer_0</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">layer_1</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">layer_2</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">layer_3</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">layer_4</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">defn</span> <span class="n">forward_with_shortcut</span><span class="p">({</span><span class="n">layer_0</span><span class="p">,</span> <span class="n">layer_1</span><span class="p">,</span> <span class="n">layer_2</span><span class="p">,</span> <span class="n">layer_3</span><span class="p">,</span> <span class="n">layer_4</span><span class="p">},</span> <span class="n">x</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">x</span>
    <span class="o">|&gt;</span> <span class="n">residual_linear_gelu</span><span class="p">(</span><span class="n">layer_0</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">residual_linear_gelu</span><span class="p">(</span><span class="n">layer_1</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">residual_linear_gelu</span><span class="p">(</span><span class="n">layer_2</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">residual_linear_gelu</span><span class="p">(</span><span class="n">layer_3</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">layer_4</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">defn</span> <span class="n">backward_without_shortcut</span><span class="p">(</span><span class="n">layers</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">value_and_grad</span><span class="p">(</span><span class="n">layers</span><span class="p">,</span> <span class="k">fn</span> <span class="n">layers</span> <span class="o">-&gt;</span>
      <span class="n">layers</span>
      <span class="o">|&gt;</span> <span class="n">forward_without_shortcut</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">SimpleGradient</span><span class="o">.</span><span class="n">mse_loss</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">defn</span> <span class="n">backward_with_shortcut</span><span class="p">(</span><span class="n">layers</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">value_and_grad</span><span class="p">(</span><span class="n">layers</span><span class="p">,</span> <span class="k">fn</span> <span class="n">layers</span> <span class="o">-&gt;</span>
      <span class="n">layers</span>
      <span class="o">|&gt;</span> <span class="n">forward_with_shortcut</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">SimpleGradient</span><span class="o">.</span><span class="n">mse_loss</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">defnp</span> <span class="n">residual_linear_gelu</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">layer</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">x</span> <span class="o">+</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">layer</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">defnp</span> <span class="n">linear_gelu</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">layer</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">x</span>
    <span class="o">|&gt;</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">linear_defn</span><span class="p">(</span><span class="n">layer</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">GELU</span><span class="o">.</span><span class="n">forward_defn</span><span class="p">()</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And tests that demonstate gradient values without and with using shortcuts:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"deep neural network asserts layer weight gradient means without shortcuts"</span> <span class="k">do</span>
    <span class="n">layer_sizes</span> <span class="o">=</span> <span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span>
    <span class="n">sample_input</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([[</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.0</span><span class="p">]],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>
    <span class="n">target</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([[</span><span class="mf">0.0</span><span class="p">]],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">layer_sizes</span><span class="p">,</span> <span class="ss">use_shortcut:</span> <span class="no">false</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="p">{</span><span class="n">loss</span><span class="p">,</span> <span class="n">gradients</span><span class="p">}</span> <span class="o">=</span> <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">backward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">sample_input</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
    <span class="n">gradient_means</span> <span class="o">=</span> <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">weight_gradient_means</span><span class="p">(</span><span class="n">gradients</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">sample_input</span><span class="p">))</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">}</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">loss</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">2.9010625e-6</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">}),</span> <span class="ss">atol:</span> <span class="mf">1.0e-12</span><span class="p">)</span>

    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">1.9688005e-5</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-10</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">8.4962267e-6</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-10</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">1.0740861e-5</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-10</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">1.0382744e-5</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-10</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">4</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">5.419739e-6</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-10</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="n">test</span> <span class="s2">"deep neural network asserts larger gradient flow with shortcuts"</span> <span class="k">do</span>
    <span class="n">layer_sizes</span> <span class="o">=</span> <span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span>
    <span class="n">sample_input</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([[</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.0</span><span class="p">]],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>
    <span class="n">target</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">([[</span><span class="mf">0.0</span><span class="p">]],</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">})</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">layer_sizes</span><span class="p">,</span> <span class="ss">use_shortcut:</span> <span class="no">true</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="p">{</span><span class="n">loss</span><span class="p">,</span> <span class="n">gradients</span><span class="p">}</span> <span class="o">=</span> <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">backward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">sample_input</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
    <span class="n">gradient_means</span> <span class="o">=</span> <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">weight_gradient_means</span><span class="p">(</span><span class="n">gradients</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">sample_input</span><span class="p">))</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">}</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="n">loss</span><span class="p">,</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.02248338</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:f</span><span class="p">,</span> <span class="mi">32</span><span class="p">}),</span> <span class="ss">atol:</span> <span class="mf">1.0e-8</span><span class="p">)</span>

    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.0034418134</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-8</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.008453862</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-8</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.0049184095</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-8</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.003547175</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-8</span><span class="p">)</span>
    <span class="n">assert_close</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="mi">4</span><span class="p">),</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.00962014</span><span class="p">),</span> <span class="ss">atol:</span> <span class="mf">1.0e-8</span><span class="p">)</span>

    <span class="n">without_shortcut</span> <span class="o">=</span>
      <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">layer_sizes</span><span class="p">,</span> <span class="ss">use_shortcut:</span> <span class="no">false</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="p">{</span><span class="n">_loss</span><span class="p">,</span> <span class="n">without_shortcut_gradients</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">backward</span><span class="p">(</span><span class="n">without_shortcut</span><span class="p">,</span> <span class="n">sample_input</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

    <span class="n">without_shortcut_gradient_means</span> <span class="o">=</span>
      <span class="no">ExampleDeepNeuralNetwork</span><span class="o">.</span><span class="n">weight_gradient_means</span><span class="p">(</span><span class="n">without_shortcut_gradients</span><span class="p">)</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">zip</span><span class="p">(</span><span class="n">gradient_means</span><span class="p">,</span> <span class="n">without_shortcut_gradient_means</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">each</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">with_shortcut_mean</span><span class="p">,</span> <span class="n">without_shortcut_mean</span><span class="p">}</span> <span class="o">-&gt;</span>
      <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">greater</span><span class="p">(</span><span class="n">with_shortcut_mean</span><span class="p">,</span> <span class="n">without_shortcut_mean</span><span class="p">)</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span> <span class="o">==</span> <span class="mi">1</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="transformer-block">Transformer block</h2>

<p>Let’s see how all the components that we have built so far fit into a transformer block, the fundamental building block of the GPT LLM architecture.</p>

<div class="language-plaintext highlighter-rouge"><div class="highlight"><pre class="highlight"><code>                     [[-0.0256, ..., 0.6890],
                      [-0.0178, ..., 0.7431],
                      [ 0.4558, ..., 0.7814],
                      [ 0.0702, ..., 0.7134]]
                                  ^
                                  |
                              +---+---+
                              |   +   | &lt;-----------------------------+
                              +---+---+                               |
                                  ^                                   |
                                  |                                   |
                          +-------+-------+                           |
                          |    Dropout    |                           |
                          +-------+-------+                           |
                                  ^                                   |
                                  |                                   |
                          +-------+-------+                           |
                          | Feed forward  |                           |
                          +-------+-------+                           |
                                  ^                                   |
                                  |                                   |
                          +-------+-------+                           |
                          | LayerNorm 2   |                           |
                          +-------+-------+                           |  
                                  ^                                   |
                                  |-----------------------------------+
                              +---+---+
                              |   +   | &lt;-----------------------------+
                              +---+---+                               |
                                  ^                                   |
                                  |                                   |
                          +-------+-------+                           |
                          |    Dropout    |                           |
                          +-------+-------+                           |
                                  ^                                   |
                                  |                                   |
                          +-------+-------+                           |
                          | Masked multi- |                           |
                          | head attention|                           |
                          +-------+-------+                           |
                                  ^                                   |
                                  |                                   |
                          +-------+-------+                           |
                          | LayerNorm 1   |                           |
                          +-------+-------+                           |
                                  ^                                   |
                                  | ----------------------------------+
                     [[0.2961, ..., 0.4604],
                      [0.2238, ..., 0.7598],
                      [0.6945, ..., 0.5963],
                      [0.0890, ..., 0.5833]]

                                                    +---------------------------+
                                                    |   Feed forward block      |
                                                    |                           |
                                                    |      +--------------+     |
                                                    |      | Linear layer |     |
                                                    |      +------+-------+     |
                                                    |             ^             |
                                                    |             |             |
                                                    |      +------+-------+     |
                                                    |      | GELU activation    |
                                                    |      +------+-------+     |
                                                    |             ^             |
                                                    |             |             |
                                                    |      +------+-------+     |
                                                    |      | Linear layer |     |
                                                    |      +--------------+     |
                                                    +---------------------------+
</code></pre></div></div>

<p>Each token has a dimension of 768, based on the configuration that we are using. The transformer block has these properties:</p>

<ul>
  <li>the output has the same dimensions as the input</li>
  <li>the multi-head attention block identifies relationships between tokens based on the provided input</li>
  <li>the feed-forward block changes each token representation individually</li>
  <li>layer normalization is applied before the multi-head attention and feed-forward blocks, as this provides better training dynamics than applying it after those blocks; this is called pre-layer normalization</li>
</ul>

<p>The ASCII art is self-descriptive, so let’s code this transformer block as an Elixir module. Note that we changed the multi-head attention module by making <code class="language-plaintext highlighter-rouge">maybe_dropout</code> a public function so we can reuse it in the transformer block.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">TransformerBlock</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Real GPT-style transformer block implemented with Nx tensors.

  This mirrors the book's PyTorch `TransformerBlock`:

      shortcut = x
      x = norm1(x)
      x = attention(x)
      x = dropout(x)
      x = x + shortcut

      shortcut = x
      x = norm2(x)
      x = feed_forward(x)
      x = dropout(x)
      x = x + shortcut

  The block uses pre-normalization, causal multi-head self-attention, a
  position-wise feed-forward network, and residual shortcut connections.
  """</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span>
    <span class="no">DummyLayerNorm</span><span class="p">,</span>
    <span class="no">FeedForward</span><span class="p">,</span>
    <span class="no">GPTConfig</span><span class="p">,</span>
    <span class="no">MultiheadAttention</span>
  <span class="p">}</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:att</span><span class="p">,</span> <span class="ss">:ff</span><span class="p">,</span> <span class="ss">:norm1</span><span class="p">,</span> <span class="ss">:norm2</span><span class="p">,</span> <span class="ss">:drop_shortcut</span><span class="p">,</span> <span class="ss">:cfg</span><span class="p">]</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">att:</span> <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">ff:</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">norm1:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">norm2:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">drop_shortcut:</span> <span class="n">float</span><span class="p">(),</span>
          <span class="ss">cfg:</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">()</span> <span class="o">|</span> <span class="n">map</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">()</span> <span class="o">|</span> <span class="n">map</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Creates a GPT transformer block from a config.

  `cfg` can be a `%LlmScratch.GPTConfig{}` or a map with atom/string keys:
  `:emb_dim`, `:context_length`, `:n_heads`, `:drop_rate`, and `:qkv_bias`.

  ## Options

    * `:seed` - deterministic seed used for attention and feed-forward
      initialization.
    * `:norm_eps` - epsilon for both layer-normalization modules.
  """</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">params</span> <span class="o">=</span> <span class="n">params_from_cfg!</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">))</span>
    <span class="n">norm_eps</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:norm_eps</span><span class="p">,</span> <span class="mf">1.0e-5</span><span class="p">)</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">cfg:</span> <span class="n">cfg</span><span class="p">,</span>
      <span class="ss">att:</span>
        <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">new</span><span class="p">(</span>
          <span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">n_heads</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">,</span>
          <span class="ss">seed:</span> <span class="n">seed</span>
        <span class="p">),</span>
      <span class="ss">ff:</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
      <span class="ss">norm1:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">eps:</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">norm2:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">eps:</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">drop_shortcut:</span> <span class="n">params</span><span class="o">.</span><span class="n">drop_rate</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Runs a forward pass over hidden states shaped `{batch_size, seq_len, emb_dim}`.

  Forward options:

    * `:mode` - `:train` or `:inference`. Dropout is skipped in inference mode.
    * `:key` - optional `Nx.Random` key used by dropout in train mode.
    * `:attention_key`, `:shortcut_key1`, `:shortcut_key2` - optional keys for
      controlling individual dropout calls.
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">block</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">shortcut</span> <span class="o">=</span> <span class="n">x</span>

    <span class="n">x</span> <span class="o">=</span>
      <span class="n">block</span><span class="o">.</span><span class="n">norm1</span>
      <span class="o">|&gt;</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="p">,</span> <span class="nv">&amp;1</span><span class="p">,</span> <span class="n">attention_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">)))</span>
      <span class="o">|&gt;</span> <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">maybe_dropout</span><span class="p">(</span>
        <span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
        <span class="n">shortcut_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:shortcut_key1</span><span class="p">)</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">shortcut</span><span class="p">)</span>

    <span class="n">shortcut</span> <span class="o">=</span> <span class="n">x</span>

    <span class="n">block</span><span class="o">.</span><span class="n">norm2</span>
    <span class="o">|&gt;</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="no">FeedForward</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="p">,</span> <span class="nv">&amp;1</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">maybe_dropout</span><span class="p">(</span>
      <span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span><span class="p">),</span>
      <span class="n">shortcut_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:shortcut_key2</span><span class="p">)</span>
    <span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">shortcut</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">call</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Alias for `forward/3` using default forward options.

  This mirrors the callable-style API used by other modules in the project.
  Dropout is skipped because `forward/3` defaults to inference mode.
  """</span>
  <span class="k">def</span> <span class="n">call</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">x</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">forward</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">attention_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">opts</span>
    <span class="o">|&gt;</span> <span class="n">opts_with_default_mode</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="ss">:key</span><span class="p">,</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:attention_key</span><span class="p">,</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:key</span><span class="p">)))</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reject</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">_key</span><span class="p">,</span> <span class="n">value</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="n">is_nil</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">shortcut_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="n">key_name</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">opts</span>
    <span class="o">|&gt;</span> <span class="n">opts_with_default_mode</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="ss">:key</span><span class="p">,</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="n">key_name</span><span class="p">,</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:key</span><span class="p">)))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">seed</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="p">%{</span><span class="ss">dropout:</span> <span class="n">block</span><span class="o">.</span><span class="n">drop_shortcut</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span><span class="p">}</span>

  <span class="k">defp</span> <span class="n">opts_with_default_mode</span><span class="p">(</span><span class="n">opts</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="p">[</span><span class="ss">:mode</span><span class="p">]),</span> <span class="ss">:mode</span><span class="p">,</span> <span class="n">mode!</span><span class="p">(</span><span class="n">opts</span><span class="p">))</span>

  <span class="k">defp</span> <span class="n">mode!</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:mode</span><span class="p">,</span> <span class="ss">:inference</span><span class="p">)</span> <span class="k">do</span>
      <span class="n">mode</span> <span class="ow">when</span> <span class="n">mode</span> <span class="ow">in</span> <span class="p">[</span><span class="ss">:train</span><span class="p">,</span> <span class="ss">:inference</span><span class="p">]</span> <span class="o">-&gt;</span> <span class="n">mode</span>
      <span class="n">mode</span> <span class="o">-&gt;</span> <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"mode must be :train or :inference, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">mode</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">params_from_cfg!</span><span class="p">(%</span><span class="no">GPTConfig</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span>
      <span class="ss">emb_dim:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:emb_dim</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">),</span>
      <span class="ss">context_length:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:context_length</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">),</span>
      <span class="ss">n_heads:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:n_heads</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">n_heads</span><span class="p">),</span>
      <span class="ss">drop_rate:</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">),</span>
      <span class="ss">qkv_bias:</span> <span class="n">validate_boolean!</span><span class="p">(</span><span class="ss">:qkv_bias</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">params_from_cfg!</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span>
      <span class="ss">emb_dim:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:emb_dim</span><span class="p">),</span>
      <span class="ss">context_length:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:context_length</span><span class="p">),</span>
      <span class="ss">n_heads:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:n_heads</span><span class="p">),</span>
      <span class="ss">drop_rate:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:drop_rate</span><span class="p">),</span>
      <span class="ss">qkv_bias:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:qkv_bias</span><span class="p">)</span>
    <span class="p">}</span>
    <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="k">fn</span> <span class="n">params</span> <span class="o">-&gt;</span>
      <span class="p">%{</span>
        <span class="ss">emb_dim:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:emb_dim</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">),</span>
        <span class="ss">context_length:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:context_length</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">context_length</span><span class="p">),</span>
        <span class="ss">n_heads:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:n_heads</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">n_heads</span><span class="p">),</span>
        <span class="ss">drop_rate:</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">params</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">),</span>
        <span class="ss">qkv_bias:</span> <span class="n">validate_boolean!</span><span class="p">(</span><span class="ss">:qkv_bias</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span>
      <span class="p">}</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="no">Atom</span><span class="o">.</span><span class="n">to_string</span><span class="p">(</span><span class="n">key</span><span class="p">)))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="n">_field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="n">value</span>

  <span class="k">defp</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="n">field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected </span><span class="si">#{</span><span class="n">field</span><span class="si">}</span><span class="s2"> to be a positive integer, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_number</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&lt;</span> <span class="mi">1</span><span class="p">,</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">value</span> <span class="o">*</span> <span class="mf">1.0</span>

  <span class="k">defp</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected drop_rate to be a number in [0, 1), got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_boolean!</span><span class="p">(</span><span class="n">_field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_boolean</span><span class="p">(</span><span class="n">value</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">value</span>

  <span class="k">defp</span> <span class="n">validate_boolean!</span><span class="p">(</span><span class="n">field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected </span><span class="si">#{</span><span class="n">field</span><span class="si">}</span><span class="s2"> to be a boolean, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">System</span><span class="o">.</span><span class="n">unique_integer</span><span class="p">([</span><span class="ss">:positive</span><span class="p">])</span>
  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">seed</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">seed</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"seed must be an integer or nil, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And test that is using transformer block:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"transformer block preserves GPT-124M hidden state shape"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="p">{</span><span class="n">x</span><span class="p">,</span> <span class="n">_key</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Random</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">123</span><span class="p">),</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="ss">shape:</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">768</span><span class="p">})</span>
    <span class="n">block</span> <span class="o">=</span> <span class="no">TransformerBlock</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>
    <span class="n">output</span> <span class="o">=</span> <span class="no">TransformerBlock</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">768</span><span class="p">}</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">output</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">768</span><span class="p">}</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="coding-the-gpt-model">Coding the GPT model</h2>

<p>Let’s turn our <code class="language-plaintext highlighter-rouge">DummyGPTModel</code> into a real <code class="language-plaintext highlighter-rouge">GPTModel</code>. Here is the full <code class="language-plaintext highlighter-rouge">GPTModel</code> architecture:</p>

<div class="language-plaintext highlighter-rouge"><div class="highlight"><pre class="highlight"><code>                              A 4 x 50,257-dimensional tensor

                     [[-0.0055, ..., -0.4747],
                      [ 0.2663, ..., -0.4224],
                      [ 1.1146, ...,  0.0276],
                      [-0.8239, ..., -0.3993]]
                                  ^
                                  |
                +-----------------+-----------------+
                |                                   |
                |             GPT model             |
                |                                   |
                |        +-----------------+        |
                |        | Linear output   |        |
                |        | layer           |        |
                |        +--------+--------+        |
                |                 ^                 |
                |                 |                 |
                |        +--------+--------+        |
                |        | Final LayerNorm |        |
                |        +--------+--------+        |
                |                 ^                 |
                |                 |                 |
                |        +--------+--------+        |
                |        | Transformer     |        |
                |        | block           |        |
                |        |                 |        |
                |        | repeated 12x    |        |
                |        +--------+--------+        |
                |                 ^                 |
                |                 |                 |
                |        +--------+--------+        |
                |        | Dropout         |        |
                |        +--------+--------+        |
                |                 ^                 |
                |                 |                 |
                |        +--------+--------+        |
                |        | Positional      |        |
                |        | embedding layer |        |
                |        +--------+--------+        |
                |                 ^                 |
                |                 |                 |
                |        +--------+--------+        |
                |        | Token embedding |        |
                |        | layer           |        |
                |        +--------+--------+        |
                |                                   |
                +-----------------+-----------------+
                                  ^
                                  |
                        +---------+---------+
                        | Tokenized text    |
                        +---------+---------+
                                  ^
                                  |
                     Every  effort  moves  you
</code></pre></div></div>

<p>Here is our Elixir code for <code class="language-plaintext highlighter-rouge">GPTModel</code>:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">GPTModel</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  GPT model implemented with Nx tensors and the real transformer block.

  This mirrors the book's PyTorch `GPTModel`:

      tok_emb = nn.Embedding(vocab_size, emb_dim)
      pos_emb = nn.Embedding(context_length, emb_dim)
      drop_emb = nn.Dropout(drop_rate)
      trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(n_layers)])
      final_norm = LayerNorm(emb_dim)
      out_head = nn.Linear(emb_dim, vocab_size, bias=False)

  `LlmScratch.DummyGPTModel` remains unchanged; this module uses
  `LlmScratch.TransformerBlock` for the actual attention and feed-forward stack.
  """</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span>
    <span class="no">DummyLayerNorm</span><span class="p">,</span>
    <span class="no">EmbeddingNative</span><span class="p">,</span>
    <span class="no">GPTConfig</span><span class="p">,</span>
    <span class="no">MultiheadAttention</span><span class="p">,</span>
    <span class="no">SelfAttentionV2</span><span class="p">,</span>
    <span class="no">TransformerBlock</span>
  <span class="p">}</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:cfg</span><span class="p">,</span> <span class="ss">:tok_emb</span><span class="p">,</span> <span class="ss">:pos_emb</span><span class="p">,</span> <span class="ss">:drop_emb</span><span class="p">,</span> <span class="ss">:trf_blocks</span><span class="p">,</span> <span class="ss">:final_norm</span><span class="p">,</span> <span class="ss">:out_head</span><span class="p">]</span>

  <span class="nv">@type</span> <span class="n">linear_no_bias</span> <span class="p">::</span> <span class="p">%{</span><span class="ss">kernel:</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()}</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">cfg:</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">tok_emb:</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">pos_emb:</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">drop_emb:</span> <span class="n">float</span><span class="p">(),</span>
          <span class="ss">trf_blocks:</span> <span class="p">[</span><span class="no">TransformerBlock</span><span class="o">.</span><span class="n">t</span><span class="p">()],</span>
          <span class="ss">final_norm:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">out_head:</span> <span class="n">linear_no_bias</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Creates a GPT model from a `%LlmScratch.GPTConfig{}`.

  ## Options

    * `:seed` - deterministic seed used for embeddings, transformer blocks,
      and the output projection.
    * `:norm_eps` - epsilon used by the final layer norm and transformer block
      layer norms.
  """</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(%</span><span class="no">GPTConfig</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">))</span>
    <span class="n">norm_eps</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:norm_eps</span><span class="p">,</span> <span class="mf">1.0e-5</span><span class="p">)</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">cfg:</span> <span class="n">cfg</span><span class="p">,</span>
      <span class="ss">tok_emb:</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span><span class="p">),</span>
      <span class="ss">pos_emb:</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
      <span class="ss">drop_emb:</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">),</span>
      <span class="ss">trf_blocks:</span> <span class="n">transformer_blocks</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span><span class="p">,</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">final_norm:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">eps:</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">out_head:</span> <span class="n">init_out_head</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span> <span class="o">+</span> <span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Runs a forward pass over token ids shaped `{batch_size, seq_len}`.

  The returned logits have shape `{batch_size, seq_len, vocab_size}`.

  Forward options:

    * `:mode` - `:inference` or `:train`. Defaults to `:inference`.
    * `:key` - optional `Nx.Random` key used by dropout in train mode.
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">in_idx</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">}</span> <span class="o">=</span> <span class="n">validate_input_shape!</span><span class="p">(</span><span class="n">in_idx</span><span class="p">)</span>
    <span class="n">validate_context_length!</span><span class="p">(</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">model</span><span class="o">.</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">)</span>

    <span class="n">tok_embeds</span> <span class="o">=</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">tok_emb</span><span class="p">,</span> <span class="n">in_idx</span><span class="p">)</span>

    <span class="n">pos_embeds</span> <span class="o">=</span>
      <span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span>
      <span class="o">|&gt;</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">positional_indices</span><span class="p">(</span><span class="n">seq_len</span><span class="p">))</span>

    <span class="n">x</span> <span class="o">=</span>
      <span class="n">tok_embeds</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">pos_embeds</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">maybe_dropout</span><span class="p">(</span>
        <span class="p">%{</span><span class="ss">dropout:</span> <span class="n">model</span><span class="o">.</span><span class="n">drop_emb</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span><span class="o">.</span><span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">},</span>
        <span class="n">dropout_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span>
      <span class="p">)</span>

    <span class="n">x</span> <span class="o">=</span>
      <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="k">fn</span> <span class="n">block</span><span class="p">,</span> <span class="n">acc</span> <span class="o">-&gt;</span>
        <span class="no">TransformerBlock</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">acc</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
      <span class="k">end</span><span class="p">)</span>

    <span class="n">x</span> <span class="o">=</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">final_norm</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

    <span class="n">linear</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">call</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Alias for `forward/3` using default forward options.
  """</span>
  <span class="k">def</span> <span class="n">call</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">in_idx</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">in_idx</span><span class="p">)</span>

  <span class="nv">@spec</span> <span class="n">total_parameters</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Counts all trainable parameters in the GPT model.

  The count includes token embeddings, positional embeddings, every
  transformer block, the final layer norm, and the output projection.
  """</span>
  <span class="k">def</span> <span class="n">total_parameters</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">model</span><span class="o">.</span><span class="n">tok_emb</span><span class="o">.</span><span class="n">weight</span>
    <span class="o">|&gt;</span> <span class="n">tensor_parameters</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">tensor_parameters</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span><span class="o">.</span><span class="n">weight</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;2</span> <span class="o">+</span> <span class="n">transformer_block_parameters</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">))))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">layer_norm_parameters</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">final_norm</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">tensor_parameters</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="o">.</span><span class="n">kernel</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">transformer_block_parameters</span><span class="p">(</span><span class="n">map</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Counts parameters in one transformer block.
  """</span>
  <span class="k">def</span> <span class="n">transformer_block_parameters</span><span class="p">(</span><span class="n">block</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">attention_parameters</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">feed_forward_parameters</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">layer_norm_parameters</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">norm1</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">layer_norm_parameters</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">norm2</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">attention_parameters</span><span class="p">(</span><span class="n">map</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Counts parameters in a multi-head attention module.
  """</span>
  <span class="k">def</span> <span class="n">attention_parameters</span><span class="p">(</span><span class="n">attention</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">attention</span><span class="o">.</span><span class="n">w_q</span>
    <span class="o">|&gt;</span> <span class="n">dense_parameters</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">dense_parameters</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_k</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">dense_parameters</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_v</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">dense_parameters</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">out_proj</span><span class="p">,</span> <span class="no">true</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">feed_forward_parameters</span><span class="p">(</span><span class="n">map</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Counts parameters in a GPT feed-forward module.
  """</span>
  <span class="k">def</span> <span class="n">feed_forward_parameters</span><span class="p">(</span><span class="n">feed_forward</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">dense_parameters</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">dense_parameters</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">layer_norm_parameters</span><span class="p">(</span><span class="n">map</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Counts parameters in a layer norm module.
  """</span>
  <span class="k">def</span> <span class="n">layer_norm_parameters</span><span class="p">(</span><span class="n">layer_norm</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">tensor_parameters</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">scale</span><span class="p">)</span> <span class="o">+</span> <span class="n">tensor_parameters</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">shift</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">tensor_parameters</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Counts scalar values in a tensor.
  """</span>
  <span class="k">def</span> <span class="n">tensor_parameters</span><span class="p">(</span><span class="n">tensor</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">transformer_blocks</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">_seed</span><span class="p">,</span> <span class="n">_norm_eps</span><span class="p">)</span> <span class="ow">when</span> <span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">==</span> <span class="mi">0</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="p">[]</span>

  <span class="k">defp</span> <span class="n">transformer_blocks</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span><span class="p">,</span> <span class="n">norm_eps</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">for</span> <span class="n">layer_idx</span> <span class="o">&lt;-</span> <span class="mi">0</span><span class="o">..</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span> <span class="k">do</span>
      <span class="no">TransformerBlock</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span> <span class="o">+</span> <span class="n">layer_idx</span><span class="p">,</span> <span class="ss">norm_eps:</span> <span class="n">norm_eps</span><span class="p">)</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">init_out_head</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}</span> <span class="o">=</span>
      <span class="no">SelfAttentionV2</span><span class="o">.</span><span class="n">init_dense_weights</span><span class="p">(</span>
        <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span>
        <span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span>
        <span class="n">seed</span><span class="p">,</span>
        <span class="no">false</span><span class="p">,</span>
        <span class="s2">"out_head"</span>
      <span class="p">)</span>

    <span class="p">%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">linear</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="p">%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">}),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">dot</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">kernel</span><span class="p">,</span> <span class="p">[</span><span class="mi">0</span><span class="p">])</span>

  <span class="k">defp</span> <span class="n">dense_parameters</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">,</span> <span class="ss">bias:</span> <span class="n">bias</span><span class="p">},</span> <span class="no">true</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">tensor_parameters</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span> <span class="o">+</span> <span class="n">tensor_parameters</span><span class="p">(</span><span class="n">bias</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">dense_parameters</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">},</span> <span class="no">false</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">tensor_parameters</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">positional_indices</span><span class="p">(</span><span class="n">seq_len</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">iota</span><span class="p">({</span><span class="n">seq_len</span><span class="p">},</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span>

  <span class="k">defp</span> <span class="n">dropout_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">opts</span>
    <span class="o">|&gt;</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">take</span><span class="p">([</span><span class="ss">:mode</span><span class="p">,</span> <span class="ss">:key</span><span class="p">])</span>
    <span class="o">|&gt;</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="ss">:mode</span><span class="p">,</span> <span class="n">mode!</span><span class="p">(</span><span class="n">opts</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">mode!</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:mode</span><span class="p">,</span> <span class="ss">:inference</span><span class="p">)</span> <span class="k">do</span>
      <span class="n">mode</span> <span class="ow">when</span> <span class="n">mode</span> <span class="ow">in</span> <span class="p">[</span><span class="ss">:train</span><span class="p">,</span> <span class="ss">:inference</span><span class="p">]</span> <span class="o">-&gt;</span> <span class="n">mode</span>
      <span class="n">mode</span> <span class="o">-&gt;</span> <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"mode must be :train or :inference, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">mode</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_input_shape!</span><span class="p">(</span><span class="n">in_idx</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">in_idx</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="p">{</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">}</span>

      <span class="n">shape</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected input shape {batch_size, seq_len}, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_context_length!</span><span class="p">(</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">context_length</span><span class="p">)</span> <span class="ow">when</span> <span class="n">seq_len</span> <span class="o">&lt;=</span> <span class="n">context_length</span><span class="p">,</span>
    <span class="k">do</span><span class="p">:</span> <span class="ss">:ok</span>

  <span class="k">defp</span> <span class="n">validate_context_length!</span><span class="p">(</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">context_length</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"seq_len (</span><span class="si">#{</span><span class="n">seq_len</span><span class="si">}</span><span class="s2">) exceeds context_length (</span><span class="si">#{</span><span class="n">context_length</span><span class="si">}</span><span class="s2">)"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_number</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&lt;</span> <span class="mi">1</span><span class="p">,</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">value</span> <span class="o">*</span> <span class="mf">1.0</span>

  <span class="k">defp</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected drop_rate to be a number in [0, 1), got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">System</span><span class="o">.</span><span class="n">unique_integer</span><span class="p">([</span><span class="ss">:positive</span><span class="p">])</span>
  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">seed</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">seed</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"seed must be an integer or nil, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And here is the test that uses <code class="language-plaintext highlighter-rouge">GPTModel</code>. Note that we are counting the number of parameters and weights. <code class="language-plaintext highlighter-rouge">weight tying</code> is a concept where we reuse token embedding weights in the output layer. Sebastian does not do that in the book, as it produces worse model performance.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"GPT-124M model returns logits for input batch"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">batch</span> <span class="o">=</span>
      <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span>
        <span class="p">[</span>
          <span class="p">[</span><span class="mi">6109</span><span class="p">,</span> <span class="mi">3626</span><span class="p">,</span> <span class="mi">6100</span><span class="p">,</span> <span class="mi">345</span><span class="p">],</span>
          <span class="p">[</span><span class="mi">6109</span><span class="p">,</span> <span class="mi">1110</span><span class="p">,</span> <span class="mi">6622</span><span class="p">,</span> <span class="mi">257</span><span class="p">]</span>
        <span class="p">],</span>
        <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
      <span class="p">)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>
    <span class="n">total_params</span> <span class="o">=</span> <span class="no">GPTModeltotal_parameters</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
    <span class="n">total_size_mb</span> <span class="o">=</span> <span class="n">total_params</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">/</span> <span class="mi">1024</span> <span class="o">/</span> <span class="mi">1024</span>
    <span class="n">total_params_with_weight_tying</span> <span class="o">=</span> <span class="n">total_params</span> <span class="o">-</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">tensor_parameters</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="o">.</span><span class="n">kernel</span><span class="p">)</span>
    <span class="n">out</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">batch</span><span class="p">)</span>

    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span>
             <span class="n">batch</span><span class="p">,</span>
             <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span>
               <span class="p">[</span>
                 <span class="p">[</span><span class="mi">6109</span><span class="p">,</span> <span class="mi">3626</span><span class="p">,</span> <span class="mi">6100</span><span class="p">,</span> <span class="mi">345</span><span class="p">],</span>
                 <span class="p">[</span><span class="mi">6109</span><span class="p">,</span> <span class="mi">1110</span><span class="p">,</span> <span class="mi">6622</span><span class="p">,</span> <span class="mi">257</span><span class="p">]</span>
               <span class="p">],</span>
               <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">}</span>
             <span class="p">)</span>
           <span class="p">)</span>
           <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">all</span><span class="p">()</span>
           <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_number</span><span class="p">()</span> <span class="o">==</span> <span class="mi">1</span>

    <span class="n">assert</span> <span class="n">total_params</span> <span class="o">==</span> <span class="mi">163_009_536</span>
    <span class="n">assert</span> <span class="n">total_size_mb</span> <span class="o">==</span> <span class="mf">621.83203125</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">out</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">50_257</span><span class="p">}</span>
    <span class="n">assert</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">out</span>
    <span class="n">assert</span> <span class="n">total_params_with_weight_tying</span> <span class="o">==</span> <span class="mi">124_412_160</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">total_parameters</span><span class="p">(%</span><span class="no">GPTModel</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">model</span><span class="o">.</span><span class="n">tok_emb</span><span class="o">.</span><span class="n">weight</span>
    <span class="o">|&gt;</span> <span class="n">count_tensor</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">count_tensor</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">pos_emb</span><span class="o">.</span><span class="n">weight</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="o">&amp;</span><span class="p">(</span><span class="nv">&amp;2</span> <span class="o">+</span> <span class="n">count_transformer_block</span><span class="p">(</span><span class="nv">&amp;1</span><span class="p">))))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">count_layer_norm</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">final_norm</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">count_tensor</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">out_head</span><span class="o">.</span><span class="n">kernel</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">count_transformer_block</span><span class="p">(</span><span class="n">block</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">count_attention</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">count_feed_forward</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">count_layer_norm</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">norm1</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">count_layer_norm</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">norm2</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">count_attention</span><span class="p">(</span><span class="n">attention</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">attention</span><span class="o">.</span><span class="n">w_q</span>
    <span class="o">|&gt;</span> <span class="n">count_dense</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">count_dense</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_k</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">count_dense</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">w_v</span><span class="p">,</span> <span class="n">attention</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">Kernel</span><span class="o">.+</span><span class="p">(</span><span class="n">count_dense</span><span class="p">(</span><span class="n">attention</span><span class="o">.</span><span class="n">out_proj</span><span class="p">,</span> <span class="no">true</span><span class="p">))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">count_feed_forward</span><span class="p">(</span><span class="n">feed_forward</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">count_dense</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">first</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">count_dense</span><span class="p">(</span><span class="n">feed_forward</span><span class="o">.</span><span class="n">layers</span><span class="o">.</span><span class="n">second</span><span class="p">,</span> <span class="no">true</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">count_layer_norm</span><span class="p">(</span><span class="n">layer_norm</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">count_tensor</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">scale</span><span class="p">)</span> <span class="o">+</span> <span class="n">count_tensor</span><span class="p">(</span><span class="n">layer_norm</span><span class="o">.</span><span class="n">shift</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">count_dense</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">,</span> <span class="ss">bias:</span> <span class="n">bias</span><span class="p">},</span> <span class="no">true</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">count_tensor</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span> <span class="o">+</span> <span class="n">count_tensor</span><span class="p">(</span><span class="n">bias</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">count_dense</span><span class="p">(%{</span><span class="ss">kernel:</span> <span class="n">kernel</span><span class="p">},</span> <span class="no">false</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">count_tensor</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">count_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">Nx</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>
</code></pre></div></div>

<p>Sebastian provides two exercises.</p>

<p>Exercise 4.1 Number of parameters in feed-forward and attention modules<br />
Calculate and compare the number of parameters contained in the feed-forward module and the multi-head attention module.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"exercise 4.1 compares feed forward and multi-head attention parameter counts"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">block</span> <span class="o">=</span> <span class="no">TransformerBlock</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="n">feed_forward_params</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">feed_forward_parameters</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="p">)</span>
    <span class="n">attention_params</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">attention_parameters</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">feed_forward_params</span> <span class="o">==</span> <span class="mi">4_722_432</span>
    <span class="n">assert</span> <span class="n">attention_params</span> <span class="o">==</span> <span class="mi">2_360_064</span>
    <span class="n">assert</span> <span class="n">feed_forward_params</span> <span class="o">==</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">attention_params</span> <span class="o">+</span> <span class="mi">2_304</span>
    <span class="n">assert</span> <span class="n">feed_forward_params</span> <span class="o">&gt;</span> <span class="n">attention_params</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Exercise 4.2 Initializing larger GPT models<br />
We initialized a 124-million-parameter GPT model, which is known as “GPT-2 small.” Without making any code modifications besides updating the configuration file, use the <code class="language-plaintext highlighter-rouge">GPTModel</code> class to implement GPT-2 medium (using 1,024-dimensional embeddings, 24 transformer blocks, and 16 multi-head attention heads), GPT-2 large (1,280-dimensional embeddings, 36 transformer blocks, and 20 multi-head attention heads), and GPT-2 XL (1,600-dimensional embeddings, 48 transformer blocks, and 25 multi-head attention heads). As a bonus, calculate the total number of parameters in each GPT model.</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"exercise 4.2 initializes GPT-2 medium config and counts parameters"</span> <span class="k">do</span>
    <span class="n">gpt2_medium_config</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">16</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">24</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">assert</span> <span class="n">gpt2_medium_config</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">==</span> <span class="mi">1024</span>
    <span class="n">assert</span> <span class="n">gpt2_medium_config</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">==</span> <span class="mi">24</span>
    <span class="n">assert</span> <span class="n">gpt2_medium_config</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">==</span> <span class="mi">16</span>
    <span class="n">assert</span> <span class="n">rem</span><span class="p">(</span><span class="n">gpt2_medium_config</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">gpt2_medium_config</span><span class="o">.</span><span class="n">n_heads</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">estimated_parameter_count</span><span class="p">(</span><span class="n">gpt2_medium_config</span><span class="p">)</span> <span class="o">==</span> <span class="mi">406_212_608</span>
  <span class="k">end</span>

  <span class="n">test</span> <span class="s2">"exercise 4.2 initializes GPT-2 large config and counts parameters"</span> <span class="k">do</span>
    <span class="n">gpt2_large_config</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">1280</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">20</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">36</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">assert</span> <span class="n">gpt2_large_config</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">==</span> <span class="mi">1280</span>
    <span class="n">assert</span> <span class="n">gpt2_large_config</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">==</span> <span class="mi">36</span>
    <span class="n">assert</span> <span class="n">gpt2_large_config</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">==</span> <span class="mi">20</span>
    <span class="n">assert</span> <span class="n">rem</span><span class="p">(</span><span class="n">gpt2_large_config</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">gpt2_large_config</span><span class="o">.</span><span class="n">n_heads</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">estimated_parameter_count</span><span class="p">(</span><span class="n">gpt2_large_config</span><span class="p">)</span> <span class="o">==</span> <span class="mi">838_220_800</span>
  <span class="k">end</span>

  <span class="n">test</span> <span class="s2">"exercise 4.2 initializes GPT-2 XL config and counts parameters"</span> <span class="k">do</span>
    <span class="n">gpt2_xl_config</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">1600</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">25</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">48</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">assert</span> <span class="n">gpt2_xl_config</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">==</span> <span class="mi">1600</span>
    <span class="n">assert</span> <span class="n">gpt2_xl_config</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">==</span> <span class="mi">48</span>
    <span class="n">assert</span> <span class="n">gpt2_xl_config</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">==</span> <span class="mi">25</span>
    <span class="n">assert</span> <span class="n">rem</span><span class="p">(</span><span class="n">gpt2_xl_config</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="n">gpt2_xl_config</span><span class="o">.</span><span class="n">n_heads</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span>
    <span class="n">assert</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">estimated_parameter_count</span><span class="p">(</span><span class="n">gpt2_xl_config</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1_637_792_000</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="generating-text">Generating text</h2>

<p>Finally, it is time to bear the fruits of our hard labor and generate text from input tokens. The question is: how do we get from output tensors to generated text?<br />
A GPT model predicts the next token from input embeddings in iterations. Let’s look at an example that starts with this text:<br />
Hello I am …</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>
1st iteration: Hello I am =&gt; Hello I am a  
2nd iteration: Hello I am a =&gt; Hello I am a model  
3rd iteration: Hello I am a model ready  
4th iteration: Hello I am a model ready to  
5th iteration: Hello I am a model ready to help  
6th iteration: Hello I am a model ready to help.  
</code></pre></div></div>

<p>So how do we do this in GPTModel?</p>

<div class="language-text highlighter-rouge"><div class="highlight"><pre class="highlight"><code>

                         GPT TEXT GENERATION LOOP

┌─────────────────────────────────────────────────────────────────────┐
│ 1. Encode input text into token IDs                                  │
│                                                                     │
│   "Hello"   ","   "I"   "am"                                        │
│      │       │     │     │                                          │
│      ▼       ▼     ▼     ▼                                          │
│    15496    11    314   716                                         │
│                                                                     │
│   Input IDs: [15496, 11, 314, 716]                                  │
└───────────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
                         ┌─────────────┐
                         │     GPT     │
                         └──────┬──────┘
                                │
                                ▼
┌─────────────────────────────────────────────────────────────────────┐
│ 2. Model output matrix                                               │
│                                                                     │
│   One row per input token.                                           │
│   One column per vocabulary token.                                   │
│                                                                     │
│   [ [-0.2949, ..., -0.8141],                                         │
│     [ 1.2199, ..., -0.3599],                                         │
│     [ 1.0446, ...,  0.0020],                                         │
│     [-0.4929, ..., -0.6093] ]  ◄── last row predicts next token       │
└───────────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
┌─────────────────────────────────────────────────────────────────────┐
│ 3. Extract last vector: logits                                       │
│                                                                     │
│   Logits: [-0.4929, ..., 2.4812, ..., -0.6093]                      │
└───────────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
┌─────────────────────────────────────────────────────────────────────┐
│ 4. Apply softmax                                                     │
│                                                                     │
│   Converts logits into probabilities.                                │
│                                                                     │
│   Probabilities: [0.0001, ..., 0.0200, ..., 0.0001]                 │
└───────────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
┌─────────────────────────────────────────────────────────────────────┐
│ 5. Choose next token                                                 │
│                                                                     │
│   Largest probability is at index 257.                               │
│                                                                     │
│   index 257  ──►  token ID 257  ──►  "a"                            │
└───────────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
┌─────────────────────────────────────────────────────────────────────┐
│ 6. Append generated token to input                                   │
│                                                                     │
│   Before: [15496, 11, 314, 716]                                      │
│   After:  [15496, 11, 314, 716, 257]                                 │
│                                                                     │
│   Text: "Hello, I am a"                                              │
└───────────────────────────────┬─────────────────────────────────────┘
                                │
                                │ repeat
                                ▼
                         ┌─────────────┐
                         │     GPT     │
                         └─────────────┘
</code></pre></div></div>

<p>Here is the module that implements this algorithm:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">TextGeneration</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Greedy token generation helpers for GPT-style models.

  `generate_text_simple/4` mirrors the Python loop from the book:

      idx_cond = idx[:, -context_size:]
      logits = model(idx_cond)
      logits = logits[:, -1, :]
      probas = torch.softmax(logits, dim=-1)
      idx_next = torch.argmax(probas, dim=-1, keepdim=True)
      idx = torch.cat((idx, idx_next), dim=1)

  The model is expected to be a struct whose module exports `forward/2`.
  Existing modules such as `LlmScratch.GPTModel` and
  `LlmScratch.DummyGPTModel` satisfy that contract.
  """</span>

  <span class="nv">@spec</span> <span class="n">generate_text_simple</span><span class="p">(</span><span class="n">struct</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">non_neg_integer</span><span class="p">(),</span> <span class="n">pos_integer</span><span class="p">())</span> <span class="p">::</span>
          <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Generates new token ids with greedy decoding.

  ## Arguments

    * `model` - a GPT-style model struct whose module exports `forward/2`.
      `forward/2` receives token ids shaped `{batch_size, seq_len}` and must
      return logits shaped `{batch_size, seq_len, vocab_size}`.

    * `idx` - token ids shaped `{batch_size, seq_len}`. The returned tensor
      keeps the same batch size and token id type.

    * `max_new_tokens` - number of new tokens to generate and append. Use `0`
      to return `idx` unchanged.

    * `context_size` - maximum number of latest tokens to pass to the model on
      each generation step.

  ## Returns

  A token id tensor shaped `{batch_size, seq_len + max_new_tokens}`.

  The next token is selected with `argmax` over the final-position softmax
  probabilities, then appended to `idx` along the sequence axis.

  ## Examples

      generated =
        LlmScratch.TextGeneration.generate_text_simple(
          model,
          Nx.tensor([[6109, 3626]], type: {:s, 64}),
          5,
          1024
        )

      Nx.shape(generated)
      #=&gt; {1, 7}
  """</span>
  <span class="k">def</span> <span class="n">generate_text_simple</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">idx</span><span class="p">,</span> <span class="n">max_new_tokens</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
      <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">max_new_tokens</span><span class="p">)</span> <span class="ow">and</span> <span class="n">max_new_tokens</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">context_size</span><span class="p">)</span> <span class="ow">and</span>
             <span class="n">context_size</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">do</span>
    <span class="n">validate_idx_shape!</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>

    <span class="no">Enum</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="mi">1</span><span class="o">..</span><span class="n">max_new_tokens</span><span class="o">//</span><span class="mi">1</span><span class="p">,</span> <span class="n">idx</span><span class="p">,</span> <span class="k">fn</span> <span class="n">_step</span><span class="p">,</span> <span class="n">acc</span> <span class="o">-&gt;</span>
      <span class="n">idx_cond</span> <span class="o">=</span> <span class="n">last_tokens</span><span class="p">(</span><span class="n">acc</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
      <span class="n">logits</span> <span class="o">=</span> <span class="n">forward!</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">idx_cond</span><span class="p">)</span>
      <span class="n">logits</span> <span class="o">=</span> <span class="n">last_position_logits</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span>
      <span class="n">probas</span> <span class="o">=</span> <span class="no">Axon</span><span class="o">.</span><span class="no">Activations</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>

      <span class="n">idx_next</span> <span class="o">=</span>
        <span class="n">probas</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="ss">axis:</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="ss">keep_axis:</span> <span class="no">true</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">as_type</span><span class="p">(</span><span class="no">Nx</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">acc</span><span class="p">))</span>

      <span class="no">Nx</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">acc</span><span class="p">,</span> <span class="n">idx_next</span><span class="p">],</span> <span class="ss">axis:</span> <span class="mi">1</span><span class="p">)</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">last_tokens</span><span class="p">(</span><span class="n">idx</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">}</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>
    <span class="n">length</span> <span class="o">=</span> <span class="n">min</span><span class="p">(</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">context_size</span><span class="p">)</span>
    <span class="n">start</span> <span class="o">=</span> <span class="n">seq_len</span> <span class="o">-</span> <span class="n">length</span>

    <span class="no">Nx</span><span class="o">.</span><span class="n">slice_along_axis</span><span class="p">(</span><span class="n">idx</span><span class="p">,</span> <span class="n">start</span><span class="p">,</span> <span class="n">length</span><span class="p">,</span> <span class="ss">axis:</span> <span class="mi">1</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">last_position_logits</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">_vocab_size</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="n">logits</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">slice_along_axis</span><span class="p">(</span><span class="n">seq_len</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="ss">axis:</span> <span class="mi">1</span><span class="p">)</span>
        <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="ss">axes:</span> <span class="p">[</span><span class="mi">1</span><span class="p">])</span>

      <span class="n">shape</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span>
              <span class="s2">"expected logits shape {batch_size, seq_len, vocab_size}, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">forward!</span><span class="p">(%</span><span class="n">module</span><span class="p">{}</span> <span class="o">=</span> <span class="n">model</span><span class="p">,</span> <span class="n">idx_cond</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">if</span> <span class="n">function_exported?</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="ss">:forward</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span> <span class="k">do</span>
      <span class="n">module</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">idx_cond</span><span class="p">)</span>
    <span class="k">else</span>
      <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">module</span><span class="p">)</span><span class="si">}</span><span class="s2"> to export forward/2"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_idx_shape!</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span> <span class="k">do</span>
      <span class="p">{</span><span class="n">_batch_size</span><span class="p">,</span> <span class="n">_seq_len</span><span class="p">}</span> <span class="o">-&gt;</span>
        <span class="ss">:ok</span>

      <span class="n">shape</span> <span class="o">-&gt;</span>
        <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected idx shape {batch_size, seq_len}, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And here is the corresponding test that produces gibberish text, since we still have not trained our model:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"generate_text_simple generates text from a GPT-124M start context"</span> <span class="k">do</span>
    <span class="n">previous_backend</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">()</span>
    <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="no">EXLA</span><span class="o">.</span><span class="no">Backend</span><span class="p">)</span>
    <span class="n">on_exit</span><span class="p">(</span><span class="k">fn</span> <span class="o">-&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">default_backend</span><span class="p">(</span><span class="n">previous_backend</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>

    <span class="n">gpt_config_124m</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">50_257</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">1024</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">768</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.0</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">start_context</span> <span class="o">=</span> <span class="s2">"Hello, I am"</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">encoded</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">,</span> <span class="n">start_context</span><span class="p">)</span>
    <span class="n">encoded_tensor</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">encoded</span><span class="p">,</span> <span class="ss">type:</span> <span class="p">{</span><span class="ss">:s</span><span class="p">,</span> <span class="mi">64</span><span class="p">})</span> <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">new_axis</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">gpt_config_124m</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>

    <span class="n">out</span> <span class="o">=</span>
      <span class="no">TextGeneration</span><span class="o">.</span><span class="n">generate_text_simple</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">encoded_tensor</span><span class="p">,</span>
        <span class="mi">6</span><span class="p">,</span>
        <span class="n">gpt_config_124m</span><span class="o">.</span><span class="n">context_length</span>
      <span class="p">)</span>

    <span class="n">output_tokens</span> <span class="o">=</span> <span class="no">Nx</span><span class="o">.</span><span class="n">to_flat_list</span><span class="p">(</span><span class="n">out</span><span class="p">)</span>
    <span class="p">{</span><span class="ss">:ok</span><span class="p">,</span> <span class="n">decoded_text</span><span class="p">}</span> <span class="o">=</span> <span class="no">Tiktoken</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="s2">"code-davinci-002"</span><span class="p">,</span> <span class="n">output_tokens</span><span class="p">)</span>

    <span class="n">assert</span> <span class="n">encoded</span> <span class="o">==</span> <span class="p">[</span><span class="mi">15496</span><span class="p">,</span> <span class="mi">11</span><span class="p">,</span> <span class="mi">314</span><span class="p">,</span> <span class="mi">716</span><span class="p">]</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">encoded_tensor</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">4</span><span class="p">}</span>
    <span class="n">assert</span> <span class="n">output_tokens</span> <span class="o">==</span> <span class="p">[</span><span class="mi">15496</span><span class="p">,</span> <span class="mi">11</span><span class="p">,</span> <span class="mi">314</span><span class="p">,</span> <span class="mi">716</span><span class="p">,</span> <span class="mi">8026</span><span class="p">,</span> <span class="mi">707</span><span class="p">,</span> <span class="mi">38647</span><span class="p">,</span> <span class="mi">1336</span><span class="p">,</span> <span class="mi">10296</span><span class="p">,</span> <span class="mi">24545</span><span class="p">]</span>
    <span class="n">assert</span> <span class="no">Nx</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">out</span><span class="p">)</span> <span class="o">==</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">10</span><span class="p">}</span>
    <span class="n">assert</span> <span class="n">length</span><span class="p">(</span><span class="n">output_tokens</span><span class="p">)</span> <span class="o">==</span> <span class="mi">10</span>
    <span class="n">assert</span> <span class="n">decoded_text</span> <span class="o">==</span> <span class="s2">"Hello, I am Stoneaw Eas full gentle Rhode"</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>Sepbastina prepared for us another exercise:</p>

<p>Exercise 4.3 Using separate dropout parameters<br />
At the beginning of this chapter, we defined a global drop_rate setting in the GPT_ CONFIG_124M dictionary to set the dropout rate in various places throughout the GPTModel architecture. Change the code to specify a separate dropout value for the various dropout layers throughout the model architecture. (Hint: there are three dis- tinct places where we used dropout layers: the embedding layer, shortcut layer, and multi-head attention module.)</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">GPTConfig</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  GPT model configuration.

  This struct mirrors the config keys used by the Python examples.

  ## Fields

    * `:vocab_size` - number of tokens in the tokenizer vocabulary. This
      determines the number of rows in the token embedding table and the final
      logits dimension produced by the output head. For GPT-2 this is commonly
      `50_257`.

    * `:context_length` - maximum number of token positions the model can
      process in one forward pass. This determines the size of the positional
      embedding table. Inputs longer than this cannot be represented by the
      configured positional embeddings.

    * `:emb_dim` - width of each token and positional embedding vector. This is
      also the model's hidden size: tensors flowing through transformer blocks
      have shape `{batch_size, seq_len, emb_dim}`.

    * `:n_heads` - number of attention heads used by real transformer blocks.
      The dummy GPT model stores this value for config compatibility, but its
      placeholder transformer blocks do not use it yet.

    * `:n_layers` - number of transformer blocks in the model. In
      `LlmScratch.DummyGPTModel`, this controls how many dummy identity blocks
      are created.

    * `:drop_rate` - backward-compatible default dropout probability. It is
      used when a more specific dropout field is not set.

    * `:emb_drop_rate` - dropout probability used after adding token and
      positional embeddings.

    * `:shortcut_drop_rate` - dropout probability used on transformer-block
      shortcut branches after attention and feed-forward outputs.

    * `:attn_drop_rate` - dropout probability used on attention weights inside
      multi-head attention.

    * `:qkv_bias` - whether query, key, and value projections include bias
      terms in real attention layers. The dummy GPT model keeps this field for
      compatibility with the full GPT config shape, but does not use it because
      its transformer blocks are placeholders.
  """</span>

  <span class="nv">@enforce_keys</span> <span class="p">[</span>
    <span class="ss">:vocab_size</span><span class="p">,</span>
    <span class="ss">:context_length</span><span class="p">,</span>
    <span class="ss">:emb_dim</span><span class="p">,</span>
    <span class="ss">:n_heads</span><span class="p">,</span>
    <span class="ss">:n_layers</span><span class="p">,</span>
    <span class="ss">:drop_rate</span><span class="p">,</span>
    <span class="ss">:qkv_bias</span>
  <span class="p">]</span>

  <span class="k">defstruct</span> <span class="p">[</span>
    <span class="ss">:vocab_size</span><span class="p">,</span>
    <span class="ss">:context_length</span><span class="p">,</span>
    <span class="ss">:emb_dim</span><span class="p">,</span>
    <span class="ss">:n_heads</span><span class="p">,</span>
    <span class="ss">:n_layers</span><span class="p">,</span>
    <span class="ss">:drop_rate</span><span class="p">,</span>
    <span class="ss">:emb_drop_rate</span><span class="p">,</span>
    <span class="ss">:shortcut_drop_rate</span><span class="p">,</span>
    <span class="ss">:attn_drop_rate</span><span class="p">,</span>
    <span class="ss">:qkv_bias</span>
  <span class="p">]</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">vocab_size:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">context_length:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">emb_dim:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">n_heads:</span> <span class="n">pos_integer</span><span class="p">(),</span>
          <span class="ss">n_layers:</span> <span class="n">non_neg_integer</span><span class="p">(),</span>
          <span class="ss">drop_rate:</span> <span class="n">float</span><span class="p">(),</span>
          <span class="ss">emb_drop_rate:</span> <span class="n">float</span><span class="p">()</span> <span class="o">|</span> <span class="no">nil</span><span class="p">,</span>
          <span class="ss">shortcut_drop_rate:</span> <span class="n">float</span><span class="p">()</span> <span class="o">|</span> <span class="no">nil</span><span class="p">,</span>
          <span class="ss">attn_drop_rate:</span> <span class="n">float</span><span class="p">()</span> <span class="o">|</span> <span class="no">nil</span><span class="p">,</span>
          <span class="ss">qkv_bias:</span> <span class="n">boolean</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">embedding_dropout</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">float</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Returns the embedding dropout rate, falling back to `drop_rate` when unset.
  """</span>
  <span class="k">def</span> <span class="n">embedding_dropout</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">dropout_value!</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_drop_rate</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">)</span>

  <span class="nv">@spec</span> <span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">float</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Returns the shortcut dropout rate, falling back to `drop_rate` when unset.
  """</span>
  <span class="k">def</span> <span class="n">shortcut_dropout</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">dropout_value!</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">shortcut_drop_rate</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">)</span>

  <span class="nv">@spec</span> <span class="n">attention_dropout</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">float</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Returns the attention dropout rate, falling back to `drop_rate` when unset.
  """</span>
  <span class="k">def</span> <span class="n">attention_dropout</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">dropout_value!</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">attn_drop_rate</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">)</span>

  <span class="nv">@spec</span> <span class="n">estimated_parameter_count</span><span class="p">(</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="n">non_neg_integer</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Estimates GPT model parameters directly from the configuration.

  This mirrors the architecture used by `LlmScratch.GPTModel` without
  allocating tensors:

    * token embeddings: `vocab_size * emb_dim`
    * positional embeddings: `context_length * emb_dim`
    * untied output projection: `emb_dim * vocab_size`
    * each transformer block: attention projections, feed-forward projections,
      and two layer norms
    * final layer norm: scale and shift, each shaped `{emb_dim}`

  Use `LlmScratch.GPTModel.total_parameters/1` when you already have a
  concrete model and want to count its actual tensors.
  """</span>
  <span class="k">def</span> <span class="n">estimated_parameter_count</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">token_embedding_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span>
    <span class="n">positional_embedding_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span>
    <span class="n">output_projection_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span>
    <span class="n">block_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span> <span class="o">*</span> <span class="n">estimated_transformer_block_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
    <span class="n">final_norm_params</span> <span class="o">=</span> <span class="n">estimated_layer_norm_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span>

    <span class="n">token_embedding_params</span> <span class="o">+</span> <span class="n">positional_embedding_params</span> <span class="o">+</span> <span class="n">output_projection_params</span> <span class="o">+</span> <span class="n">block_params</span> <span class="o">+</span>
      <span class="n">final_norm_params</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">estimated_transformer_block_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">estimated_attention_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="o">+</span>
      <span class="n">estimated_feed_forward_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span> <span class="o">+</span>
      <span class="mi">2</span> <span class="o">*</span> <span class="n">estimated_layer_norm_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">estimated_attention_parameter_count</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">qkv_kernel_params</span> <span class="o">=</span> <span class="mi">3</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span>
    <span class="n">qkv_bias_params</span> <span class="o">=</span> <span class="k">if</span> <span class="n">cfg</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="mi">3</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="k">else</span><span class="p">:</span> <span class="mi">0</span>
    <span class="n">output_projection_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span> <span class="o">+</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span>

    <span class="n">qkv_kernel_params</span> <span class="o">+</span> <span class="n">qkv_bias_params</span> <span class="o">+</span> <span class="n">output_projection_params</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">estimated_feed_forward_parameter_count</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">expanded_dim</span> <span class="o">=</span> <span class="mi">4</span> <span class="o">*</span> <span class="n">emb_dim</span>
    <span class="n">first_projection_params</span> <span class="o">=</span> <span class="n">emb_dim</span> <span class="o">*</span> <span class="n">expanded_dim</span> <span class="o">+</span> <span class="n">expanded_dim</span>
    <span class="n">second_projection_params</span> <span class="o">=</span> <span class="n">expanded_dim</span> <span class="o">*</span> <span class="n">emb_dim</span> <span class="o">+</span> <span class="n">emb_dim</span>

    <span class="n">first_projection_params</span> <span class="o">+</span> <span class="n">second_projection_params</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">estimated_layer_norm_parameter_count</span><span class="p">(</span><span class="n">emb_dim</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">emb_dim</span>

  <span class="k">defp</span> <span class="n">dropout_value!</span><span class="p">(</span><span class="no">nil</span><span class="p">,</span> <span class="n">fallback</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">fallback</span><span class="p">)</span>
  <span class="k">defp</span> <span class="n">dropout_value!</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="n">_fallback</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">value</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_number</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&lt;</span> <span class="mi">1</span><span class="p">,</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">value</span> <span class="o">*</span> <span class="mf">1.0</span>

  <span class="k">defp</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected dropout value to be a number in [0, 1), got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>GPTModel has changed its <code class="language-plaintext highlighter-rouge">new</code> function:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="k">def</span> <span class="n">new</span><span class="p">(%</span><span class="no">GPTConfig</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">))</span>
    <span class="n">norm_eps</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:norm_eps</span><span class="p">,</span> <span class="mf">1.0e-5</span><span class="p">)</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">cfg:</span> <span class="n">cfg</span><span class="p">,</span>
      <span class="ss">tok_emb:</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">vocab_size</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span><span class="p">),</span>
      <span class="ss">pos_emb:</span> <span class="no">EmbeddingNative</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
      <span class="ss">drop_emb:</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">embedding_dropout</span><span class="p">(</span><span class="n">cfg</span><span class="p">),</span>
      <span class="ss">trf_blocks:</span> <span class="n">transformer_blocks</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span><span class="p">,</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">final_norm:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">eps:</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">out_head:</span> <span class="n">init_out_head</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span> <span class="o">+</span> <span class="n">cfg</span><span class="o">.</span><span class="n">n_layers</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>
</code></pre></div></div>

<p>And TransformerBlock has two new attributes, att_drop_reate and shortcut_drop_rate:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code>  <span class="k">defmodule</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="no">TransformerBlock</span> <span class="k">do</span>
  <span class="nv">@moduledoc</span> <span class="sd">"""
  Real GPT-style transformer block implemented with Nx tensors.

  This mirrors the book's PyTorch `TransformerBlock`:

      shortcut = x
      x = norm1(x)
      x = attention(x)
      x = dropout(x)
      x = x + shortcut

      shortcut = x
      x = norm2(x)
      x = feed_forward(x)
      x = dropout(x)
      x = x + shortcut

  The block uses pre-normalization, causal multi-head self-attention, a
  position-wise feed-forward network, and residual shortcut connections.
  """</span>

  <span class="n">alias</span> <span class="no">LlmScratch</span><span class="o">.</span><span class="p">{</span>
    <span class="no">DummyLayerNorm</span><span class="p">,</span>
    <span class="no">FeedForward</span><span class="p">,</span>
    <span class="no">GPTConfig</span><span class="p">,</span>
    <span class="no">MultiheadAttention</span>
  <span class="p">}</span>

  <span class="k">defstruct</span> <span class="p">[</span><span class="ss">:att</span><span class="p">,</span> <span class="ss">:ff</span><span class="p">,</span> <span class="ss">:norm1</span><span class="p">,</span> <span class="ss">:norm2</span><span class="p">,</span> <span class="ss">:drop_shortcut</span><span class="p">,</span> <span class="ss">:cfg</span><span class="p">]</span>

  <span class="nv">@type</span> <span class="n">t</span> <span class="p">::</span> <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
          <span class="ss">att:</span> <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">ff:</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">norm1:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">norm2:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span>
          <span class="ss">drop_shortcut:</span> <span class="n">float</span><span class="p">(),</span>
          <span class="ss">cfg:</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">()</span> <span class="o">|</span> <span class="n">map</span><span class="p">()</span>
        <span class="p">}</span>

  <span class="nv">@spec</span> <span class="n">new</span><span class="p">(</span><span class="no">GPTConfig</span><span class="o">.</span><span class="n">t</span><span class="p">()</span> <span class="o">|</span> <span class="n">map</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Creates a GPT transformer block from a config.

  `cfg` can be a `%LlmScratch.GPTConfig{}` or a map with atom/string keys:
  `:emb_dim`, `:context_length`, `:n_heads`, `:drop_rate`, and `:qkv_bias`.
  `:attn_drop_rate` and `:shortcut_drop_rate` may be supplied to control those
  dropout sites independently; otherwise `:drop_rate` is used for compatibility.

  ## Options

    * `:seed` - deterministic seed used for attention and feed-forward
      initialization.
    * `:norm_eps` - epsilon for both layer-normalization modules.
  """</span>
  <span class="k">def</span> <span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="ow">when</span> <span class="n">is_map</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_list</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">params</span> <span class="o">=</span> <span class="n">params_from_cfg!</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:seed</span><span class="p">))</span>
    <span class="n">norm_eps</span> <span class="o">=</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:norm_eps</span><span class="p">,</span> <span class="mf">1.0e-5</span><span class="p">)</span>

    <span class="p">%</span><span class="bp">__MODULE__</span><span class="p">{</span>
      <span class="ss">cfg:</span> <span class="n">cfg</span><span class="p">,</span>
      <span class="ss">att:</span>
        <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">new</span><span class="p">(</span>
          <span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">context_length</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">attn_drop_rate</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">n_heads</span><span class="p">,</span>
          <span class="n">params</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">,</span>
          <span class="ss">seed:</span> <span class="n">seed</span>
        <span class="p">),</span>
      <span class="ss">ff:</span> <span class="no">FeedForward</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
      <span class="ss">norm1:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">eps:</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">norm2:</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">,</span> <span class="ss">eps:</span> <span class="n">norm_eps</span><span class="p">),</span>
      <span class="ss">drop_shortcut:</span> <span class="n">params</span><span class="o">.</span><span class="n">shortcut_drop_rate</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">forward</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">keyword</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Runs a forward pass over hidden states shaped `{batch_size, seq_len, emb_dim}`.

  Forward options:

    * `:mode` - `:train` or `:inference`. Dropout is skipped in inference mode.
    * `:key` - optional `Nx.Random` key used by dropout in train mode.
    * `:attention_key`, `:shortcut_key1`, `:shortcut_key2` - optional keys for
      controlling individual dropout calls.
  """</span>
  <span class="k">def</span> <span class="n">forward</span><span class="p">(%</span><span class="bp">__MODULE__</span><span class="p">{}</span> <span class="o">=</span> <span class="n">block</span><span class="p">,</span> <span class="p">%</span><span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="p">{}</span> <span class="o">=</span> <span class="n">x</span><span class="p">,</span> <span class="n">opts</span> <span class="p">\\</span> <span class="p">[])</span> <span class="k">do</span>
    <span class="n">shortcut</span> <span class="o">=</span> <span class="n">x</span>

    <span class="n">x</span> <span class="o">=</span>
      <span class="n">block</span><span class="o">.</span><span class="n">norm1</span>
      <span class="o">|&gt;</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
      <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="p">,</span> <span class="nv">&amp;1</span><span class="p">,</span> <span class="n">attention_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">)))</span>
      <span class="o">|&gt;</span> <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">maybe_dropout</span><span class="p">(</span>
        <span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">seed</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
        <span class="n">shortcut_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:shortcut_key1</span><span class="p">)</span>
      <span class="p">)</span>
      <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">shortcut</span><span class="p">)</span>

    <span class="n">shortcut</span> <span class="o">=</span> <span class="n">x</span>

    <span class="n">block</span><span class="o">.</span><span class="n">norm2</span>
    <span class="o">|&gt;</span> <span class="no">DummyLayerNorm</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
    <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="o">&amp;</span><span class="no">FeedForward</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">block</span><span class="o">.</span><span class="n">ff</span><span class="p">,</span> <span class="nv">&amp;1</span><span class="p">))</span>
    <span class="o">|&gt;</span> <span class="no">MultiheadAttention</span><span class="o">.</span><span class="n">maybe_dropout</span><span class="p">(</span>
      <span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">seed</span> <span class="o">+</span> <span class="mi">2</span><span class="p">),</span>
      <span class="n">shortcut_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:shortcut_key2</span><span class="p">)</span>
    <span class="p">)</span>
    <span class="o">|&gt;</span> <span class="no">Nx</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">shortcut</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="nv">@spec</span> <span class="n">call</span><span class="p">(</span><span class="n">t</span><span class="p">(),</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">())</span> <span class="p">::</span> <span class="no">Nx</span><span class="o">.</span><span class="no">Tensor</span><span class="o">.</span><span class="n">t</span><span class="p">()</span>
  <span class="nv">@doc</span> <span class="sd">"""
  Alias for `forward/3` using default forward options.

  This mirrors the callable-style API used by other modules in the project.
  Dropout is skipped because `forward/3` defaults to inference mode.
  """</span>
  <span class="k">def</span> <span class="n">call</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">x</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">forward</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">attention_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">opts</span>
    <span class="o">|&gt;</span> <span class="n">opts_with_default_mode</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="ss">:key</span><span class="p">,</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:attention_key</span><span class="p">,</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:key</span><span class="p">)))</span>
    <span class="o">|&gt;</span> <span class="no">Enum</span><span class="o">.</span><span class="n">reject</span><span class="p">(</span><span class="k">fn</span> <span class="p">{</span><span class="n">_key</span><span class="p">,</span> <span class="n">value</span><span class="p">}</span> <span class="o">-&gt;</span> <span class="n">is_nil</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">shortcut_opts</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="n">key_name</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">opts</span>
    <span class="o">|&gt;</span> <span class="n">opts_with_default_mode</span><span class="p">()</span>
    <span class="o">|&gt;</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="ss">:key</span><span class="p">,</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="n">key_name</span><span class="p">,</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:key</span><span class="p">)))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">block</span><span class="p">,</span> <span class="n">seed</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="p">%{</span><span class="ss">dropout:</span> <span class="n">block</span><span class="o">.</span><span class="n">drop_shortcut</span><span class="p">,</span> <span class="ss">seed:</span> <span class="n">seed</span><span class="p">}</span>

  <span class="k">defp</span> <span class="n">opts_with_default_mode</span><span class="p">(</span><span class="n">opts</span><span class="p">),</span>
    <span class="k">do</span><span class="p">:</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="no">Keyword</span><span class="o">.</span><span class="n">take</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="p">[</span><span class="ss">:mode</span><span class="p">]),</span> <span class="ss">:mode</span><span class="p">,</span> <span class="n">mode!</span><span class="p">(</span><span class="n">opts</span><span class="p">))</span>

  <span class="k">defp</span> <span class="n">mode!</span><span class="p">(</span><span class="n">opts</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">case</span> <span class="no">Keyword</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opts</span><span class="p">,</span> <span class="ss">:mode</span><span class="p">,</span> <span class="ss">:inference</span><span class="p">)</span> <span class="k">do</span>
      <span class="n">mode</span> <span class="ow">when</span> <span class="n">mode</span> <span class="ow">in</span> <span class="p">[</span><span class="ss">:train</span><span class="p">,</span> <span class="ss">:inference</span><span class="p">]</span> <span class="o">-&gt;</span> <span class="n">mode</span>
      <span class="n">mode</span> <span class="o">-&gt;</span> <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"mode must be :train or :inference, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">mode</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
    <span class="k">end</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">params_from_cfg!</span><span class="p">(%</span><span class="no">GPTConfig</span><span class="p">{}</span> <span class="o">=</span> <span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span>
      <span class="ss">emb_dim:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:emb_dim</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">),</span>
      <span class="ss">context_length:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:context_length</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">context_length</span><span class="p">),</span>
      <span class="ss">n_heads:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:n_heads</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">n_heads</span><span class="p">),</span>
      <span class="ss">attn_drop_rate:</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">attention_dropout</span><span class="p">(</span><span class="n">cfg</span><span class="p">),</span>
      <span class="ss">shortcut_drop_rate:</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">cfg</span><span class="p">),</span>
      <span class="ss">qkv_bias:</span> <span class="n">validate_boolean!</span><span class="p">(</span><span class="ss">:qkv_bias</span><span class="p">,</span> <span class="n">cfg</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span>
    <span class="p">}</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">params_from_cfg!</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="k">do</span>
    <span class="p">%{</span>
      <span class="ss">emb_dim:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:emb_dim</span><span class="p">),</span>
      <span class="ss">context_length:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:context_length</span><span class="p">),</span>
      <span class="ss">n_heads:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:n_heads</span><span class="p">),</span>
      <span class="ss">drop_rate:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:drop_rate</span><span class="p">),</span>
      <span class="ss">attn_drop_rate:</span> <span class="n">cfg_value</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:attn_drop_rate</span><span class="p">),</span>
      <span class="ss">shortcut_drop_rate:</span> <span class="n">cfg_value</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:shortcut_drop_rate</span><span class="p">),</span>
      <span class="ss">qkv_bias:</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">:qkv_bias</span><span class="p">)</span>
    <span class="p">}</span>
    <span class="o">|&gt;</span> <span class="n">then</span><span class="p">(</span><span class="k">fn</span> <span class="n">params</span> <span class="o">-&gt;</span>
      <span class="n">drop_rate</span> <span class="o">=</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="ss">:drop_rate</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">drop_rate</span><span class="p">)</span>

      <span class="p">%{</span>
        <span class="ss">emb_dim:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:emb_dim</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">emb_dim</span><span class="p">),</span>
        <span class="ss">context_length:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:context_length</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">context_length</span><span class="p">),</span>
        <span class="ss">n_heads:</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="ss">:n_heads</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">n_heads</span><span class="p">),</span>
        <span class="ss">attn_drop_rate:</span> <span class="n">dropout_or_default!</span><span class="p">(</span><span class="ss">:attn_drop_rate</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">attn_drop_rate</span><span class="p">,</span> <span class="n">drop_rate</span><span class="p">),</span>
        <span class="ss">shortcut_drop_rate:</span>
          <span class="n">dropout_or_default!</span><span class="p">(</span><span class="ss">:shortcut_drop_rate</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">shortcut_drop_rate</span><span class="p">,</span> <span class="n">drop_rate</span><span class="p">),</span>
        <span class="ss">qkv_bias:</span> <span class="n">validate_boolean!</span><span class="p">(</span><span class="ss">:qkv_bias</span><span class="p">,</span> <span class="n">params</span><span class="o">.</span><span class="n">qkv_bias</span><span class="p">)</span>
      <span class="p">}</span>
    <span class="k">end</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">cfg_value</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span> <span class="k">do</span>
    <span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="no">Map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="no">Atom</span><span class="o">.</span><span class="n">to_string</span><span class="p">(</span><span class="n">key</span><span class="p">)))</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">cfg_value!</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span> <span class="k">do</span>
    <span class="n">cfg_value</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="n">_field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">,</span> <span class="k">do</span><span class="p">:</span> <span class="n">value</span>

  <span class="k">defp</span> <span class="n">validate_positive_integer!</span><span class="p">(</span><span class="n">field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected </span><span class="si">#{</span><span class="n">field</span><span class="si">}</span><span class="s2"> to be a positive integer, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">dropout_or_default!</span><span class="p">(</span><span class="n">_field</span><span class="p">,</span> <span class="no">nil</span><span class="p">,</span> <span class="n">default</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">default</span>
  <span class="k">defp</span> <span class="n">dropout_or_default!</span><span class="p">(</span><span class="n">field</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">_default</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>

  <span class="k">defp</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">_field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_number</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">value</span> <span class="o">&lt;</span> <span class="mi">1</span><span class="p">,</span>
    <span class="k">do</span><span class="p">:</span> <span class="n">value</span> <span class="o">*</span> <span class="mf">1.0</span>

  <span class="k">defp</span> <span class="n">validate_dropout!</span><span class="p">(</span><span class="n">field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected </span><span class="si">#{</span><span class="n">field</span><span class="si">}</span><span class="s2"> to be a number in [0, 1), got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">validate_boolean!</span><span class="p">(</span><span class="n">_field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_boolean</span><span class="p">(</span><span class="n">value</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">value</span>

  <span class="k">defp</span> <span class="n">validate_boolean!</span><span class="p">(</span><span class="n">field</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"expected </span><span class="si">#{</span><span class="n">field</span><span class="si">}</span><span class="s2"> to be a boolean, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="no">nil</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="no">System</span><span class="o">.</span><span class="n">unique_integer</span><span class="p">([</span><span class="ss">:positive</span><span class="p">])</span>
  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="ow">when</span> <span class="n">is_integer</span><span class="p">(</span><span class="n">seed</span><span class="p">),</span> <span class="k">do</span><span class="p">:</span> <span class="n">seed</span>

  <span class="k">defp</span> <span class="n">normalize_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span> <span class="k">do</span>
    <span class="k">raise</span> <span class="no">ArgumentError</span><span class="p">,</span> <span class="s2">"seed must be an integer or nil, got: </span><span class="si">#{</span><span class="n">inspect</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
  <span class="k">end</span>
<span class="k">end</span>
</code></pre></div></div>

<p>And test that is using those new dropouts:</p>

<div class="language-elixir highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="n">test</span> <span class="s2">"exercise 4.3 supports separate dropout rates for GPT dropout sites"</span> <span class="k">do</span>
    <span class="n">cfg</span> <span class="o">=</span> <span class="p">%</span><span class="no">GPTConfig</span><span class="p">{</span>
      <span class="ss">vocab_size:</span> <span class="mi">32</span><span class="p">,</span>
      <span class="ss">context_length:</span> <span class="mi">8</span><span class="p">,</span>
      <span class="ss">emb_dim:</span> <span class="mi">12</span><span class="p">,</span>
      <span class="ss">n_heads:</span> <span class="mi">3</span><span class="p">,</span>
      <span class="ss">n_layers:</span> <span class="mi">2</span><span class="p">,</span>
      <span class="ss">drop_rate:</span> <span class="mf">0.1</span><span class="p">,</span>
      <span class="ss">emb_drop_rate:</span> <span class="mf">0.2</span><span class="p">,</span>
      <span class="ss">shortcut_drop_rate:</span> <span class="mf">0.3</span><span class="p">,</span>
      <span class="ss">attn_drop_rate:</span> <span class="mf">0.4</span><span class="p">,</span>
      <span class="ss">qkv_bias:</span> <span class="no">false</span>
    <span class="p">}</span>

    <span class="n">model</span> <span class="o">=</span> <span class="no">GPTModel</span><span class="o">.</span><span class="n">new</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="ss">seed:</span> <span class="mi">123</span><span class="p">)</span>
    <span class="p">[</span><span class="n">first_block</span> <span class="o">|</span> <span class="n">_rest</span><span class="p">]</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">trf_blocks</span>

    <span class="n">assert</span> <span class="n">model</span><span class="o">.</span><span class="n">drop_emb</span> <span class="o">==</span> <span class="mf">0.2</span>
    <span class="n">assert</span> <span class="n">first_block</span><span class="o">.</span><span class="n">drop_shortcut</span> <span class="o">==</span> <span class="mf">0.3</span>
    <span class="n">assert</span> <span class="n">first_block</span><span class="o">.</span><span class="n">att</span><span class="o">.</span><span class="n">dropout</span> <span class="o">==</span> <span class="mf">0.4</span>
    <span class="n">assert</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">embedding_dropout</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="o">==</span> <span class="mf">0.2</span>
    <span class="n">assert</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">shortcut_dropout</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="o">==</span> <span class="mf">0.3</span>
    <span class="n">assert</span> <span class="no">GPTConfig</span><span class="o">.</span><span class="n">attention_dropout</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span> <span class="o">==</span> <span class="mf">0.4</span>
  <span class="k">end</span>
</code></pre></div></div>

<h2 id="conclusion">Conclusion</h2>

<p>This chapter was not as hard as the previous one on the attention mechanism. I started combining the pieces together and got the first gibberish text from my untrained GPTModel!</p>]]></content><author><name></name></author><category term="llm-from-scratch" /><category term="gpt" /><category term="tutorials" /><category term="llm" /><category term="gpt" /><category term="workbook" /><summary type="html"><![CDATA[This post implements a GPT-style model from scratch in Elixir/Nx.]]></summary></entry><entry><title type="html">AWS error: not authorized to perform</title><link href="https://karlosmid.com/2026/05/amazon-error-not-authorized-to-perform/" rel="alternate" type="text/html" title="AWS error: not authorized to perform" /><published>2026-05-01T00:00:00+00:00</published><updated>2026-05-01T00:00:00+00:00</updated><id>https://karlosmid.com/2026/05/amazon-error-not-authorized-to-perform</id><content type="html" xml:base="https://karlosmid.com/2026/05/amazon-error-not-authorized-to-perform/"><![CDATA[<h2 id="tldr">TL;DR</h2>

<p>How I resolved the cryptic AWS error <code class="language-plaintext highlighter-rouge">User: arn:aws:iam::user_id:user/user_name is not authorized to perform: connect:* on resource: * with an explicit deny in a resource-based policy</code> using some tricks from context-driven testing school.</p>

<h2 id="the-problem">The Problem</h2>

<p>When you get this error, the first thing is not to panic. Calm down and start exploring, because every error that you receive in your daily work is a learning opportunity. We got this error when we made an API call to search contacts in the Amazon Connect service. Connect is the name of the service. The error message led me to think that we had a problem with the Amazon Connect policy settings.</p>

<h2 id="the-investigation">The Investigation</h2>

<p>The obvious choice was to put the existing policy into an LLM. Do not forget to strip out your AWS account info, such as the user ID and user name. But I had already gone through a round of that conversation with an LLM, because I first got a pretty specific policy error with no general information with <code class="language-plaintext highlighter-rouge">*</code>. Also, other AWS Connect API calls were successful. Something was specific to the search contacts API call. So I used an old trick from context-driven testing school: a comparable product. I asked the LLM to produce an AWS CLI call with the same parameters for search contacts:</p>

<pre><code class="language-aws">aws connect search-contacts \\n  --region eu-central-1 \\n  --instance-id instance_id \\n  --time-range \\n    Type=INITIATION_TIMESTAMP,\\nStartTime=2026-03-10T07:52:44Z,\\nEndTime=2026-04-09T07:52:44Z \\n  --max-results 5
</code></pre>
<p>And I got a response, so I could rule out the Amazon Connect policy issue. Something was different with the AWS API call that we made for search contacts. The next step was the Amazon CloudTrail service, where you can see all requests, both accepted and rejected. I quickly noticed the difference in this part:</p>

<div class="language-json highlighter-rouge"><div class="highlight"><pre class="highlight"><code><span class="nl">"requestParameters"</span><span class="p">:</span><span class="w"> </span><span class="p">{</span><span class="w">
       </span><span class="nl">"TimeRange"</span><span class="p">:</span><span class="w"> </span><span class="p">{</span><span class="w">
           </span><span class="nl">"Type"</span><span class="p">:</span><span class="w"> </span><span class="s2">"INITIATION_TIMESTAMP"</span><span class="p">,</span><span class="w">
           </span><span class="nl">"StartTime"</span><span class="p">:</span><span class="w"> </span><span class="mi">1773129164</span><span class="p">,</span><span class="w">
           </span><span class="nl">"EndTime"</span><span class="p">:</span><span class="w"> </span><span class="mi">1775721164</span><span class="w">
       </span><span class="p">},</span><span class="w">
</span></code></pre></div></div>

<p>Above snippet show aws cli successful call. Our API call had string values for <code class="language-plaintext highlighter-rouge">StartTime</code> and <code class="language-plaintext highlighter-rouge">EndTime</code>, not Linux epoch timestamps as integers. Notice that the <code class="language-plaintext highlighter-rouge">aws connect</code> CLI also used string values, but the CLI correctly transformed those into epoch integers.</p>

<h2 id="conclusion">Conclusion</h2>

<p>AWS responded with a cryptic error. An error that pointed to <code class="language-plaintext highlighter-rouge">StartTime</code> and <code class="language-plaintext highlighter-rouge">EndTime</code> would be more helpful, but on the other hand, we should consult the AWS documentation. I guess AWS responded with kitchen-sink error handling code, where they have one big try/catch block, and for each error case they respond with the same error. Your takeaway is to use the comparable product trick and observability. For AWS, this is the CloudTrail product. Give logs a minute or two to show up in CloudTrail.</p>]]></content><author><name></name></author><category term="aws" /><category term="policies" /><category term="api" /><category term="aws" /><category term="policies" /><category term="api" /><summary type="html"><![CDATA[A misleading Amazon Connect authorization error turned out to be an API request formatting issue. Comparing the call with AWS CLI output and checking CloudTrail revealed that StartTime and EndTime needed epoch timestamp values.]]></summary></entry></feed>