explain-10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3189ca11-1f29-4be9-b0e7-22a99b80aeb8",
   "metadata": {},
   "source": [
    "### Now we have covered *MOST* of the important technobabble in AI\n",
    "  - Easy! BUT!\n",
    "#### RESPECT THE CODE!\n",
    "  - There are plenty of gotchas!\n",
    "#### The simplified formulations I show **are** the basis for the industrial scale code being used by OpoenAI, Google, Meta, etc...\n",
    "#### ~~The Devil~~ Bugs are in the details!!\n",
    "  - You can get pretty far without understanding the math, and adapting the boiler-plate code for your purposes, BUT!\n",
    "  - You **WILL** write buggy code, which has fundamental *math* flaws.\n",
    "  - You **WILL** write buggy code, that uses pytorch classes incorrectly.\n",
    "  - Only way to NOT waste weeks/months of effort is to (slowly?) **LEARN THE MATH**\n",
    "  - And figure out what each pytorch function you use is actually doing!\n",
    "  - **Use the source, Luke!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f60a602-ae2d-494b-b5d8-0db75072299b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch import nn\n",
    "from tinytorch.MyStuff import *\n",
    "from tinytorch.tensorhelpers import *\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from matplotlib.patches import Ellipse\n",
    "from matplotlib.text import OffsetFrom\n",
    "import os\n",
    "import os.path as path\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42fa8af7-cd51-4e54-8005-de0e3b2fe988",
   "metadata": {},
   "source": [
    "### Now for SimpleLlama3/LLama 3.1/3.2\n",
    "  - Arguably one of the top open LLM models out there.\n",
    "  - Uses a slightly modified `transformer` architecture\n",
    "  - Meta did a service to humanity by producing this model, and releasing it for anyone to use and modify.\n",
    "  - Encouraged Microsoft to do the same with their Phi models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b0dbc73-92c0-432f-8dd4-b35214ddaf3c",
   "metadata": {},
   "source": [
    "### (https://github.com/meta-llama/llama-models.git)\n",
    "  - `llama-models/models/llama3/api/args.py`       -  `ModelArgs`\n",
    "  - `llama-models/models/llama3/api/tokenizer.model`  - the tokenizer model file!\n",
    "  - `llama-models/models/llama3/api/tokenizer.py`  -  the tokenizer class!\n",
    "  - `llama-models/models/llama3/reference_impl/model.py` - the code\n",
    "  - `llama-models/models/llama3/reference_impl/generation.py` <- build/load the model\n",
    "#### Download the -Instruct models! (3.2, 1B or 3B only) or llama3.0 (8B only) from (https://www.llama.com/llama-downloads/)\n",
    "  - You will need a supported gpu that works with pytorch (almost any nvidia card, with at least 16gb VRAM)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d26e0af-d200-4504-98ff-21e775e8ceb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass, asdict\n",
    "from typing import Optional\n",
    "import inspect\n",
    "\n",
    "from LogRelay import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a68ffa2-a833-41e0-890b-e8b0ee8de169",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelArgs:\n",
    "  dim: int = 4096\n",
    "  n_layers: int = 32\n",
    "  n_heads: int = 32\n",
    "  n_kv_heads: Optional[int] = None\n",
    "  vocab_size: int = -1 # Later set in the build method\n",
    "  multiple_of: int = 256\n",
    "  ffn_dim_multiplier: Optional[float] = None\n",
    "  norm_eps: float = 1e-5\n",
    "  rope_theta: float = 500000\n",
    "  # Needed for KV cache\n",
    "  max_batch_size: int = 32\n",
    "  max_seq_len: int = 2048 # Llama3 has 8192. LLama 3.1 suppports upto 128K\n",
    "  device: str = None\n",
    "  bias: bool = False\n",
    "  dropout: float = 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9b1ec58-036b-471d-9927-49463ff32307",
   "metadata": {},
   "source": [
    "### FeedForward layer\n",
    "  - [torch.nn.Module](https://pytorch.org/docs/2.2/generated/torch.nn.Module.html#torch.nn.Module)\n",
    "  - [torch.nn.Linear](https://pytorch.org/docs/2.2/generated/torch.nn.Linear.html)\n",
    "  - [torch.nn.parameter.Parameter](https://pytorch.org/docs/2.2/generated/torch.nn.parameter.Parameter.html)\n",
    "  - [torch.nn.functinoal.silu](https://pytorch.org/docs/2.2/generated/torch.nn.functional.silu.html)\n",
    "    - [The SILU paper](https://arxiv.org/abs/1702.03118)\n",
    "    - [GELU](https://arxiv.org/abs/1606.08415)\n",
    "    - [SwISH](https://arxiv.org/abs/1710.05941v1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ded5a22e-4fb5-447a-bbd8-8bf2918cce0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Why are we using '_' in names? so that we can load up the SimpleLlama3 model without changes!\n",
    "class FeedForward(nn.Module):\n",
    "  LG = startDebug(__name__)\n",
    "  def __init__(\n",
    "    self,\n",
    "    args: ModelArgs\n",
    "  ):\n",
    "    nn.Module.__init__(self)\n",
    "\n",
    "    hidden_dim = 4 * args.dim\n",
    "    hidden_dim = int(2 * hidden_dim / 3)\n",
    "    if args.ffn_dim_multiplier is not None:\n",
    "      hidden_dim = int(args.ffn_dim_multiplier * hidden_dim)\n",
    "    # Round the hidden_dim to the nearest multiple of the multiple_of parameter\n",
    "    hidden_dim = args.multiple_of * ((hidden_dim + args.multiple_of - 1) // args.multiple_of)\n",
    "\n",
    "    self.w1 = nn.Linear(args.dim, hidden_dim, bias=False)\n",
    "    self.w2 = nn.Linear(hidden_dim, args.dim, bias=False)\n",
    "    self.w3 = nn.Linear(args.dim, hidden_dim, bias=False)\n",
    "    # assert str(self.w3.device).startswith('cuda')\n",
    "    \n",
    "  def forward(self, x: torch.Tensor):\n",
    "    # (B, Seq_Len, Dim) --> (B, Seq_Len, Hidden_Dim)\n",
    "    swish = F.silu(self.w1(x))\n",
    "    # (B, Seq_Len, Dim) --> (B, Seq_Len, Hidden_Dim)\n",
    "    x_V = self.w3(x)\n",
    "    # (B, Seq_Len, Hidden_Dim) * (B, Seq_Len, Hidden_Dim) --> (B, Seq_Len, Hidden_Dim)\n",
    "    x = swish * x_V\n",
    "    # (B, Seq_Len, Hidden_Dim) --> (B, Seq_Len, Dim)\n",
    "    x = self.w2(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ea94e82-d66a-4a32-9f89-f23a80476741",
   "metadata": {},
   "source": [
    "### For Rotary Positional Encoding\n",
    "  - [torch.outer](https://pytorch.org/docs/2.2/generated/torch.outer.html)\n",
    "  - [torch.arange](https://pytorch.org/docs/2.2/generated/torch.arange.html)\n",
    "  - [torch.polar](https://pytorch.org/docs/2.2/generated/torch.polar.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0516ce4c-99d0-4e43-8b88-235a3a9042e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def precompute_theta_pos_frequencies(head_dim: int, seq_len: int, theta: float = 10000.0, device: str = ''):\n",
    "  # As written in the paragraph 3.2.2 of the paper\n",
    "  # >> In order to generalize our results in 2D to any xi ∈ Rd where **d is even**, [...]\n",
    "  assert head_dim % 2 == 0, \"Dimension must be divisible by 2\"\n",
    "  # Build the theta parameter\n",
    "  # According to the formula theta_i = 10000^(-2(i-1)/dim) for i = [1, 2, ... dim/2]\n",
    "  # Shape: (Head_Dim / 2)\n",
    "  theta_numerator = torch.arange(0, head_dim, 2).float()\n",
    "  # Shape: (Head_Dim / 2)\n",
    "  theta = 1.0 / (theta ** (theta_numerator / head_dim)).to(device) # (Dim / 2)\n",
    "  # Construct the positions (the \"m\" parameter)\n",
    "  # Shape: (Seq_Len)\n",
    "  m = torch.arange(seq_len, device=device)\n",
    "  # Multiply each theta by each position using the outer product.\n",
    "  # Shape: (Seq_Len) outer_product* (Head_Dim / 2) -> (Seq_Len, Head_Dim / 2)\n",
    "  freqs = torch.outer(m, theta).float()\n",
    "  # We can compute complex numbers in the polar form c = R * exp(m * theta), where R = 1 as follows:\n",
    "  # (Seq_Len, Head_Dim / 2) -> (Seq_Len, Head_Dim / 2)\n",
    "  freqs_complex = torch.polar(torch.ones_like(freqs), freqs)\n",
    "  return freqs_complex\n",
    "\n",
    "def apply_rotary_embeddings(x: torch.Tensor, freqs_complex: torch.Tensor, device: str):\n",
    "  # Separate the last dimension pairs of two values, representing the real and imaginary parts of the complex number\n",
    "  # Two consecutive values will become a single complex number\n",
    "  # (B, Seq_Len, H, Head_Dim) -> (B, Seq_Len, H, Head_Dim/2)\n",
    "  x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))\n",
    "  # Reshape the freqs_complex tensor to match the shape of the x_complex tensor. So we need to add the batch dimension and the head dimension\n",
    "  # (Seq_Len, Head_Dim/2) --> (1, Seq_Len, 1, Head_Dim/2)\n",
    "  freqs_complex = freqs_complex.unsqueeze(0).unsqueeze(2).to(device)\n",
    "  # Multiply each complex number in the x_complex tensor by the corresponding complex number in the freqs_complex tensor\n",
    "  # Which results in the rotation of the complex number as shown in the Figure 1 of the paper\n",
    "  # (B, Seq_Len, H, Head_Dim/2) * (1, Seq_Len, 1, Head_Dim/2) = (B, Seq_Len, H, Head_Dim/2)\n",
    "  assert x_complex.device == freqs_complex.device, f'{x_complex.device=} == {freqs_complex.device=}'\n",
    "  x_rotated = x_complex * freqs_complex\n",
    "  # Convert the complex number back to the real number\n",
    "  # (B, Seq_Len, H, Head_Dim/2) -> (B, Seq_Len, H, Head_Dim/2, 2)\n",
    "  x_out = torch.view_as_real(x_rotated)\n",
    "  # (B, Seq_Len, H, Head_Dim/2, 2) -> (B, Seq_Len, H, Head_Dim)\n",
    "  x_out = x_out.reshape(*x.shape)\n",
    "  return x_out.type_as(x).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84ba83b9-e38b-43ff-a956-2caeb4189013",
   "metadata": {},
   "source": [
    "### RMSNorm Layer\n",
    "  - [LLama suggests RMSNorm instead](https://arxiv.org/abs/1910.07467)\n",
    "  - [torch.nn.Parameter](https://pytorch.org/docs/2.2/generated/torch.nn.parameter.Parameter.html)\n",
    "  - [torch.rsqrt](https://pytorch.org/docs/2.2/generated/torch.rsqrt.html)\n",
    "  - [torch.pow](https://pytorch.org/docs/2.2/generated/torch.pow.html)\n",
    "  - [torch.mean](https://pytorch.org/docs/2.2/generated/torch.mean.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c807fe5a-1e58-4b4e-9e34-c823d5583ef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RMSNorm(nn.Module):\n",
    "  LG = startDebug(__name__)\n",
    "  def __init__(self, dim: int, eps: float = 1e-6):\n",
    "    nn.Module.__init__(self)\n",
    "    self.eps = eps\n",
    "    # The gamma parameter\n",
    "    self.weight = nn.Parameter(torch.ones(dim))\n",
    "\n",
    "  def _norm(self, x: torch.Tensor):\n",
    "    # (B, Seq_Len, Dim) * (B, Seq_Len, 1) = (B, Seq_Len, Dim)\n",
    "    # rsqrt: 1 / sqrt(x)\n",
    "    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)\n",
    "\n",
    "  def forward(self, x: torch.Tensor):\n",
    "    # (Dim) * (B, Seq_Len, Dim) = (B, Seq_Len, Dim)\n",
    "    return self.weight * self._norm(x.float()).type_as(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c1ef048-d204-4811-b9db-ec8702184206",
   "metadata": {},
   "source": [
    "### repeat_kv (used for Grouped Multi-Head Attention\n",
    "  - [torch.Tensor.expand](https://pytorch.org/docs/stable/generated/torch.Tensor.expand)\n",
    "  - [torch.Tensor.reshape](https://pytorch.org/docs/stable/generated/torch.Tensor.reshape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5710c4f-7b8d-436a-a3f7-e96aa072eed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:\n",
    "  batch_size, seq_len, n_kv_heads, head_dim = x.shape\n",
    "  if n_rep == 1:\n",
    "    return x\n",
    "  return (\n",
    "    # (B, Seq_Len, N_KV_Heads, 1, Head_Dim)\n",
    "    x[:, :, :, None, :]\n",
    "    # (B, Seq_Len, N_KV_Heads, N_Rep, Head_Dim)\n",
    "    .expand(batch_size, seq_len, n_kv_heads, n_rep, head_dim)\n",
    "    # (B, Seq_Len, N_KV_Heads * N_Rep, Head_Dim)\n",
    "    .reshape(batch_size, seq_len, n_kv_heads * n_rep, head_dim)\n",
    "  )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3995897e-bbd1-4e07-b9ae-8b9cdd13da8c",
   "metadata": {},
   "source": [
    "### SelfAttention Layer\n",
    "  - [torch.zeros](https://pytorch.org/docs/2.2/generated/torch.zeros.html)\n",
    "  - [torch.view](https://pytorch.org/docs/2.2/generated/torch.Tensor.view.html)\n",
    "  - [torch.transpose](https://pytorch.org/docs/2.2/generated/torch.transpose.html)\n",
    "  - [torch.nn.functional.softmax](https://pytorch.org/docs/2.2/generated/torch.nn.functional.softmax.html)\n",
    "  - CAREFUL! [Pytorch broadcasting semantics](https://pytorch.org/tutorials/beginner/introyt/tensors_deeper_tutorial.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72107414-484c-4cf5-bf5b-0641471a7f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SelfAttention(nn.Module):\n",
    "  LG = startDebug(__name__)\n",
    "  def __init__(self, args: ModelArgs):\n",
    "    nn.Module.__init__(self)\n",
    "    self.args = args\n",
    "    # Indicates the number of heads for the Keys and Values\n",
    "    self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads\n",
    "    # Indicates the number of heads for the Queries\n",
    "    self.n_heads_q = args.n_heads\n",
    "    # Indicates how many times the Keys and Values should be repeated\n",
    "    self.n_rep = self.n_heads_q // self.n_kv_heads\n",
    "    # Indicates the dimension of each head, that is, the part of the embedding that each head will be responsible for\n",
    "    self.head_dim = args.dim // args.n_heads\n",
    "    self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)\n",
    "    self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)\n",
    "    self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)\n",
    "    self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)\n",
    "    assert isinstance(args, ModelArgs)\n",
    "    assert args.device is not None\n",
    "    self.cache_k = torch.zeros((args.max_batch_size, args.max_seq_len, self.n_kv_heads, self.head_dim)).to(args.device)\n",
    "    self.cache_v = torch.zeros((args.max_batch_size, args.max_seq_len, self.n_kv_heads, self.head_dim)).to(args.device)\n",
    "  def forward(\n",
    "      self,\n",
    "      x: torch.Tensor,\n",
    "      start_pos: int,\n",
    "      freqs_complex: torch.Tensor,\n",
    "      mask: Optional[torch.Tensor]\n",
    "  ):\n",
    "    batch_size, seq_len, _ = x.shape  # (B, 1, Dim)\n",
    "\n",
    "    # (B, 1, Dim) -> (B, 1, H_Q * Head_Dim)\n",
    "    xq = self.wq(x)\n",
    "    # (B, 1, Dim) -> (B, 1, H_KV * Head_Dim)\n",
    "    xk = self.wk(x)\n",
    "    # (B, 1, Dim) -> (B, 1, H_KV * Head_Dim)\n",
    "    xv = self.wv(x)\n",
    "\n",
    "    # (B, 1, H_Q * Head_Dim) -> (B, 1, H_Q, Head_Dim)\n",
    "    xq = xq.view(batch_size, seq_len, self.n_heads_q, self.head_dim)\n",
    "    # (B, 1, H_KV * Head_Dim) -> (B, 1, H_KV, Head_Dim)\n",
    "    xk = xk.view(batch_size, seq_len, self.n_kv_heads, self.head_dim)\n",
    "    # (B, 1, H_KV * Head_Dim) -> (B, 1, H_KV, Head_Dim)\n",
    "    xv = xv.view(batch_size, seq_len, self.n_kv_heads, self.head_dim)\n",
    "\n",
    "    # (B, 1, H_Q, Head_Dim) --> (B, 1, H_Q, Head_Dim)\n",
    "    xq = apply_rotary_embeddings(xq, freqs_complex, device=x.device)\n",
    "    # (B, 1, H_KV, Head_Dim) --> (B, 1, H_KV, Head_Dim)\n",
    "    xk = apply_rotary_embeddings(xk, freqs_complex, device=x.device)\n",
    "\n",
    "    # Replace the entry in the cache\n",
    "    self.cache_k[:batch_size, start_pos : start_pos + seq_len] = xk\n",
    "    self.cache_v[:batch_size, start_pos : start_pos + seq_len] = xv\n",
    "\n",
    "    # (B, Seq_Len_KV, H_KV, Head_Dim)\n",
    "    keys = self.cache_k[:batch_size, : start_pos + seq_len]\n",
    "    # (B, Seq_Len_KV, H_KV, Head_Dim)\n",
    "    values = self.cache_v[:batch_size, : start_pos + seq_len]\n",
    "\n",
    "    # Since every group of Q shares the same K and V heads, just repeat the K and V heads for every Q in the same group.\n",
    "    # (B, Seq_Len_KV, H_KV, Head_Dim) --> (B, Seq_Len_KV, H_Q, Head_Dim)\n",
    "    keys = repeat_kv(keys, self.n_rep)\n",
    "    # (B, Seq_Len_KV, H_KV, Head_Dim) --> (B, Seq_Len_KV, H_Q, Head_Dim)\n",
    "    values = repeat_kv(values, self.n_rep)\n",
    "\n",
    "    # (B, 1, H_Q, Head_Dim) -> (B, H_Q, 1, Head_Dim)\n",
    "    xq = xq.transpose(1, 2)\n",
    "    # (B, Seq_Len_KV, H_Q, Head_Dim) -> (B, H_Q, Seq_Len_KV, Head_Dim)\n",
    "    keys = keys.transpose(1, 2)\n",
    "    # (B, Seq_Len_KV, H_Q, Head_Dim) -> (B, H_Q, Seq_Len_KV, Head_Dim)\n",
    "    values = values.transpose(1, 2)\n",
    "    # (B, H_Q, 1, Head_Dim) @ (B, H_Q, Head_Dim, Seq_Len_KV) -> (B, H_Q, 1, Seq_Len_KV)\n",
    "    scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)\n",
    "    if mask is not None:\n",
    "      scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)\n",
    "\n",
    "    # (B, H_Q, 1, Seq_Len_KV) -> (B, H_Q, 1, Seq_Len_KV)\n",
    "    scores = F.softmax(scores.float(), dim=-1).type_as(xq)\n",
    "\n",
    "    # (B, H_Q, 1, Seq_Len) @ (B, H_Q, Seq_Len_KV, Head_Dim) -> (B, H_Q, 1, Head_Dim)\n",
    "    output = torch.matmul(scores, values)\n",
    "    # (B, H_Q, 1, Head_Dim) -> (B, 1, H_Q, Head_Dim) -> (B, 1, Dim)\n",
    "    output = (output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1))\n",
    "    return self.wo(output) # (B, 1, Dim) -> (B, 1, Dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "499c19ce-46f6-47f0-9832-dddeb207ca75",
   "metadata": {},
   "source": [
    "### EncoderBlock\n",
    "  - SelfAttention\n",
    "  - FeedForward\n",
    "  - RMSNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4b87d48-31e4-4603-a7a3-840157cfd933",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "  LG = startDebug(__name__)\n",
    "  def __init__(self, args: ModelArgs):\n",
    "    nn.Module.__init__(self)\n",
    "\n",
    "    self.n_heads = args.n_heads\n",
    "    self.dim = args.dim\n",
    "    self.head_dim = args.dim // args.n_heads\n",
    "\n",
    "    self.attention = SelfAttention(args)\n",
    "    self.feed_forward = FeedForward(args)\n",
    "\n",
    "    # Normalization BEFORE the attention block\n",
    "    self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)\n",
    "    # Normalization BEFORE the feed forward block\n",
    "    self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)\n",
    "  \n",
    "  def forward(self, x: torch.Tensor, start_pos: int, freqs_complex: torch.Tensor,\n",
    "              mask: Optional[torch.Tensor]):\n",
    "    assert isinstance(x, torch.Tensor) and isinstance(start_pos, int)\n",
    "    # assert x.device == mask.device, f'{x.device=} == {mask.device=}'\n",
    "    # (B, Seq_Len, Dim) + (B, Seq_Len, Dim) --> (B, Seq_Len, Dim)\n",
    "    h = x + self.attention.forward(\n",
    "      self.attention_norm(x), start_pos, freqs_complex, mask\n",
    "    )\n",
    "    # (B, Seq_Len, Dim) + (B, Seq_Len, Dim) --> (B, Seq_Len, Dim)\n",
    "    out = h + self.feed_forward.forward(self.ffn_norm(h))\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "077d4c19-6273-49bd-8254-c49d36f96db3",
   "metadata": {},
   "source": [
    "### Transformer\n",
    "  - Uses all of the modules defined above\n",
    "  - [torch.nn.Embedding](https://pytorch.org/docs/2.2/generated/torch.nn.Embedding.html)\n",
    "  - [torch.nn.functional.softmax](https://pytorch.org/docs/2.2/generated/torch.nn.functional.softmax.html)\n",
    "  - [torch.nn.ModuleList](https://pytorch.org/docs/2.2/generated/torch.nn.ModuleList.html)\n",
    "  - [torch.hstack](https://pytorch.org/docs/2.2/generated/torch.hstack.html)\n",
    "  - [torch.cat](https://pytorch.org/docs/2.2/generated/torch.cat.html)\n",
    "  - [torch.full](https://pytorch.org/docs/2.2/generated/torch.full.html)\n",
    "  - [torch.triu](https://pytorch.org/docs/2.2/generated/torch.triu.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9ae8b1e-508f-4a17-9500-91dffb70da9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "### quick example of torch.cat and torch.stack\n",
    "A = torch.arange(6).view(3,-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c34233a7-18cc-4889-b319-d1c43876875c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(Shapes(A=A,cat0=torch.cat([A, A], dim=0), cat1=torch.cat([A, A], dim=1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc46e9be-4dc6-43e7-9bbf-e65e28ca914a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(Shapes(A=A, stack0=torch.stack([A, A], dim=0), stack1=torch.stack([A, A], dim=1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "138a5c17-90e7-4592-804d-4ad938099228",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "  LG = startDebug(__name__)\n",
    "  def __init__(self, args: ModelArgs):\n",
    "    nn.Module.__init__(self)\n",
    "    setupLogRelay(self)\n",
    "    assert args.device is not None\n",
    "    self.info(f'Transformer initializing with {asdict(args)}')\n",
    "    assert args.vocab_size != -1, \"Vocab size must be set\"\n",
    "    self.params = args\n",
    "    self.config = args\n",
    "    self.args = args\n",
    "    self.vocab_size = args.vocab_size\n",
    "    self.n_layers = args.n_layers\n",
    "    self.tok_embeddings = nn.Embedding(self.vocab_size, args.dim)\n",
    "\n",
    "    self.layers = nn.ModuleList()\n",
    "    for layer_id in range(args.n_layers):\n",
    "      self.layers.append(EncoderBlock(args))\n",
    "\n",
    "    self.norm = RMSNorm(args.dim, eps=args.norm_eps)\n",
    "    self.output = nn.Linear(args.dim, self.vocab_size, bias=False)\n",
    "\n",
    "    self.freqs_complex = precompute_theta_pos_frequencies(self.args.dim // self.args.n_heads, self.args.max_seq_len * 2,\n",
    "                                args.rope_theta,\n",
    "                                device=self.args.device)\n",
    "    self.info(\"number of parameters: %.2fM\" % (self.get_num_params()/1e6,))\n",
    "\n",
    "  def forward(self, tokens: torch.Tensor, start_pos: int, targets = None):\n",
    "    # (B, Seq_Len)\n",
    "    batch_size, seq_len = tokens.shape\n",
    "    # assert seq_len > 1, f\"Only one token at a time can be processed {seq_len=}\"\n",
    "\n",
    "    # (B, Seq_Len) -> (B, Seq_Len, Dim)\n",
    "    h = self.tok_embeddings(tokens)\n",
    "    assert isinstance(h, torch.Tensor)\n",
    "    assert isinstance(start_pos, int), f'{start_pos=} not int'\n",
    "    # Retrieve the pairs (m, theta) corresponding to the positions [start_pos, start_pos + seq_len]\n",
    "    freqs_complex = self.freqs_complex[start_pos:start_pos + seq_len]\n",
    "\n",
    "    mask = None\n",
    "    if seq_len > 1:\n",
    "      mask = torch.full((seq_len, seq_len), float(\"-inf\"), device=tokens.device)\n",
    "\n",
    "      mask = torch.triu(mask, diagonal=1)\n",
    "\n",
    "      # When performing key-value caching, we compute the attention scores\n",
    "      # only for the new sequence. Thus, the matrix of scores is of size\n",
    "      # (seq_len, cache_len + seq_len), and the only masked entries are (i, j) for\n",
    "      # j > cache_len + i, since row i corresponds to token cache_len + i.\n",
    "      mask = torch.hstack(\n",
    "        [torch.zeros((seq_len, start_pos), device=tokens.device), mask]\n",
    "      ).type_as(h)\n",
    "\n",
    "    # Consecutively apply all the encoder layers\n",
    "    for layer in self.layers:\n",
    "      h = layer(h, start_pos, freqs_complex, mask)\n",
    "    h = self.norm(h)\n",
    "    output = self.output(h).float()\n",
    "    loss = None\n",
    "    if targets is not None:\n",
    "      logits = output\n",
    "      loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)\n",
    "\n",
    "    return output, loss\n",
    "\n",
    "  def get_num_params(self, non_embedding=True):\n",
    "    \"\"\"\n",
    "    Return the number of parameters in the model.\n",
    "    For non-embedding count (default), the position embeddings get subtracted.\n",
    "    The token embeddings would too, except due to the parameter sharing these\n",
    "    params are actually used as weights in the final layer, so we include them.\n",
    "    \"\"\"\n",
    "    n_params = sum(p.numel() for p in self.parameters())\n",
    "    # if non_embedding:\n",
    "    #   n_params -= self.transformer.wpe.weight.numel()\n",
    "    return n_params\n",
    "\n",
    "  def _init_weights(self, module):\n",
    "    if isinstance(module, nn.Linear):\n",
    "      torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "      if module.bias is not None:\n",
    "        torch.nn.init.zeros_(module.bias)\n",
    "    elif isinstance(module, nn.Embedding):\n",
    "      torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c19841a-ac81-40ae-9ad7-ffd2525d8dfb",
   "metadata": {},
   "source": [
    "### Now for the tokenizer\n",
    "  - SimpleLlama3/SimpleLlama3.1 uses OpenAI's tiktokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d82ced5-8cd2-4ee7-ab08-1018cb5c5d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "from logging import getLogger\n",
    "from pathlib import Path\n",
    "from typing import (\n",
    "  AbstractSet,\n",
    "  cast,\n",
    "  Collection,\n",
    "  Dict,\n",
    "  Iterator,\n",
    "  List,\n",
    "  Literal,\n",
    "  Sequence,\n",
    "  TypedDict,\n",
    "  Union,\n",
    ")\n",
    "\n",
    "import tiktoken\n",
    "from tiktoken.load import load_tiktoken_bpe\n",
    "\n",
    "\n",
    "logger = getLogger(__name__)\n",
    "\n",
    "\n",
    "Role = Literal[\"system\", \"user\", \"assistant\"]\n",
    "\n",
    "\n",
    "class Message(TypedDict):\n",
    "  role: Role\n",
    "  content: str\n",
    "\n",
    "Dialog = Sequence[Message]\n",
    "\n",
    "\n",
    "class Tokenizer:\n",
    "  \"\"\"\n",
    "  Tokenizing and encoding/decoding text using the Tiktoken tokenizer.\n",
    "  \"\"\"\n",
    "\n",
    "  special_tokens: Dict[str, int]\n",
    "\n",
    "  num_reserved_special_tokens = 256\n",
    "\n",
    "  pat_str = r\"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+\"  # noqa: E501\n",
    "\n",
    "  def __init__(self, model_path: str):\n",
    "    \"\"\"\n",
    "    Initializes the Tokenizer with a Tiktoken model.\n",
    "\n",
    "    Args:\n",
    "      model_path (str): The path to the Tiktoken model file.\n",
    "    \"\"\"\n",
    "    assert os.path.isfile(model_path), model_path\n",
    "\n",
    "    mergeable_ranks = load_tiktoken_bpe(model_path)\n",
    "    num_base_tokens = len(mergeable_ranks)\n",
    "    special_tokens = [\n",
    "      \"<|begin_of_text|>\",\n",
    "      \"<|end_of_text|>\",\n",
    "      \"<|reserved_special_token_0|>\",\n",
    "      \"<|reserved_special_token_1|>\",\n",
    "      \"<|reserved_special_token_2|>\",\n",
    "      \"<|reserved_special_token_3|>\",\n",
    "      \"<|start_header_id|>\",\n",
    "      \"<|end_header_id|>\",\n",
    "      \"<|reserved_special_token_4|>\",\n",
    "      \"<|eot_id|>\",  # end of turn\n",
    "    ] + [\n",
    "      f\"<|reserved_special_token_{i}|>\"\n",
    "      for i in range(5, self.num_reserved_special_tokens - 5)\n",
    "    ]\n",
    "    self.special_tokens = {\n",
    "      token: num_base_tokens + i for i, token in enumerate(special_tokens)\n",
    "    }\n",
    "    self.model = tiktoken.Encoding(\n",
    "      name=Path(model_path).name,\n",
    "      pat_str=self.pat_str,\n",
    "      mergeable_ranks=mergeable_ranks,\n",
    "      special_tokens=self.special_tokens,\n",
    "    )\n",
    "    logger.info(f\"Reloaded tiktoken model from {model_path}\")\n",
    "\n",
    "    self.n_words: int = self.model.n_vocab\n",
    "    # BOS / EOS token IDs\n",
    "    self.bos_id: int = self.special_tokens[\"<|begin_of_text|>\"]\n",
    "    self.eos_id: int = self.special_tokens[\"<|end_of_text|>\"]\n",
    "    self.pad_id: int = -1\n",
    "    self.stop_tokens = {\n",
    "      self.special_tokens[\"<|end_of_text|>\"],\n",
    "      self.special_tokens[\"<|eot_id|>\"],\n",
    "    }\n",
    "    logger.info(\n",
    "      f\"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}\"\n",
    "    )\n",
    "\n",
    "  def encode(self, s: str,  *, bos: bool, eos: bool,\n",
    "             allowed_special: Union[Literal[\"all\"], AbstractSet[str]] = set(),\n",
    "             disallowed_special: Union[Literal[\"all\"], Collection[str]] = (),\n",
    "            ) -> List[int]:\n",
    "    \"\"\"\n",
    "    Encodes a string into a list of token IDs.\n",
    "\n",
    "    Args:\n",
    "      s (str): The input string to be encoded.\n",
    "      bos (bool): Whether to prepend the beginning-of-sequence token.\n",
    "      eos (bool): Whether to append the end-of-sequence token.\n",
    "      allowed_tokens (\"all\"|set[str]): allowed special tokens in string\n",
    "      disallowed_tokens (\"all\"|set[str]): special tokens that raise an error when in string\n",
    "\n",
    "    Returns:\n",
    "      list[int]: A list of token IDs.\n",
    "\n",
    "    By default, setting disallowed_special=() encodes a string by ignoring\n",
    "    special tokens. Specifically:\n",
    "    - Setting `disallowed_special` to () will cause all text corresponding\n",
    "      to special tokens to be encoded as natural text (insteading of raising\n",
    "      an error).\n",
    "    - Setting `allowed_special` to \"all\" will treat all text corresponding\n",
    "      to special tokens to be encoded as special tokens.\n",
    "    \"\"\"\n",
    "    assert type(s) is str\n",
    "\n",
    "    # The tiktoken tokenizer can handle <=400k chars without\n",
    "    # pyo3_runtime.PanicException.\n",
    "    TIKTOKEN_MAX_ENCODE_CHARS = 400_000\n",
    "\n",
    "    # https://github.com/openai/tiktoken/issues/195\n",
    "    # Here we iterate over subsequences and split if we exceed the limit\n",
    "    # of max consecutive non-whitespace or whitespace characters.\n",
    "    MAX_NO_WHITESPACES_CHARS = 25_000\n",
    "\n",
    "    substrs = (\n",
    "      substr\n",
    "      for i in range(0, len(s), TIKTOKEN_MAX_ENCODE_CHARS)\n",
    "      for substr in self._split_whitespaces_or_nonwhitespaces(\n",
    "        s[i : i + TIKTOKEN_MAX_ENCODE_CHARS], MAX_NO_WHITESPACES_CHARS\n",
    "      )\n",
    "    )\n",
    "    t: List[int] = []\n",
    "    for substr in substrs:\n",
    "      t.extend(\n",
    "        self.model.encode(\n",
    "          substr,\n",
    "          allowed_special=allowed_special,\n",
    "          disallowed_special=disallowed_special,\n",
    "        )\n",
    "      )\n",
    "    if bos:\n",
    "      t.insert(0, self.bos_id)\n",
    "    if eos:\n",
    "      t.append(self.eos_id)\n",
    "    return t\n",
    "\n",
    "  def decode(self, t: Sequence[int]) -> str:\n",
    "    \"\"\"\n",
    "    Decodes a list of token IDs into a string.\n",
    "\n",
    "    Args:\n",
    "      t (List[int]): The list of token IDs to be decoded.\n",
    "\n",
    "    Returns:\n",
    "      str: The decoded string.\n",
    "    \"\"\"\n",
    "    # Typecast is safe here. Tiktoken doesn't do anything list-related with the sequence.\n",
    "    return self.model.decode(cast(List[int], t))\n",
    "\n",
    "  @staticmethod\n",
    "  def _split_whitespaces_or_nonwhitespaces(\n",
    "    s: str, max_consecutive_slice_len: int\n",
    "  ) -> Iterator[str]:\n",
    "    \"\"\"\n",
    "    Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len`\n",
    "    consecutive whitespaces or consecutive non-whitespaces.\n",
    "    \"\"\"\n",
    "    current_slice_len = 0\n",
    "    current_slice_is_space = s[0].isspace() if len(s) > 0 else False\n",
    "    slice_start = 0\n",
    "\n",
    "    for i in range(len(s)):\n",
    "      is_now_space = s[i].isspace()\n",
    "\n",
    "      if current_slice_is_space ^ is_now_space:\n",
    "        current_slice_len = 1\n",
    "        current_slice_is_space = is_now_space\n",
    "      else:\n",
    "        current_slice_len += 1\n",
    "        if current_slice_len > max_consecutive_slice_len:\n",
    "          yield s[slice_start:i]\n",
    "          slice_start = i\n",
    "          current_slice_len = 1\n",
    "    yield s[slice_start:]\n",
    "\n",
    "\n",
    "class ChatFormat:\n",
    "  def __init__(self, tokenizer: Tokenizer):\n",
    "    self.tokenizer = tokenizer\n",
    "\n",
    "  def encode_header(self, message: Message) -> List[int]:\n",
    "    tokens = []\n",
    "    tokens.append(self.tokenizer.special_tokens[\"<|start_header_id|>\"])\n",
    "    tokens.extend(self.tokenizer.encode(message[\"role\"], bos=False, eos=False))\n",
    "    tokens.append(self.tokenizer.special_tokens[\"<|end_header_id|>\"])\n",
    "    tokens.extend(self.tokenizer.encode(\"\\n\\n\", bos=False, eos=False))\n",
    "    return tokens\n",
    "\n",
    "  def encode_message(self, message: Message) -> List[int]:\n",
    "    tokens = self.encode_header(message)\n",
    "    tokens.extend(\n",
    "      self.tokenizer.encode(message[\"content\"].strip(), bos=False, eos=False)\n",
    "    )\n",
    "    tokens.append(self.tokenizer.special_tokens[\"<|eot_id|>\"])\n",
    "    return tokens\n",
    "\n",
    "  def encode_dialog_prompt(self, dialog: Dialog) -> List[int]:\n",
    "    tokens = []\n",
    "    tokens.append(self.tokenizer.special_tokens[\"<|begin_of_text|>\"])\n",
    "    for message in dialog:\n",
    "      tokens.extend(self.encode_message(message))\n",
    "    # Add the start of an assistant message for the model to complete.\n",
    "    tokens.extend(self.encode_header({\"role\": \"assistant\", \"content\": \"\"}))\n",
    "    return tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73f9b2ff-4232-4dad-b47c-d1063d7862dd",
   "metadata": {},
   "source": [
    "### Inference engine\n",
    " - [torch.multinomial](https://pytorch.org/docs/2.2/generated/torch.multinomial.html)\n",
    " - [torch.argmax](https://pytorch.org/docs/2.2/generated/torch.argmax.html)\n",
    " - [torch.where](https://pytorch.org/docs/2.2/generated/torch.where.html)\n",
    " - [torch.gather](https://pytorch.org/docs/2.2/generated/torch.gather.html)\n",
    " - [torch.cumsum](https://pytorch.org/docs/2.2/generated/torch.cumsum.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c474973-785c-405a-aa9e-d21730295208",
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from official Llama3 repository\n",
    "from typing import (\n",
    "  AbstractSet,\n",
    "  cast,\n",
    "  Collection,\n",
    "  Dict,\n",
    "  Iterator,\n",
    "  List,\n",
    "  Literal,\n",
    "  Sequence,\n",
    "  TypedDict,\n",
    "  Optional,\n",
    "  Union,\n",
    "  Tuple,\n",
    ")\n",
    "import torch\n",
    "import time\n",
    "from pathlib import Path\n",
    "import json\n",
    "# from sentencepiece import SentencePieceProcessor\n",
    "from tqdm import tqdm\n",
    "import tiktoken\n",
    "from tiktoken.load import load_tiktoken_bpe\n",
    "\n",
    "class CompletionPrediction(TypedDict, total=False):\n",
    "  generation: str\n",
    "  tokens: List[str]  # not required\n",
    "  logprobs: List[float]  # not required\n",
    "\n",
    "class ChatPrediction(TypedDict, total=False):\n",
    "  generation: Message\n",
    "  tokens: List[str]  # not required\n",
    "  logprobs: List[float]  # not required\n",
    "\n",
    "class SimpleLlama3:\n",
    "  LG = startDebug(__name__)\n",
    "  def __init__(self, model: Transformer, tokenizer: Tokenizer, model_args: ModelArgs):\n",
    "    setupLogRelay(self)\n",
    "    self.model = model\n",
    "    self.tokenizer = tokenizer\n",
    "    self.args = model_args\n",
    "    self.formatter = ChatFormat(tokenizer)\n",
    "\n",
    "  @staticmethod\n",
    "  def build(checkpoints_dir: str, tokenizer_path: str, load_model: bool, max_seq_len: int, max_batch_size: int, device: str):\n",
    "    prev_time = time.time()\n",
    "    if load_model:\n",
    "      checkpoints = sorted(Path(checkpoints_dir).glob(\"*.pth\"))\n",
    "      assert len(checkpoints) > 0, f\"no checkpoint files found in {checkpoints_dir}\"\n",
    "      ckpt_path = checkpoints[0]\n",
    "      print(f'Loading checkpoint \"{ckpt_path}\"')\n",
    "      checkpoint = torch.load(ckpt_path, map_location=\"cpu\")\n",
    "      print(f\"Loaded checkpoint in {time.time() - prev_time:.2f}s\")\n",
    "      prev_time = time.time()\n",
    "    with open(Path(checkpoints_dir) / \"params.json\", \"r\") as f:\n",
    "      params = json.loads(f.read())\n",
    "\n",
    "    model_args: ModelArgs = ModelArgs(\n",
    "      max_seq_len=max_seq_len,\n",
    "      max_batch_size=max_batch_size,\n",
    "      device=device,\n",
    "      **params\n",
    "    )\n",
    "\n",
    "    tokenizer = Tokenizer(model_path=tokenizer_path)\n",
    "    # tokenizer.load(tokenizer_path)\n",
    "    model_args.vocab_size = tokenizer.n_words\n",
    "    print(f'{model_args.vocab_size=}')\n",
    "    if device == \"cuda\":\n",
    "      torch.set_default_tensor_type(torch.cuda.HalfTensor)\n",
    "    else:\n",
    "      torch.set_default_tensor_type(torch.BFloat16Tensor)\n",
    "    \n",
    "    model = Transformer(model_args).to(device)\n",
    "\n",
    "    if load_model:\n",
    "      # The only unmatched key in the checkpoint is rope.freqs. Remove it\n",
    "      # del checkpoint['rope.freqs']\n",
    "      model.load_state_dict(checkpoint, strict=True)\n",
    "      print(f\"Loaded state dict in {time.time() - prev_time:.2f}s\")\n",
    "    \n",
    "    return SimpleLlama3(model, tokenizer, model_args)\n",
    "\n",
    "  @torch.inference_mode()\n",
    "  def generate(\n",
    "    self,\n",
    "    prompt_tokens: List[List[int]],\n",
    "    max_gen_len: int,\n",
    "    temperature: float = 0.6,\n",
    "    top_p: float = 0.9,\n",
    "    logprobs: bool = False,\n",
    "    echo: bool = False,\n",
    "  ) -> Tuple[List[List[int]], Optional[List[List[float]]]]:\n",
    "    \"\"\"\n",
    "    Generate text sequences based on provided prompts using the language generation model.\n",
    "\n",
    "    Args:\n",
    "      prompt_tokens (List[List[int]]): List of tokenized prompts, where each prompt is represented as a list of integers.\n",
    "      max_gen_len (int): Maximum length of the generated text sequence.\n",
    "      temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6.\n",
    "      top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to 0.9.\n",
    "      logprobs (bool, optional): Flag indicating whether to compute token log probabilities. Defaults to False.\n",
    "      echo (bool, optional): Flag indicating whether to include prompt tokens in the generated output. Defaults to False.\n",
    "\n",
    "    Returns:\n",
    "      Tuple[List[List[int]], Optional[List[List[float]]]]: A tuple containing generated token sequences and, if logprobs is True, corresponding token log probabilities.\n",
    "\n",
    "    Note:\n",
    "      This method uses the provided prompts as a basis for generating text. It employs nucleus sampling to produce text with controlled randomness.\n",
    "      If logprobs is True, token log probabilities are computed for each generated token.\n",
    "\n",
    "    \"\"\"\n",
    "    device = self.args.device\n",
    "    params = self.model.params\n",
    "    bsz = len(prompt_tokens)\n",
    "    assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)\n",
    "\n",
    "    min_prompt_len = min(len(t) for t in prompt_tokens)\n",
    "    max_prompt_len = max(len(t) for t in prompt_tokens)\n",
    "    assert max_prompt_len <= params.max_seq_len\n",
    "    total_len = min(params.max_seq_len, max_gen_len + max_prompt_len)\n",
    "\n",
    "    pad_id = self.tokenizer.pad_id\n",
    "    tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device=device)\n",
    "    for k, t in enumerate(prompt_tokens):\n",
    "      tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long, device=device)\n",
    "    if logprobs:\n",
    "      token_logprobs = torch.zeros_like(tokens, dtype=torch.float)\n",
    "\n",
    "    prev_pos = 0\n",
    "    eos_reached = torch.tensor([False] * bsz, device=device)\n",
    "    input_text_mask = tokens != pad_id\n",
    "    if min_prompt_len == total_len:\n",
    "      logits, loss = self.model.forward(tokens, prev_pos)\n",
    "      self.debug('got logits.shape=$lg loss $LO temp=$temp topP=$topk', list(logits.shape), loss, temperature, top_p)\n",
    "      token_logprobs = -F.cross_entropy(\n",
    "        input=logits.transpose(1, 2),\n",
    "        target=tokens,\n",
    "        reduction=\"none\",\n",
    "        ignore_index=pad_id,\n",
    "      )\n",
    "\n",
    "    stop_tokens = torch.tensor(list(self.tokenizer.stop_tokens))\n",
    "\n",
    "    for cur_pos in range(min_prompt_len, total_len):\n",
    "      logits,loss = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)\n",
    "      # self.debug('pos $p got logits.shape=$L LOSS $L2', cur_pos, list(logits.shape), loss)\n",
    "      if temperature > 0:\n",
    "        probs = torch.softmax(logits[:, -1] / temperature, dim=-1)\n",
    "        next_token = self._sample_top_p(probs, top_p)\n",
    "      else:\n",
    "        next_token = torch.argmax(logits[:, -1], dim=-1)\n",
    "\n",
    "      next_token = next_token.reshape(-1)\n",
    "      # only replace token if prompt has already been generated\n",
    "      next_token = torch.where(\n",
    "        input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token\n",
    "      )\n",
    "      tokens[:, cur_pos] = next_token\n",
    "      if logprobs:\n",
    "        token_logprobs[:, prev_pos + 1 : cur_pos + 1] = -F.cross_entropy(\n",
    "          input=logits.transpose(1, 2),\n",
    "          target=tokens[:, prev_pos + 1 : cur_pos + 1],\n",
    "          reduction=\"none\",\n",
    "          ignore_index=pad_id,\n",
    "        )\n",
    "      eos_reached |= (~input_text_mask[:, cur_pos]) & (\n",
    "        torch.isin(next_token, stop_tokens)\n",
    "      )\n",
    "      prev_pos = cur_pos\n",
    "      if all(eos_reached):\n",
    "        break\n",
    "    self.debug('got total [$mv, $p]=$t', min_prompt_len, total_len, total_len-min_prompt_len)\n",
    "\n",
    "    if logprobs:\n",
    "      token_logprobs = token_logprobs.tolist()\n",
    "    out_tokens, out_logprobs = [], []\n",
    "    for i, toks in enumerate(tokens.tolist()):\n",
    "      # cut to max gen len\n",
    "      start = 0 if echo else len(prompt_tokens[i])\n",
    "      toks = toks[start : len(prompt_tokens[i]) + max_gen_len]\n",
    "      probs = None\n",
    "      if logprobs:\n",
    "        probs = token_logprobs[i][start : len(prompt_tokens[i]) + max_gen_len]\n",
    "      # cut to after eos tok if any\n",
    "      for stop_token in self.tokenizer.stop_tokens:\n",
    "        try:\n",
    "          eos_idx = toks.index(stop_token)\n",
    "          toks = toks[:eos_idx]\n",
    "          probs = probs[:eos_idx] if logprobs else None\n",
    "        except ValueError:\n",
    "          pass\n",
    "      out_tokens.append(toks)\n",
    "      out_logprobs.append(probs)\n",
    "    return (out_tokens, out_logprobs if logprobs else None)\n",
    "  \n",
    "  def _sample_top_p(self, probs, p):\n",
    "    # (B, vocab_size)\n",
    "    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)\n",
    "    # (B, vocab_size)\n",
    "    probs_sum = torch.cumsum(probs_sort, dim=-1)\n",
    "    # (B, vocab_size)\n",
    "    # (Substracting \"probs_sort\" shifts the cumulative sum by 1 position to the right before masking)\n",
    "    mask = probs_sum - probs_sort > p \n",
    "    # Zero out all the probabilities of tokens that are not selected by the Top P\n",
    "    probs_sort[mask] = 0.0 \n",
    "    # Redistribute the probabilities so that they sum up to 1.\n",
    "    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))\n",
    "    # Sample a token (its index) from the top p distribution\n",
    "    next_token = torch.multinomial(probs_sort, num_samples=1)\n",
    "    # Get the token position in the vocabulary corresponding to the sampled index\n",
    "    next_token = torch.gather(probs_idx, -1, next_token) \n",
    "    return next_token\n",
    "\n",
    "  def text_completion(\n",
    "    self,\n",
    "    prompts: List[str],\n",
    "    temperature: float = 0.6,\n",
    "    top_p: float = 0.9,\n",
    "    max_gen_len: Optional[int] = None,\n",
    "    logprobs: bool = False,\n",
    "    echo: bool = False,\n",
    "  ) -> List[CompletionPrediction]:\n",
    "    \"\"\"\n",
    "    Perform text completion for a list of prompts using the language generation model.\n",
    "\n",
    "    Args:\n",
    "      prompts (List[str]): List of text prompts for completion.\n",
    "      temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6.\n",
    "      top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to 0.9.\n",
    "      max_gen_len (Optional[int], optional): Maximum length of the generated completion sequence.\n",
    "        If not provided, it's set to the model's maximum sequence length minus 1.\n",
    "      logprobs (bool, optional): Flag indicating whether to compute token log probabilities. Defaults to False.\n",
    "      echo (bool, optional): Flag indicating whether to include prompt tokens in the generated output. Defaults to False.\n",
    "\n",
    "    Returns:\n",
    "      List[CompletionPrediction]: List of completion predictions, each containing the generated text completion.\n",
    "\n",
    "    Note:\n",
    "      This method generates text completions for the provided prompts, employing nucleus sampling to introduce controlled randomness.\n",
    "      If logprobs is True, token log probabilities are computed for each generated token.\n",
    "\n",
    "    \"\"\"\n",
    "    if max_gen_len is None:\n",
    "      max_gen_len = self.model.params.max_seq_len - 1\n",
    "    prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]\n",
    "    generation_tokens, generation_logprobs = self.generate(\n",
    "      prompt_tokens=prompt_tokens,\n",
    "      max_gen_len=max_gen_len,\n",
    "      temperature=temperature,\n",
    "      top_p=top_p,\n",
    "      logprobs=logprobs,\n",
    "      echo=echo,\n",
    "    )\n",
    "    if logprobs:\n",
    "      return [\n",
    "        {\n",
    "          \"generation\": self.tokenizer.decode(t),\n",
    "          \"tokens\": [self.tokenizer.decode([x]) for x in t],\n",
    "          \"logprobs\": logprobs_i,\n",
    "        }\n",
    "        for t, logprobs_i in zip(generation_tokens, generation_logprobs)\n",
    "      ]\n",
    "    return [{\"generation\": self.tokenizer.decode(t)} for t in generation_tokens]\n",
    "\n",
    "  def chat_completion(\n",
    "    self,\n",
    "    dialogs: List[Dialog],\n",
    "    temperature: float = 0.6,\n",
    "    top_p: float = 0.9,\n",
    "    max_gen_len: Optional[int] = None,\n",
    "    logprobs: bool = False,\n",
    "  ) -> List[ChatPrediction]:\n",
    "    \"\"\"\n",
    "    Generate assistant responses for a list of conversational dialogs using the language generation model.\n",
    "\n",
    "    Args:\n",
    "      dialogs (List[Dialog]): List of conversational dialogs, where each dialog is a list of messages.\n",
    "      temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6.\n",
    "      top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to 0.9.\n",
    "      max_gen_len (Optional[int], optional): Maximum length of the generated response sequence.\n",
    "        If not provided, it's set to the model's maximum sequence length minus 1.\n",
    "      logprobs (bool, optional): Flag indicating whether to compute token log probabilities. Defaults to False.\n",
    "\n",
    "    Returns:\n",
    "      List[ChatPrediction]: List of chat predictions, each containing the assistant's generated response.\n",
    "\n",
    "    Note:\n",
    "      This method generates assistant responses for the provided conversational dialogs.\n",
    "      It employs nucleus sampling to introduce controlled randomness in text generation.\n",
    "      If logprobs is True, token log probabilities are computed for each generated token.\n",
    "    \"\"\"\n",
    "    if max_gen_len is None:\n",
    "      max_gen_len = self.model.params.max_seq_len - 1\n",
    "\n",
    "    prompt_tokens = [\n",
    "      self.formatter.encode_dialog_prompt(dialog) for dialog in dialogs\n",
    "    ]\n",
    "    self.debug('got $L tokens $M', len(prompt_tokens[0]), \" \".join(map(str, prompt_tokens[0])))\n",
    "    generation_tokens, generation_logprobs = self.generate(\n",
    "      prompt_tokens=prompt_tokens,\n",
    "      max_gen_len=max_gen_len,\n",
    "      temperature=temperature,\n",
    "      top_p=top_p,\n",
    "      logprobs=logprobs,\n",
    "    )\n",
    "    if logprobs:\n",
    "      return [\n",
    "        {\n",
    "          \"generation\": {\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": self.tokenizer.decode(t),\n",
    "          },\n",
    "          \"tokens\": [self.tokenizer.decode([x]) for x in t],\n",
    "          \"logprobs\": logprobs_i,\n",
    "        }\n",
    "        for t, logprobs_i in zip(generation_tokens, generation_logprobs)\n",
    "      ]\n",
    "    return [\n",
    "      {\n",
    "        \"generation\": {\n",
    "          \"role\": \"assistant\",\n",
    "          \"content\": self.tokenizer.decode(t),\n",
    "        },\n",
    "      }\n",
    "      for t in generation_tokens\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d734bd53-e501-4861-975a-ba5657cf5358",
   "metadata": {},
   "outputs": [],
   "source": [
    "## what the heck is this?!?! This is to modify any classes you may have modified above\n",
    "## but skip the reloading of the model!\n",
    "try:\n",
    "  if hasattr(model, 'model'):\n",
    "    print(\"rebuilding model from existing (because we changed the code!)\")\n",
    "    M2 = SimpleLlama3(model.model, model.tokenizer, model.args)\n",
    "    model = M2\n",
    "    print(\"rebuilt model\")\n",
    "except Exception as e:\n",
    "  print('building model from scratch')\n",
    "  T0 = time.monotonic()\n",
    "  model = SimpleLlama3.build(checkpoints_dir='Meta-Llama-3-8B-Instruct', \n",
    "                       tokenizer_path='Meta-Llama-3-8B-Instruct/tokenizer.model',\n",
    "                       load_model=True,\n",
    "                       max_seq_len=1024,\n",
    "                       max_batch_size=8,\n",
    "                       device='cuda')\n",
    "  T1 = time.monotonic()\n",
    "  print(f'{T1-T0} seconds to load model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86459d83-df51-48a1-89fa-59a6bf36076d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model.model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c8e7cf6-deb2-4afc-92e2-73f76fccbc4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## SimpleLlama3 code expects dialogs of this shape:\n",
    "dialogs: List[Dialog] = [\n",
    "    [{\"role\": \"user\", \"content\": \"what is the recipe of mayonnaise?\"}],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8be714a9-aac3-444b-b1c4-881ef70dfa4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def addConversation(M,role=1):\n",
    "  Roles = [ \"system\", \"user\", \"assistant\" ]\n",
    "  Rtn = []\n",
    "  Rtn.append({ \"role\" : Roles[role], \"content\": M})\n",
    "  return Rtn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b536438-20b5-4f65-84c5-47174c0e4a19",
   "metadata": {},
   "outputs": [],
   "source": [
    "Q =  addConversation('what is the recipe for vanilla rice pudding?')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e7efa37-d4c3-4898-969d-2e297cb65211",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = model.chat_completion([Q])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebe4ed8e-b866-4acd-8f47-821fe43a6d0f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(A[0]['generation']['content'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1243dd7-77e5-43e0-9fe1-4a11f0805837",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = model.chat_completion([addConversation(\"Write a poem about trees in winter. Make it like Shakespeare\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6457633d-b949-411a-a90c-562b08ce480a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(A[0]['generation']['content'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d028d5c7-163a-4f42-a8d8-d87812ef945b",
   "metadata": {},
   "outputs": [],
   "source": [
    "C = [ addConversation(\"Write a poem about trees in winter. Make it like Shakespeare\"),\n",
    "       addConversation(A[0]['generation']['content'], role=2),\n",
    "       addConversation(\"It's too long. Make it 5 lines, please\")  ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b80433aa-07da-41fc-ae21-a0913ceecf5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "A2 = model.chat_completion(C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca780d00-e8a5-4df1-8ef5-d0ba1012abf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(A2[0]['generation']['content'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad3ccd1-826a-4a1d-8a9d-3f4e4c614063",
   "metadata": {},
   "outputs": [],
   "source": [
    "A2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65c2df6a-cd13-4d1c-b8d4-e70f481b22b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(A2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0c46004-22ef-41df-8027-29d8cb4821d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "A2[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ef9af90-8f67-4d42-8eeb-06caeb0c91a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def addConversation(M,role=1):\n",
    "  Roles = [ \"system\", \"user\", \"assistant\" ]\n",
    "  return { \"role\" : Roles[role], \"content\": M}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0a2f23a-0d90-4b72-876f-737c547c3ac7",
   "metadata": {},
   "outputs": [],
   "source": [
    "C2 = [[ addConversation(\"Write a poem about trees in winter. Make it like Shakespeare\"),\n",
    "       addConversation(A[0]['generation']['content'], role=2),\n",
    "       addConversation(\"It's too long. Make it 5 lines, please\")  ]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41e9b7d6-d11a-402d-989f-74c1a681e3a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "A2 = model.chat_completion(C2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9416bf21-0163-444c-9de9-7d41e799defd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(A2))\n",
    "print(A2[0]['generation']['content'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3867f965-e597-49b7-8bb3-ebce9df3355c",
   "metadata": {},
   "source": [
    "## How to support LLAMA 3.2!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28e86696-bba2-4730-8ca0-87a252c982ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# llama 3.2!\n",
    "def apply_scaling(freqs: torch.Tensor):\n",
    "  # Values obtained from grid search\n",
    "  scale_factor = 8\n",
    "  low_freq_factor = 1\n",
    "  high_freq_factor = 4\n",
    "  old_context_len = 8192  # original llama3 length\n",
    "\n",
    "  low_freq_wavelen = old_context_len / low_freq_factor\n",
    "  high_freq_wavelen = old_context_len / high_freq_factor\n",
    "  new_freqs = []\n",
    "  for freq in freqs:\n",
    "    wavelen = 2 * math.pi / freq\n",
    "    if wavelen < high_freq_wavelen:\n",
    "      new_freqs.append(freq)\n",
    "    elif wavelen > low_freq_wavelen:\n",
    "      new_freqs.append(freq / scale_factor)\n",
    "    else:\n",
    "      assert low_freq_wavelen != high_freq_wavelen\n",
    "      smooth = (old_context_len / wavelen - low_freq_factor) / (\n",
    "        high_freq_factor - low_freq_factor\n",
    "      )\n",
    "      new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq)\n",
    "  return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device)\n",
    "\n",
    "\n",
    "def precompute_theta_pos_frequencies(head_dim: int, seq_len: int, theta: float = 10000.0, use_scaled: bool = False, device: str = ''):\n",
    "  # As written in the paragraph 3.2.2 of the paper\n",
    "  # >> In order to generalize our results in 2D to any xi ∈ Rd where **d is even**, [...]\n",
    "  assert head_dim % 2 == 0, \"Dimension must be divisible by 2\"\n",
    "  # Build the theta parameter\n",
    "  # According to the formula theta_i = 10000^(-2(i-1)/dim) for i = [1, 2, ... dim/2]\n",
    "  # Shape: (Head_Dim / 2)\n",
    "  theta_numerator = torch.arange(0, head_dim, 2).float()\n",
    "  # Shape: (Head_Dim / 2)\n",
    "  theta = 1.0 / (theta ** (theta_numerator / head_dim)).to(device) # (Dim / 2)\n",
    "  # llama-3.2\n",
    "  # only two lines to get 128K context. WTF!\n",
    "  if use_scaled:\n",
    "    theta = apply_scaling(theta)\n",
    "  # Construct the positions (the \"m\" parameter)\n",
    "  # Shape: (Seq_Len)\n",
    "  m = torch.arange(seq_len, device=device)\n",
    "  # Multiply each theta by each position using the outer product.\n",
    "  # Shape: (Seq_Len) outer_product* (Head_Dim / 2) -> (Seq_Len, Head_Dim / 2)\n",
    "  freqs = torch.outer(m, theta).float()\n",
    "  # We can compute complex numbers in the polar form c = R * exp(m * theta), where R = 1 as follows:\n",
    "  # (Seq_Len, Head_Dim / 2) -> (Seq_Len, Head_Dim / 2)\n",
    "  freqs_complex = torch.polar(torch.ones_like(freqs), freqs)\n",
    "  return freqs_complex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f77779d-2e16-4146-af5d-9fb1452f5728",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelArgs:\n",
    "  dim: int = 4096\n",
    "  n_layers: int = 32\n",
    "  n_heads: int = 32\n",
    "  n_kv_heads: Optional[int] = None\n",
    "  vocab_size: int = -1 # Later set in the build method\n",
    "  multiple_of: int = 256 # make SwiGLU hidden layer size multiple of large power of 2\n",
    "  ffn_dim_multiplier: Optional[float] = None\n",
    "  norm_eps: float = 1e-5\n",
    "  rope_theta: float = 500000\n",
    "\n",
    "  ## llama 3.2\n",
    "  use_scaled_rope: bool = False\n",
    "\n",
    "  # Needed for KV cache\n",
    "  max_batch_size: int = 32\n",
    "  max_seq_len: int = 2048\n",
    "\n",
    "  # vision model params\n",
    "  vision_chunk_size: int = -1  # image resolution for image models\n",
    "  vision_max_num_chunks: int = 4\n",
    "  vision_num_cross_attention_layers: int = -1\n",
    "\n",
    "  # others!\n",
    "  # doRPE : bool = True\n",
    "  bias: bool = False\n",
    "  dropout: float = 0.0\n",
    "\n",
    "  # for SimpleTrainer and TrainerConfig\n",
    "  dtype : torch.dtype = None\n",
    "  device: str = None\n",
    "\n",
    "  def __init__(self, **kwargs):\n",
    "    for k, v in kwargs.items():\n",
    "      if hasattr(self, k):\n",
    "        setattr(self, k, v)\n",
    "\n",
    "  def __post_init__(s):\n",
    "    if isinstance(s.dtype, str):\n",
    "      s.dtype = findByName(s.dtype, locals(), globals())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dceddd84-435f-4d2c-9a0e-7cfdb048410a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "  model = SimpleLlama3.build(checkpoints_dir='Llama3.2-3B-Instruct', \n",
    "                       tokenizer_path='Llama3.2-3B-Instruct/tokenizer.model',\n",
    "                       load_model=True,\n",
    "                       max_seq_len=1024,\n",
    "                       max_batch_size=8,\n",
    "                       device='cuda')\n",
    "  T1 = time.monotonic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "574775ea-d208-4ae0-9060-93a5b34b82b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38fd778-4dd8-4f82-8948-84ddcd055aac",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "del model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78bfd1bb-2769-415d-81b6-ad15f08e035e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3790902-690e-409e-b5ca-a5e5ff807075",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b59b6a59-fa61-4739-b62a-9bd8991d6f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87ca097f-a280-499f-a41e-42e435ecb41f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.model.args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7695a52f-5fc4-4a26-9877-fef645c353d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "Q =[[addConversation(\"What is the recipe for vanilla rice pudding?\")]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaf568ff-3174-4794-bd2d-5ebe64289dea",
   "metadata": {},
   "outputs": [],
   "source": [
    "addConversation(\"Waht is the recipe for vanilla rice pudding?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4073f44-2cb6-4779-ba66-dbd50c7ba41c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "A3 = model.chat_completion(Q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4000f444-d6e4-4df0-892b-a073b35676a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(A3[0]['generation']['content'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89ee3a8c-4da6-4aca-88dd-5fe3633de390",
   "metadata": {},
   "outputs": [],
   "source": [
    "C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "484c0716-b573-4fe2-b9f0-e5c371bbc46e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "A4 = model.chat_completion(C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d695c464-1e69-4a1a-aa7d-474f49c07e8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "C2 = [[ addConversation(\"Write a poem about trees in winter. Make it like Shakespeare\"),\n",
    "       addConversation(A[0]['generation']['content'], role=2),\n",
    "       addConversation(\"It's too long. Make it 5 lines, please\")  ]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6f10ec2-5482-430a-8bc6-bff045bbf0bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "A5 = model.chat_completion(C2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4c1137b-25e5-4218-82e4-dc48d80123c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(A5[0]['generation']['content'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}