multipacking.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "486f817a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip3.10 install -e . --no-deps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "16fce098",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "57327cab",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, T5ForConditionalGeneration\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('mesolitica/nanot5-small-malaysian-cased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0965c7a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = T5ForConditionalGeneration.from_pretrained('mesolitica/nanot5-small-malaysian-cased',\n",
    "                                                  attn_implementation = 'sdpa').to(torch.bfloat16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "224aab6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):\n",
    "    \"\"\"\n",
    "    Adapted from Mesh Tensorflow:\n",
    "    https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593\n",
    "\n",
    "    Translate relative position to a bucket number for relative attention. The relative position is defined as\n",
    "    memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to\n",
    "    position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for\n",
    "    small absolute relative_position and larger buckets for larger absolute relative_positions. All relative\n",
    "    positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.\n",
    "    This should allow for more graceful generalization to longer sequences than the model has been trained on\n",
    "\n",
    "    Args:\n",
    "        relative_position: an int32 Tensor\n",
    "        bidirectional: a boolean - whether the attention is bidirectional\n",
    "        num_buckets: an integer\n",
    "        max_distance: an integer\n",
    "\n",
    "    Returns:\n",
    "        a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)\n",
    "    \"\"\"\n",
    "    relative_buckets = 0\n",
    "    if bidirectional:\n",
    "        num_buckets //= 2\n",
    "        relative_buckets += (relative_position > 0).to(torch.long) * num_buckets\n",
    "        relative_position = torch.abs(relative_position)\n",
    "    else:\n",
    "        relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))\n",
    "    # now relative_position is in the range [0, inf)\n",
    "\n",
    "    # half of the buckets are for exact increments in positions\n",
    "    max_exact = num_buckets // 2\n",
    "    is_small = relative_position < max_exact\n",
    "\n",
    "    # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance\n",
    "    relative_position_if_large = max_exact + (\n",
    "        torch.log(relative_position.float() / max_exact)\n",
    "        / math.log(max_distance / max_exact)\n",
    "        * (num_buckets - max_exact)\n",
    "    ).to(torch.long)\n",
    "    relative_position_if_large = torch.min(\n",
    "        relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)\n",
    "    )\n",
    "\n",
    "    relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)\n",
    "    return relative_buckets\n",
    "\n",
    "def compute_bias(\n",
    "    query_length, \n",
    "    key_length,\n",
    "    relative_attention_bias,\n",
    "    bidirectional = True, \n",
    "    num_buckets = 32, \n",
    "    max_distance = 128, \n",
    "    device=None,\n",
    "):\n",
    "    \"\"\"Compute binned relative position bias\"\"\"\n",
    "    if device is None:\n",
    "        device = relative_attention_bias.weight.device\n",
    "    context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]\n",
    "    memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]\n",
    "    relative_position = memory_position - context_position  # shape (query_length, key_length)\n",
    "    relative_position_bucket = _relative_position_bucket(\n",
    "        relative_position,  # shape (query_length, key_length)\n",
    "        bidirectional=bidirectional,\n",
    "        num_buckets=num_buckets,\n",
    "        max_distance=max_distance,\n",
    "    )\n",
    "    values = relative_attention_bias(relative_position_bucket)  # shape (query_length, key_length, num_heads)\n",
    "    values = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, query_length, key_length)\n",
    "    return values\n",
    "\n",
    "def block_diagonal_concat_inverted(*masks, dtype=torch.bfloat16):\n",
    "    total_size = sum(mask.size(0) for mask in masks)\n",
    "    combined_mask = torch.zeros(total_size, total_size, dtype=dtype)\n",
    "\n",
    "    current_pos = 0\n",
    "\n",
    "    for mask in masks:\n",
    "        size = mask.size(0)\n",
    "        combined_mask[current_pos:current_pos + size, current_pos:current_pos + size] = mask\n",
    "        current_pos += size\n",
    "\n",
    "    min_value = torch.finfo(dtype).min if dtype.is_floating_point else torch.iinfo(dtype).min\n",
    "    inverted_mask = torch.where(combined_mask == 1, torch.tensor(0, dtype=dtype), min_value)\n",
    "    return inverted_mask.unsqueeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1aed023e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_attention_mask(attention_mask):\n",
    "    maxlen_right = max([attention_mask[i].shape[1] for i in range(len(attention_mask))])\n",
    "    maxlen_bottom = max([attention_mask[i].shape[0] for i in range(len(attention_mask))])\n",
    "    attention_mask = [\n",
    "        F.pad(\n",
    "            attention_mask[i],\n",
    "            (0, maxlen_right - attention_mask[i].shape[1], 0, maxlen_bottom - attention_mask[i].shape[0])) for i in range(\n",
    "            len(attention_mask))]\n",
    "    return torch.stack(attention_mask)\n",
    "\n",
    "def pad_attention_mask_4d(attention_mask):\n",
    "    maxlen_right = max([attention_mask[i].shape[-2] for i in range(len(attention_mask))])\n",
    "    maxlen_bottom = max([attention_mask[i].shape[-1] for i in range(len(attention_mask))])\n",
    "    attention_mask = [\n",
    "        F.pad(\n",
    "            attention_mask[i],\n",
    "            (0, maxlen_right - attention_mask[i].shape[-2], 0, maxlen_bottom - attention_mask[i].shape[-1])) for i in range(\n",
    "            len(attention_mask))]\n",
    "    return torch.stack(attention_mask)\n",
    "\n",
    "def block_diagonal_concat(*masks, dtype=torch.bfloat16):\n",
    "    total_size = sum(mask.size(0) for mask in masks)\n",
    "    combined_mask = torch.zeros(total_size, total_size, dtype=dtype)\n",
    "\n",
    "    current_pos = 0\n",
    "\n",
    "    for mask in masks:\n",
    "        size = mask.size(0)\n",
    "        combined_mask[current_pos:current_pos + size, current_pos:current_pos + size] = mask\n",
    "        current_pos += size\n",
    "\n",
    "    return combined_mask\n",
    "\n",
    "def block_diagonal_concat_4d(*masks, dtype=torch.bfloat16):\n",
    "    total_size = sum(mask.size(1) for mask in masks)\n",
    "    combined_mask = torch.zeros(masks[0].shape[0], \n",
    "                                total_size, total_size, dtype=dtype)\n",
    "\n",
    "    current_pos = 0\n",
    "\n",
    "    for mask in masks:\n",
    "        size = mask.size(1)\n",
    "        combined_mask[:, current_pos:current_pos + size, current_pos:current_pos + size] = mask\n",
    "        current_pos += size\n",
    "\n",
    "    return combined_mask\n",
    "\n",
    "def block_diagonal_concat_cross(*masks, dtype=torch.bfloat16):\n",
    "    total_rows = sum(mask.size(0) for mask in masks)\n",
    "    total_cols = sum(mask.size(1) for mask in masks)\n",
    "    \n",
    "    combined_mask = torch.zeros((total_rows, total_cols), dtype=dtype)\n",
    "    \n",
    "    current_row, current_col = 0, 0\n",
    "\n",
    "    for mask in masks:\n",
    "        rows, cols = mask.size()\n",
    "        combined_mask[current_row:current_row + rows, current_col:current_col + cols] = mask\n",
    "        current_row += rows\n",
    "        current_col += cols\n",
    "        \n",
    "    return combined_mask\n",
    "\n",
    "def input_ids(left, right):\n",
    "    inputs = tokenizer(left + tokenizer.eos_token)\n",
    "    inputs['labels'] = tokenizer(right + tokenizer.eos_token)['input_ids']\n",
    "    inputs.pop('token_type_ids')\n",
    "    \n",
    "    left_len = len(inputs['input_ids'])\n",
    "    right_len = len(inputs['labels'])\n",
    "    \n",
    "    inputs['attention_mask'] = torch.ones(left_len, left_len)\n",
    "    inputs['encoder_attention_mask'] = torch.ones(right_len, left_len)\n",
    "    inputs['decoder_attention_mask'] = torch.tril(torch.ones(right_len, right_len))\n",
    "    return inputs\n",
    "    \n",
    "def multipack(batch):\n",
    "    batch = [b for b in batch if b is not None]\n",
    "    \n",
    "    results = {}\n",
    "    \n",
    "    input_ids, labels = [], []\n",
    "    encoder_biases, decoder_biases = [], []\n",
    "    for b in batch:\n",
    "        input_ids.extend(b['input_ids'])\n",
    "        labels.extend(b['labels'])\n",
    "        \n",
    "        left = len(b['input_ids'])\n",
    "        right = len(b['labels'])\n",
    "        \n",
    "        encoder_bias = compute_bias(\n",
    "            left, left,\n",
    "            model.encoder.block[0].layer[0].SelfAttention.relative_attention_bias,\n",
    "            bidirectional=True,\n",
    "            num_buckets=model.config.relative_attention_num_buckets,\n",
    "            max_distance=model.config.relative_attention_max_distance,\n",
    "        )\n",
    "        encoder_biases.append(encoder_bias[0])\n",
    "        \n",
    "        decoder_bias = compute_bias(\n",
    "            right, right,\n",
    "            model.decoder.block[0].layer[0].SelfAttention.relative_attention_bias,\n",
    "            bidirectional=False,\n",
    "            num_buckets=model.config.relative_attention_num_buckets,\n",
    "            max_distance=model.config.relative_attention_max_distance,\n",
    "        )\n",
    "        decoder_biases.append(decoder_bias[0])\n",
    "\n",
    "    results['input_ids'] = input_ids\n",
    "    results['labels'] = labels\n",
    "    \n",
    "    results['position_bias'] = block_diagonal_concat_4d(*encoder_biases)\n",
    "    results['decoder_position_bias'] = block_diagonal_concat_4d(*decoder_biases)\n",
    "    \n",
    "    attention_mask = [b['attention_mask'] for b in batch]\n",
    "    results['attention_mask'] = block_diagonal_concat(*attention_mask)\n",
    "    encoder_attention_mask = [b['encoder_attention_mask'] for b in batch]\n",
    "    results['encoder_attention_mask'] = block_diagonal_concat_cross(*encoder_attention_mask)\n",
    "    decoder_attention_mask = [b['decoder_attention_mask'] for b in batch]\n",
    "    results['decoder_attention_mask'] = block_diagonal_concat(*decoder_attention_mask)\n",
    "    \n",
    "    return results\n",
    "\n",
    "def collator(batch, pad_token_id = 1, label_pad = -100):\n",
    "    max_length = max(len(l['input_ids']) for l in batch)\n",
    "    results = {}\n",
    "    results['input_ids'] = [\n",
    "        b['input_ids'] + [pad_token_id] * (max_length - len(b['input_ids']))\n",
    "        for b in batch\n",
    "    ]\n",
    "    results['input_ids'] = torch.tensor(results['input_ids'], dtype = torch.int64)\n",
    "    \n",
    "    max_length = max(len(l['labels']) for l in batch)\n",
    "    results['labels'] = [\n",
    "        b['labels'] + [label_pad] * (max_length - len(b['labels']))\n",
    "        for b in batch\n",
    "    ]\n",
    "    results['labels'] = torch.tensor(results['labels'], dtype = torch.int64)\n",
    "    \n",
    "    results['position_bias'] = pad_attention_mask_4d([b['position_bias'] for b in batch])\n",
    "    results['decoder_position_bias'] = pad_attention_mask_4d([b['decoder_position_bias'] for b in batch])\n",
    "    \n",
    "    attention_mask = [b['attention_mask'] for b in batch]\n",
    "    results['attention_mask'] = pad_attention_mask(attention_mask)\n",
    "    encoder_attention_mask = [b['encoder_attention_mask'] for b in batch]\n",
    "    results['encoder_attention_mask'] = pad_attention_mask(encoder_attention_mask)\n",
    "    decoder_attention_mask = [b['decoder_attention_mask'] for b in batch]\n",
    "    results['decoder_attention_mask'] = pad_attention_mask(decoder_attention_mask)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7b3f4208",
   "metadata": {},
   "outputs": [],
   "source": [
    "left = ['this life', 'i want to eat', 'yeah i know kind of']\n",
    "right = ['so tough', 'nasi lemak and chicken rendang', 'i guess?']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e5c6cab6",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = input_ids(left[0], right[0])\n",
    "b1 = input_ids(left[1], right[1])\n",
    "b2 = input_ids(left[2], right[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4cfe5328",
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = multipack([b, b1, b2])\n",
    "m2 = multipack([b, b2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "25faebbf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.]],\n",
       "       dtype=torch.bfloat16)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1['attention_mask']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fd68729f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1.]],\n",
       "       dtype=torch.bfloat16)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1['encoder_attention_mask']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0a111470",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1.]],\n",
       "       dtype=torch.bfloat16)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1['decoder_attention_mask']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "5ed67900",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 4.0625,  7.3750,  5.6250,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [-0.1455,  4.0625,  7.3750,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [-1.6484, -0.1455,  4.0625,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  4.0625,  7.3750,  5.6250,  4.6250,  4.0938,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000, -0.1455,  4.0625,  7.3750,  5.6250,  4.6250,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000, -1.6484, -0.1455,  4.0625,  7.3750,  5.6250,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.1089, -1.6484, -0.1455,  4.0625,  7.3750,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000, -1.0156,  0.1089, -1.6484, -0.1455,  4.0625,\n",
       "          0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          4.0625,  7.3750,  5.6250,  4.6250,  4.0938,  3.6719,  3.3125],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "         -0.1455,  4.0625,  7.3750,  5.6250,  4.6250,  4.0938,  3.6719],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "         -1.6484, -0.1455,  4.0625,  7.3750,  5.6250,  4.6250,  4.0938],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "          0.1089, -1.6484, -0.1455,  4.0625,  7.3750,  5.6250,  4.6250],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "         -1.0156,  0.1089, -1.6484, -0.1455,  4.0625,  7.3750,  5.6250],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "         -0.1738, -1.0156,  0.1089, -1.6484, -0.1455,  4.0625,  7.3750],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "         -1.0391, -0.1738, -1.0156,  0.1089, -1.6484, -0.1455,  4.0625]],\n",
       "       dtype=torch.bfloat16, grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1['position_bias'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "4c6e8b3a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.7812, 2.7812, 2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [3.9375, 2.7812, 2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [2.4062, 3.9375, 2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [1.5312, 2.4062, 3.9375, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 2.7812, 2.7812, 2.7812, 2.7812, 2.7812,\n",
       "         2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 3.9375, 2.7812, 2.7812, 2.7812, 2.7812,\n",
       "         2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 2.4062, 3.9375, 2.7812, 2.7812, 2.7812,\n",
       "         2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 1.5312, 2.4062, 3.9375, 2.7812, 2.7812,\n",
       "         2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 1.0469, 1.5312, 2.4062, 3.9375, 2.7812,\n",
       "         2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 0.4824, 1.0469, 1.5312, 2.4062, 3.9375,\n",
       "         2.7812, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0427, 0.4824, 1.0469, 1.5312, 2.4062,\n",
       "         3.9375, 2.7812, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000, 0.0000, 2.7812, 2.7812, 2.7812, 2.7812],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000, 0.0000, 3.9375, 2.7812, 2.7812, 2.7812],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000, 0.0000, 2.4062, 3.9375, 2.7812, 2.7812],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "         0.0000, 0.0000, 1.5312, 2.4062, 3.9375, 2.7812]],\n",
       "       dtype=torch.bfloat16, grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1['decoder_position_bias'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ab7c30ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = collator([m1, m2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5502fd34",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([8, 15, 15]), torch.Size([8, 10, 10]))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1['position_bias'].shape, m2['position_bias'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "908b002b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 15, 32128])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o.logits.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3.10",
   "language": "python",
   "name": "python3.10"
  },
  "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
}