gemma_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers\n",
    "!pip install datasets\n",
    "!pip install matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BOzDPUPKjC0W"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as opt\n",
    "import math\n",
    "from transformers import AutoTokenizer\n",
    "from torch.cuda import amp\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import random\n",
    "\n",
    "torch.manual_seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "max_token = 64\n",
    "device_id = 1 # GPU device id.\n",
    "epochs = 31\n",
    "batch_size = 600\n",
    "validation_batch_size = 10\n",
    "weight_decay = 1e-3\n",
    "lr = 1e-3\n",
    "num_layer = 3\n",
    "head_dim = 64\n",
    "projection_dim = 1024\n",
    "expansion_factor = 4\n",
    "checkpoint_filepath = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load llama2 token embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load from pt file (If you had already preprocessed)\n",
    "word_embeddings_tensor = torch.load('word_embeddings_tensor_llama2.pt').cuda(device_id)\n",
    "num_embeddings, embedding_dim = word_embeddings_tensor.shape\n",
    "word_embeddings_tensor.requires_grad = False\n",
    "\n",
    "model_id = \"NousResearch/Llama-2-7b-hf\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side=\"right\")\n",
    "pad_token_id = 32000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_unique_random_numbers(N, k):\n",
    "    nums = random.sample(range(N), k)\n",
    "    nums.sort()\n",
    "    return torch.tensor(nums)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load training data\n",
    "training_data = torch.load(\"llama2_wiki_64_ranked.pt\")\n",
    "\n",
    "# Take out the evaluation tensor\n",
    "random_rows = generate_unique_random_numbers(training_data.shape[0], int(training_data.shape[0] * 0.02)) # Using 2% as evaluation tensor\n",
    "validation_data = training_data[random_rows]\n",
    "\n",
    "# Truncating training data\n",
    "mask = torch.ones(training_data.size(0), dtype=torch.bool)\n",
    "mask[random_rows] = False\n",
    "training_data = training_data[mask]\n",
    "\n",
    "# Throw away the last batch\n",
    "smaller_batch = training_data.shape[0] % batch_size\n",
    "training_data = training_data[:-smaller_batch]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Randomize Trainging data\n",
    "N = training_data.shape[0]  # Total number of rows\n",
    "block_size = batch_size  # Size of each block\n",
    "N_blocks = (N + block_size - 1) // block_size  # Total number of blocks\n",
    "\n",
    "# Create a list of block indices and shuffle them\n",
    "block_indices = list(range(N_blocks))\n",
    "random.shuffle(block_indices)\n",
    "\n",
    "# Generate new row indices based on the shuffled block indices\n",
    "indices = []\n",
    "\n",
    "for block_idx in block_indices:\n",
    "    start = block_idx * block_size\n",
    "    end = min(start + block_size, N)\n",
    "    block_row_indices = torch.arange(start, end)\n",
    "    indices.append(block_row_indices)\n",
    "\n",
    "# Concatenate all indices into a single tensor\n",
    "indices = torch.cat(indices)\n",
    "\n",
    "# Rearrange training data using the new indices\n",
    "training_data = training_data[indices]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a TensorDataset\n",
    "training_data = TensorDataset(training_data)\n",
    "validation_data = TensorDataset(validation_data)\n",
    "\n",
    "# Use DataLoader for batching, etc.\n",
    "training_loader = DataLoader(training_data, batch_size=batch_size)\n",
    "validation_loader = DataLoader(validation_data, batch_size=validation_batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Instantiate LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-ommkUOkjG6r"
   },
   "outputs": [],
   "source": [
    "class ROPEEmbedding(nn.Module):\n",
    "    def __init__(self, max_token: int, dim: int, theta: int):\n",
    "        super().__init__()\n",
    "        self.pos_emb = self.create_embedding(max_token, dim, theta)\n",
    "\n",
    "    def create_embedding(self, max_token: int, dim: int, theta: int) -> torch.Tensor:\n",
    "        tensor = torch.arange(0, dim // 2)\n",
    "        tensor = torch.repeat_interleave(tensor, 2)\n",
    "        tensor = -tensor * 2 / dim\n",
    "        tensor = torch.pow(theta, tensor)\n",
    "\n",
    "        index = torch.arange(max_token).float() # This is the m in the formula\n",
    "        tensor = torch.einsum(\"i, j -> ij\", tensor, index)\n",
    "\n",
    "        cos_matrix = tensor.cos()\n",
    "        sin_matrix = tensor.sin()\n",
    "        sin_matrix[0::2] *= -1 # Flipping sign for 0, 2, 4... row of sin matrix\n",
    "\n",
    "        pos_emb = torch.cat((cos_matrix, sin_matrix), dim=0)\n",
    "        pos_emb = pos_emb.transpose(1, 0)\n",
    "        pos_emb = nn.Parameter(pos_emb, requires_grad=False)\n",
    "\n",
    "        return pos_emb\n",
    "\n",
    "    def flip_for_sin(self, tensor: torch.Tensor) -> torch.Tensor:\n",
    "        original_shape = tensor.shape\n",
    "        tensor = tensor.reshape(tensor.shape[0], tensor.shape[1], -1, 2) # Get to pairs\n",
    "        tensor = tensor[..., [1, 0]] # Swap\n",
    "        tensor = tensor.reshape(original_shape) # Get back to original shape\n",
    "        return tensor\n",
    "\n",
    "    def forward(self, tensor: torch.Tensor) -> torch.Tensor:\n",
    "        sequence_length = tensor.shape[2] # Assuming we are using batch_size, head, sequence_length and dim\n",
    "\n",
    "        tensor = torch.cat((tensor, self.flip_for_sin(tensor)), dim=-1)\n",
    "        tensor = tensor * self.pos_emb[:sequence_length, :]\n",
    "        cos, sin = tensor.chunk(chunks=2, dim=-1)\n",
    "        tensor = cos + sin\n",
    "        return tensor\n",
    "\n",
    "class MultiQueryAttention(nn.Module):\n",
    "    def __init__(self, hidden_dim: int, head_dim: int, q_head: int, kv_head: int, embedding: ROPEEmbedding):\n",
    "        super().__init__()\n",
    "        self.head_dim = head_dim\n",
    "        self.q_head = q_head\n",
    "        self.kv_head = kv_head\n",
    "        self.embedding = embedding\n",
    "        self.qkv = nn.Linear(hidden_dim, (q_head+kv_head*2)*head_dim)\n",
    "        self.o = nn.Linear(q_head*head_dim, hidden_dim)\n",
    "        self.scaler = 1/math.sqrt(head_dim)\n",
    "\n",
    "        if q_head != kv_head:\n",
    "            # If we are using multi query attention\n",
    "            assert q_head % kv_head == 0\n",
    "            self.multi_query_attention = True\n",
    "            self.q_kv_scale = q_head//kv_head\n",
    "        else:\n",
    "            self.multi_query_attention = False\n",
    "\n",
    "    def forward(self, tensor: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:\n",
    "        batch_size, seq_len, hid_dim = tensor.shape\n",
    "\n",
    "        tensor = self.qkv(tensor)\n",
    "        query, key, value = tensor.split([self.head_dim*self.q_head, self.head_dim*self.kv_head, self.head_dim*self.kv_head], dim=-1)\n",
    "\n",
    "        query = query.view(batch_size, seq_len, self.q_head, self.head_dim)\n",
    "        key = key.view(batch_size, seq_len, self.kv_head, self.head_dim)\n",
    "        value = value.view(batch_size, seq_len, self.kv_head, self.head_dim)\n",
    "\n",
    "        if self.multi_query_attention:\n",
    "            # If we are using multi query attention, duplicate key value heads\n",
    "            key = torch.repeat_interleave(key, self.q_kv_scale, dim=-2)\n",
    "            value = torch.repeat_interleave(value, self.q_kv_scale, dim=-2)\n",
    "\n",
    "        # Switch to batch_size, head, seq_len, head_dim\n",
    "        query = query.transpose(1, 2)\n",
    "        key = key.transpose(1, 2)\n",
    "        value = value.transpose(1, 2)\n",
    "\n",
    "        # Apply ROPE\n",
    "        query = self.embedding(query)\n",
    "        key = self.embedding(key)\n",
    "\n",
    "        # Classic self attention\n",
    "        attention_raw = torch.matmul(query, key.transpose(2, 3))\n",
    "        attention_scaled = attention_raw * self.scaler\n",
    "        if attention_mask != None:\n",
    "            attention_scaled += attention_mask\n",
    "        attention_score = torch.softmax(attention_scaled, dim=-1)\n",
    "        value = torch.matmul(attention_score, value)\n",
    "\n",
    "        # Reshape back to batch_size, seq_len, hid_dim\n",
    "        value = value.transpose(1, 2).contiguous()\n",
    "        value = value.view(batch_size, seq_len, hid_dim)\n",
    "\n",
    "        # Output layer\n",
    "        output = self.o(value)\n",
    "\n",
    "        return output\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, hidden_size: int, inner_size: int, dropout_ratio: float = 0.1):\n",
    "        super().__init__()\n",
    "        self.gate_and_up = nn.Linear(hidden_size, inner_size * 2)\n",
    "        self.down = nn.Linear(inner_size, hidden_size)\n",
    "        self.dropout = nn.Dropout(p=dropout_ratio)\n",
    "\n",
    "    def forward(self, tensor: torch.Tensor) -> torch.Tensor:\n",
    "        tensor = self.gate_and_up(tensor)\n",
    "        gate, up = tensor.chunk(chunks=2, dim=-1)\n",
    "        gate = F.gelu(gate, approximate=\"tanh\")\n",
    "        tensor = gate * up\n",
    "        tensor = self.dropout(tensor)\n",
    "        tensor = self.down(tensor)\n",
    "        return tensor\n",
    "\n",
    "class GemmaLayer(nn.Module):\n",
    "    def __init__(self, hidden_dim: int, inner_size: int, head_dim: int, q_head: int, kv_head: int, embedding: ROPEEmbedding, dropout_ratio: float = 0.1):\n",
    "        super().__init__()\n",
    "        self.norm1 = nn.LayerNorm(hidden_dim)\n",
    "        self.mqa = MultiQueryAttention(hidden_dim, head_dim, q_head, kv_head, embedding)\n",
    "\n",
    "        self.norm2 = nn.LayerNorm(hidden_dim)\n",
    "        self.ffn = FeedForward(hidden_dim, inner_size, dropout_ratio)\n",
    "\n",
    "    def forward(self, tensor: torch.Tensor, attention_mask: torch.Tensor = None):\n",
    "        skip_connection = tensor\n",
    "        tensor = self.norm1(tensor)\n",
    "        tensor = self.mqa(tensor, attention_mask)\n",
    "        tensor += skip_connection\n",
    "\n",
    "        skip_connection = tensor\n",
    "        tensor = self.norm2(tensor)\n",
    "        tensor = self.ffn(tensor)\n",
    "        tensor += skip_connection\n",
    "\n",
    "        return tensor\n",
    "\n",
    "class Gemma(nn.Module):\n",
    "    def __init__(self, \n",
    "                 num_layer: int, \n",
    "                 vocab_size: int, \n",
    "                 max_token: int, \n",
    "                 hidden_dim: int, \n",
    "                 inner_size: int, \n",
    "                 head_dim: int, \n",
    "                 q_head: int = None, \n",
    "                 kv_head: int = None, \n",
    "                 dropout_ratio: float = 0.1, \n",
    "                 theta: int = 10000, \n",
    "                 projection_dim: int = None):\n",
    "        super().__init__()\n",
    "        self.embedding = ROPEEmbedding(max_token, head_dim, theta)\n",
    "        self.num_layer = num_layer\n",
    "\n",
    "        # Because of computational power limit, we might want to project input token embedding down.\n",
    "        self.projection = projection_dim != None\n",
    "        if self.projection:\n",
    "            self.projection_matrix = nn.Linear(hidden_dim, projection_dim)\n",
    "            hidden_dim = projection_dim\n",
    "\n",
    "        if q_head == None:\n",
    "            q_head = (hidden_dim // head_dim)\n",
    "\n",
    "        if kv_head == None:\n",
    "            kv_head = (hidden_dim // head_dim)\n",
    "\n",
    "        if hidden_dim % (head_dim * q_head) != 0 or hidden_dim % (head_dim * kv_head):\n",
    "            raise ValueError(\"Error: hidden_dim or projection_dim (if specified) must be divisible by the product of the number of q or kv heads and the head dimension.\")\n",
    "\n",
    "        self.transformer = nn.ModuleList()\n",
    "        for _ in range(self.num_layer):\n",
    "            self.transformer.append(GemmaLayer(hidden_dim, inner_size, head_dim, q_head, kv_head, self.embedding, dropout_ratio))\n",
    "        self.output_norm = nn.LayerNorm(hidden_dim)\n",
    "\n",
    "        self.classifier = nn.Linear(hidden_dim, vocab_size)\n",
    "\n",
    "    def forward(self, tensor: torch.Tensor) -> torch.Tensor:\n",
    "        # If projecting input embeddings\n",
    "        if self.projection:\n",
    "            tensor = self.projection_matrix(tensor)\n",
    "\n",
    "        # Create causal mask\n",
    "        causal_mask = torch.triu(torch.ones(tensor.shape[1], tensor.shape[1]) * float('-inf'), diagonal=1).cuda(device_id)\n",
    "        causal_mask.requires_grad = False\n",
    "        \n",
    "        seq_len = tensor.shape[1]\n",
    "        for layer in self.transformer:\n",
    "            tensor = layer(tensor, causal_mask)\n",
    "\n",
    "        tensor = self.output_norm(tensor)\n",
    "\n",
    "        # Classification\n",
    "        tensor = self.classifier(tensor)\n",
    "        return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_checkpoint(model, optimizer, epoch, loss):\n",
    "    if isinstance(model, nn.DataParallel):\n",
    "        model_to_save = model.module\n",
    "    else:\n",
    "        model_to_save = model\n",
    "    checkpoint = {\n",
    "        'epoch': epoch,\n",
    "        'model_state_dict': model_to_save.state_dict(),\n",
    "        'optimizer_state_dict': optimizer.state_dict(),\n",
    "        'loss': loss,\n",
    "    }\n",
    "    \n",
    "    timestamp = time.strftime('%Y%m%d_%H%M%S')\n",
    "    filename = f'checkpoint_{epoch}_{timestamp}.pth.tar'\n",
    "    torch.save(checkpoint, filename)\n",
    "    print(f'Checkpoint saved at epoch {epoch} as {filename}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_checkpoint(model, optimizer, filename) -> int:\n",
    "    checkpoint = torch.load(filename)\n",
    "    \n",
    "    model.load_state_dict(checkpoint['model_state_dict'])\n",
    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "    \n",
    "    epoch = checkpoint['epoch']\n",
    "    loss = checkpoint['loss']\n",
    "    print(f'Checkpoint loaded from epoch {epoch} with loss {loss}')\n",
    "    return epoch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gemma = Gemma(num_layer, num_embeddings, max_token, embedding_dim, embedding_dim*expansion_factor, head_dim, projection_dim=projection_dim).cuda(device_id)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.AdamW(gemma.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs*len(training_loader), eta_min=1e-5)\n",
    "\n",
    "if checkpoint_filepath != None and checkpoint_filepath != \"\":\n",
    "    current_epoch = load_checkpoint(gemma, optimizer, checkpoint_filepath) + 1\n",
    "else:\n",
    "    current_epoch = 0\n",
    "\n",
    "print(\"This model has\", sum(p.numel() for p in gemma.parameters()), \"parameters.\")\n",
    "scaler = amp.GradScaler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_train = []\n",
    "loss_valid = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chunk data\n",
    "# For example, in a batch, the longest sentence length(including padding token) is 34. Then the training data become shape [batch_size, 34]\n",
    "def trim_padding(input_tensor):\n",
    "    # Create a mask where tokens are not equal to the pad_token\n",
    "    mask = input_tensor != pad_token_id  # Shape: [batch_size, max_seq_length]\n",
    "\n",
    "    # Calculate the lengths of each sentence (number of non-padding tokens)\n",
    "    lengths = mask.sum(dim=1)  # Shape: [batch_size]\n",
    "\n",
    "    # Find the maximum and minimum sentence lengths\n",
    "    max_length = lengths.max().item()\n",
    "    min_length = lengths.min().item()\n",
    "\n",
    "    # Check if the difference between the longest and shortest sentence is 2 or more\n",
    "    if max_length - min_length >= 2:\n",
    "        return None\n",
    "\n",
    "    # Trim the input tensor to the maximum sentence length\n",
    "    trimmed_tensor = input_tensor[:, :max_length]\n",
    "\n",
    "    return trimmed_tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(current_epoch, epochs):\n",
    "    loss_train_epoch = []\n",
    "    loss_val_epoch = []\n",
    "    \n",
    "    gemma.train()\n",
    "    for data in tqdm(training_loader):\n",
    "        # Teacher forcing\n",
    "        data = trim_padding(data[0])\n",
    "        if data == None:\n",
    "            continue\n",
    "            \n",
    "        input_data = data[:, :-1].long().cuda(device_id)\n",
    "        target_data = data[:, 1:].long().cuda(device_id)\n",
    "\n",
    "        # Convert to embedding.\n",
    "        input_embeddings = word_embeddings_tensor[input_data]\n",
    "\n",
    "        # Forward pass\n",
    "        with amp.autocast():\n",
    "            prediction = gemma(input_embeddings)\n",
    "\n",
    "            # Change shape for loss calculation\n",
    "            prediction = prediction.view(-1, num_embeddings)\n",
    "            target_data = target_data.reshape(-1)\n",
    "\n",
    "            mask = target_data != pad_token_id\n",
    "            prediction = prediction[mask]\n",
    "            target_data = target_data[mask]\n",
    "\n",
    "            loss = criterion(prediction, target_data) # Calculate loss\n",
    "        # Backward pass\n",
    "        scaler.scale(loss).backward()\n",
    "        torch.nn.utils.clip_grad_norm_(gemma.parameters(), max_norm=1.0)\n",
    "        scaler.step(optimizer)\n",
    "        scaler.update()\n",
    "        # Clear out grad\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Record loss\n",
    "        loss_train_epoch.append(loss.item())\n",
    "        scheduler.step()\n",
    "\n",
    "    loss_train.append(np.mean(loss_train_epoch))\n",
    "\n",
    "    gemma.eval()\n",
    "    with torch.no_grad():\n",
    "        for data in tqdm(validation_loader):\n",
    "            # Teacher forcing\n",
    "            data = trim_padding(data[0])\n",
    "            if data == None:\n",
    "                continue\n",
    "            input_data = data[:, :-1].long().cuda(device_id)\n",
    "            target_data = data[:, 1:].long().cuda(device_id)\n",
    "    \n",
    "            # Convert to embedding.\n",
    "            input_embeddings = word_embeddings_tensor[input_data]\n",
    "    \n",
    "            # Forward pass\n",
    "            with amp.autocast():\n",
    "                prediction = gemma(input_embeddings)\n",
    "    \n",
    "                # Change shape for loss calculation\n",
    "                prediction = prediction.view(-1, num_embeddings)\n",
    "                target_data = target_data.reshape(-1)\n",
    "\n",
    "                mask = target_data != pad_token_id\n",
    "                prediction = prediction[mask]\n",
    "                target_data = target_data[mask]\n",
    "                    \n",
    "                loss = criterion(prediction, target_data) # Calculate loss\n",
    "    \n",
    "            # Record loss\n",
    "            loss_val_epoch.append(loss.item())\n",
    "    \n",
    "        loss_valid.append(np.mean(loss_val_epoch))\n",
    "\n",
    "    # Save checkpoint\n",
    "    save_checkpoint(gemma, optimizer, epoch, loss_valid[-1])\n",
    "\n",
    "    plt.plot(loss_train, label=\"Training loss\")\n",
    "    plt.plot(loss_valid, label=\"Validation loss\")\n",
    "    print(\"Training loss: \", loss_train[-1])\n",
    "    print(\"Validation loss: \", loss_valid[-1])\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temperature = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence = \"arXiv is an open-access\"\n",
    "tokenized_sentence = tokenizer(sentence)[\"input_ids\"]\n",
    "if tokenized_sentence[-1] == 2:\n",
    "    tokenized_sentence = tokenized_sentence[:-1]\n",
    "gemma.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    while(tokenized_sentence[-1] != tokenizer.eos_token_id and len(tokenized_sentence) < max_token): # Keep iterating until reaches end of sentence or max token limit\n",
    "        # Preparing input\n",
    "        tokenized_sentence_tensor = torch.tensor(tokenized_sentence).cuda(device_id)\n",
    "        sentence_embedding = word_embeddings_tensor[tokenized_sentence_tensor]\n",
    "        sentence_embedding = sentence_embedding.unsqueeze(0).cuda(device_id)\n",
    "\n",
    "        # Make prediction\n",
    "        with amp.autocast():\n",
    "            prediction = gemma(sentence_embedding)\n",
    "        prediction = prediction[0][-1] # We only care about last token\n",
    "        prediction = prediction / temperature\n",
    "        prediction = F.softmax(prediction, dim=-1)\n",
    "        output_token = torch.multinomial(prediction, 1)\n",
    "\n",
    "        # Append to conversation history\n",
    "        tokenized_sentence.append(output_token.item())\n",
    "\n",
    "tokens = tokenizer.decode(tokenized_sentence, skip_special_tokens=True)\n",
    "print(tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(gemma, 'gemma3point156.pth')"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}