{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ['CUDA_VISIBLE_DEVICES'] = \"2\"\n", "\n", "import math\n", "import torch\n", "import einsum\n", "import numpy as np\n", "from torch import nn, einsum\n", "import torch.optim as optim\n", "from tqdm import tqdm\n", "from torch.nn.functional import mse_loss\n", "\n", "from suno_utils.utils.text import read_jsonl" ] }, { "cell_type": "code", "execution_count": 42, "metadata": {}, "outputs": [], "source": [ "\n", "def apply_delay_pattern(x: torch.Tensor, delay: int = 1, pad_token: int = 0):\n", " batch_size, seq_length, n_codebooks = x.size()\n", " \n", " # Calculate the maximum shift\n", " max_shift = delay * (n_codebooks - 1)\n", " \n", " # Create a new tensor filled with pad_token\n", " result = torch.full((batch_size, seq_length + max_shift, n_codebooks), \n", " pad_token, dtype=x.dtype, device=x.device)\n", " \n", " for i in range(n_codebooks):\n", " shift = delay * i\n", " result[:, shift:shift+seq_length, i] = x[:, :, i]\n", " \n", " return result\n", "\n", "def restore_original_alignment(x: torch.Tensor, delay: int = 1, pad_token: int = 0):\n", " batch_size, extended_seq_length, n_codebooks = x.size()\n", " \n", " # Calculate the original sequence length\n", " original_seq_length = extended_seq_length - delay * (n_codebooks - 1)\n", " \n", " # Create a new tensor to store the result\n", " result = torch.zeros(batch_size, original_seq_length, n_codebooks, \n", " dtype=x.dtype, device=x.device)\n", " \n", " for i in range(n_codebooks):\n", " shift = delay * i\n", " result[:, :, i] = x[:, shift:shift+original_seq_length, i]\n", " \n", " return result" ] }, { "cell_type": "code", "execution_count": 45, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "tensor([[[ 1., 2., 3.],\n", " [ 4., 5., 6.],\n", " [ 7., 8., 9.],\n", " [10., 11., 12.]]]) torch.Size([1, 4, 3])\n", "tensor([[[ 1., 100., 100.],\n", " [ 4., 100., 100.],\n", " [ 7., 100., 100.],\n", " [ 10., 2., 100.],\n", " [100., 5., 100.],\n", " [100., 8., 100.],\n", " [100., 11., 3.],\n", " [100., 100., 6.],\n", " [100., 100., 9.],\n", " [100., 100., 12.]]]) torch.Size([1, 10, 3])\n", "tensor([[[ 1., 2., 3.],\n", " [ 4., 5., 6.],\n", " [ 7., 8., 9.],\n", " [10., 11., 12.]]]) torch.Size([1, 4, 3])\n" ] } ], "source": [ "x = torch.tensor([[[1, 2], [3, 4], [5, 6], [7, 8]]], dtype=torch.float32)\n", "\n", "# create a test case with a batch size of 1, sequence length of 4, and 3 codebooks\n", "x = torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]], dtype=torch.float32)\n", "\n", "print(x, x.shape)\n", "\n", "delay = 3\n", "pad_token = 100\n", "\n", "x_delayed = apply_delay_pattern(x, delay=delay, pad_token=pad_token)\n", "print(x_delayed, x_delayed.shape)\n", "\n", "x_restored = restore_original_alignment(x_delayed, delay=delay, pad_token=pad_token)\n", "print(x_restored, x_restored.shape)" ] }, { "cell_type": "code", "execution_count": 58, "metadata": {}, "outputs": [], "source": [ "# 2. Dataset\n", "class MemmapDataset(torch.utils.data.Dataset):\n", " def __init__(\n", " self,\n", " acoustic_tokens_memmap_path: str,\n", " n_tokens_memmap: int,\n", " n_acoustic_codebooks: int,\n", " acoustic_codebook_size: int,\n", " ):\n", " acoustic_tokens = np.memmap(\n", " acoustic_tokens_memmap_path, dtype=np.uint16, mode=\"r\"\n", " )\n", " acoustic_tokens = acoustic_tokens.reshape(\n", " -1, n_tokens_memmap, n_acoustic_codebooks\n", " )\n", "\n", " print(f\"Acoustic tokens shape: {acoustic_tokens.shape}\")\n", " self.acoustic_tokens = acoustic_tokens\n", " self.acoustic_codebook_size = acoustic_codebook_size\n", " self.pad_token = acoustic_codebook_size + 1\n", "\n", " def __len__(self):\n", " return self.acoustic_tokens.shape[0]\n", "\n", " def __getitem__(self, idx):\n", " input_seq = torch.from_numpy(self.acoustic_tokens[idx, ...].copy())\n", " # input_seq: (n_tokens_memmap, n_acoustic_codebooks)\n", "\n", " # Apply delay pattern to input sequence\n", " input_seq = apply_delay_pattern(input_seq.unsqueeze(0), delay=1, pad_token=self.pad_token).squeeze(0)\n", "\n", " # create target sequence by shifting input sequence by 1\n", " target_seq = input_seq.clone()\n", " target_seq[:-1] = input_seq[1:]\n", " target_seq[-1] = self.pad_token\n", "\n", " return input_seq, target_seq" ] }, { "cell_type": "code", "execution_count": 61, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Acoustic tokens shape: (16650, 250, 12)\n", "torch.Size([261, 12]) torch.Size([261, 12])\n", "tensor([[ 732, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 527, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 128, 101, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 128, 1241, 58, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 449, 1649, 1508, 684, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 220, 381, 1242, 1914, 957, 1040, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 485, 746, 824, 957, 1585, 1731, 4097, 4097, 4097, 4097, 4097],\n", " [1080, 2021, 351, 226, 957, 1421, 1918, 114, 4097, 4097, 4097, 4097],\n", " [1769, 1668, 1487, 226, 1100, 11, 971, 114, 743, 4097, 4097, 4097],\n", " [1229, 1488, 1651, 358, 784, 11, 339, 806, 1372, 1712, 4097, 4097],\n", " [ 620, 273, 712, 1881, 325, 11, 920, 402, 743, 1873, 552, 4097],\n", " [ 620, 156, 1331, 1802, 1049, 1077, 19, 556, 1367, 1873, 1334, 670]],\n", " dtype=torch.uint16)\n", "tensor([[ 732, 527, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 128, 101, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 128, 1241, 58, 4097, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 449, 1649, 1508, 684, 4097, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 220, 381, 1242, 1914, 957, 1040, 4097, 4097, 4097, 4097, 4097, 4097],\n", " [ 732, 485, 746, 824, 957, 1585, 1731, 4097, 4097, 4097, 4097, 4097],\n", " [1080, 2021, 351, 226, 957, 1421, 1918, 114, 4097, 4097, 4097, 4097],\n", " [1769, 1668, 1487, 226, 1100, 11, 971, 114, 743, 4097, 4097, 4097],\n", " [1229, 1488, 1651, 358, 784, 11, 339, 806, 1372, 1712, 4097, 4097],\n", " [ 620, 273, 712, 1881, 325, 11, 920, 402, 743, 1873, 552, 4097],\n", " [ 620, 156, 1331, 1802, 1049, 1077, 19, 556, 1367, 1873, 1334, 670],\n", " [ 506, 116, 1497, 1374, 1070, 399, 1036, 1325, 686, 1818, 506, 670]],\n", " dtype=torch.uint16)\n" ] } ], "source": [ "acoustic_tokens_memmap_path = \"/app/suno/data/chirp_v4/vae/data_codec_val.bin\"\n", "n_tokens_memmap = 250\n", "n_acoustic_codebooks = 12\n", "acoustic_codebook_size = 4096\n", "\n", "dataset = MemmapDataset(acoustic_tokens_memmap_path, n_tokens_memmap, n_acoustic_codebooks, acoustic_codebook_size)\n", "\n", "for batch in dataset:\n", " input_seq, target_seq = batch\n", "\n", " print(input_seq.shape, target_seq.shape)\n", " print(input_seq[0:12, :])\n", " print(target_seq[0:12, :])\n", " break" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Simple sunoGPT" ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [], "source": [ "class CausalSelfAttention(nn.Module):\n", " def __init__(self, embed_size, num_heads):\n", " super().__init__()\n", " self.multihead_attn = nn.MultiheadAttention(embed_size, num_heads, batch_first=True)\n", " \n", " def forward(self, x):\n", " seq_length = x.size(1)\n", " # Create a causal mask\n", " mask = torch.triu(torch.ones(seq_length, seq_length), diagonal=1).bool()\n", " mask = mask.to(x.device)\n", " \n", " # Apply causal self-attention\n", " attn_output, _ = self.multihead_attn(x, x, x, attn_mask=mask)\n", " return attn_output\n", "\n", "# 1. Model Architecture\n", "class GPTModel(nn.Module):\n", " def __init__(self, acoustic_codebook_size: int, n_acoustic_codebooks: int, embed_size: int, num_heads: int, num_layers: int, max_seq_length: int,):\n", " super(GPTModel, self).__init__()\n", " self.token_embedding = nn.Embedding(acoustic_codebook_size, embed_size)\n", " self.position_embedding = nn.Embedding(max_seq_length, embed_size)\n", " \n", " self.layers = nn.ModuleList([\n", " nn.Sequential(\n", " CausalSelfAttention(embed_size, num_heads),\n", " nn.LayerNorm(embed_size),\n", " nn.Linear(embed_size, embed_size * 4),\n", " nn.GELU(),\n", " nn.Linear(embed_size * 4, embed_size),\n", " nn.LayerNorm(embed_size)\n", " )\n", " for _ in range(num_layers)\n", " ])\n", " \n", " self.fc_out = nn.Linear(embed_size, acoustic_codebook_size)\n", "\n", " def forward(self, x):\n", " seq_length = x.size(1)\n", " position_ids = torch.arange(seq_length, device=x.device).unsqueeze(0)\n", " \n", " token_embeds = self.token_embedding(x)\n", " position_embeds = self.position_embedding(position_ids)\n", " x = token_embeds + position_embeds\n", " \n", " for layer in self.layers:\n", " x = x + layer(x) # Residual connection\n", " \n", " return self.fc_out(x)\n", "\n", "# 2. Dataset\n", "class TextDataset(torch.utils.data.Dataset):\n", " def __init__(self, text, seq_length, vocab):\n", " self.text = text\n", " self.seq_length = seq_length\n", " self.vocab = vocab\n", "\n", " def __len__(self):\n", " return len(self.text) - self.seq_length\n", "\n", " def __getitem__(self, idx):\n", " input_seq = self.text[idx:idx+self.seq_length]\n", " target_seq = self.text[idx+1:idx+self.seq_length+1]\n", " return torch.tensor([self.vocab[c] for c in input_seq]), torch.tensor([self.vocab[c] for c in target_seq])\n", "\n", "# 3. Training Loop\n", "def train(model, dataloader, optimizer, criterion, device):\n", " model.train()\n", " total_loss = 0\n", " pbar = tqdm(dataloader, total=len(dataloader))\n", " for input_seq, target_seq in pbar:\n", " input_seq, target_seq = input_seq.to(device), target_seq.to(device)\n", "\n", " optimizer.zero_grad()\n", " output = model(input_seq)\n", " loss = criterion(output.view(-1, output.size(-1)), target_seq.view(-1))\n", " loss.backward()\n", " optimizer.step()\n", " \n", " total_loss += loss.item()\n", " pbar.set_description(f\"Loss: {loss.item():.4f}\")\n", " \n", " return total_loss / len(dataloader)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Train" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Number of GPT parameters: 5.28M\n", "Number of characters in text: 5.46M\n", "91\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " 0%| | 0/170566 [00:00 40\u001b[0m loss \u001b[38;5;241m=\u001b[39m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataloader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcriterion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n", "Cell \u001b[0;32mIn[14], line 74\u001b[0m, in \u001b[0;36mtrain\u001b[0;34m(model, dataloader, optimizer, criterion, device)\u001b[0m\n\u001b[1;32m 71\u001b[0m input_seq, target_seq \u001b[38;5;241m=\u001b[39m input_seq\u001b[38;5;241m.\u001b[39mto(device), target_seq\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m 73\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[0;32m---> 74\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_seq\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 75\u001b[0m loss \u001b[38;5;241m=\u001b[39m criterion(output\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, output\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)), target_seq\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m))\n\u001b[1;32m 76\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1530\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1539\u001b[0m \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1540\u001b[0m \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1544\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n", "Cell \u001b[0;32mIn[14], line 46\u001b[0m, in \u001b[0;36mGPTModel.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 43\u001b[0m x \u001b[38;5;241m=\u001b[39m token_embeds \u001b[38;5;241m+\u001b[39m position_embeds\n\u001b[1;32m 45\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlayers:\n\u001b[0;32m---> 46\u001b[0m x \u001b[38;5;241m=\u001b[39m x \u001b[38;5;241m+\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;66;03m# Residual connection\u001b[39;00m\n\u001b[1;32m 48\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfc_out(x)\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1530\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1539\u001b[0m \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1540\u001b[0m \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1544\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/container.py:217\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m 215\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m 216\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 217\u001b[0m \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 218\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1530\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1539\u001b[0m \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1540\u001b[0m \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 1544\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n", "Cell \u001b[0;32mIn[14], line 9\u001b[0m, in \u001b[0;36mCausalSelfAttention.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 7\u001b[0m seq_length \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m 8\u001b[0m \u001b[38;5;66;03m# Create a causal mask\u001b[39;00m\n\u001b[0;32m----> 9\u001b[0m mask \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtriu\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mones\u001b[49m\u001b[43m(\u001b[49m\u001b[43mseq_length\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mseq_length\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdiagonal\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mbool()\n\u001b[1;32m 10\u001b[0m mask \u001b[38;5;241m=\u001b[39m mask\u001b[38;5;241m.\u001b[39mto(x\u001b[38;5;241m.\u001b[39mdevice)\n\u001b[1;32m 12\u001b[0m \u001b[38;5;66;03m# Apply causal self-attention\u001b[39;00m\n", "\u001b[0;31mKeyboardInterrupt\u001b[0m: " ] } ], "source": [ " # Hyperparameters\n", "vocab_size = 1000 # Example value\n", "embed_size = 256\n", "num_heads = 8\n", "num_layers = 6\n", "max_seq_length = 100\n", "batch_size = 32\n", "num_epochs = 10\n", "learning_rate = 0.001\n", "\n", "# Device configuration\n", "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", "\n", "# Create model\n", "model = GPTModel(vocab_size, embed_size, num_heads, num_layers, max_seq_length).to(device)\n", "\n", "# print number of model parameters in millions\n", "num_params = sum(p.numel() for p in model.parameters()) / 1_000_000\n", "print(f'Number of GPT parameters: {num_params:.2f}M')\n", "\n", "# Create dataset and dataloader\n", "# Note: You'll need to implement text preprocessing and vocabulary creation\n", "with open(\"/home/christian/code/christian/data/shakespeare.txt\") as f:\n", " text = f.read()\n", "\n", "# print number of characters in millions in text\n", "print(f'Number of characters in text: {len(text) / 1_000_000:.2f}M')\n", "\n", "vocab = {char: i for i, char in enumerate(set(text))}\n", "print(len(vocab))\n", "dataset = TextDataset(text, max_seq_length, vocab)\n", "dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)\n", "\n", "# Loss and optimizer\n", "criterion = nn.CrossEntropyLoss()\n", "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n", "\n", "# Training loop\n", "for epoch in range(num_epochs):\n", " loss = train(model, dataloader, optimizer, criterion, device)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Inference" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def generate(model, start_sequence, max_length, temperature=1.0, top_k=None, top_p=None):\n", " model.eval() # Set the model to evaluation mode\n", " current_sequence = start_sequence.clone()\n", " \n", " for _ in range(max_length):\n", " with torch.no_grad(): # No need to track gradients for generation\n", " logits = model(current_sequence)\n", " next_token_logits = logits[:, -1, :] / temperature\n", " \n", " # Apply top-k filtering\n", " if top_k is not None:\n", " top_k = min(top_k, next_token_logits.size(-1))\n", " indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]\n", " next_token_logits[indices_to_remove] = float('-inf')\n", " \n", " # Apply top-p (nucleus) filtering\n", " if top_p is not None:\n", " sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)\n", " cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)\n", " sorted_indices_to_remove = cumulative_probs > top_p\n", " sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()\n", " sorted_indices_to_remove[..., 0] = 0\n", " indices_to_remove = sorted_indices[sorted_indices_to_remove]\n", " next_token_logits[indices_to_remove] = float('-inf')\n", " \n", " # Sample from the filtered distribution\n", " probs = torch.nn.functional.softmax(next_token_logits, dim=-1)\n", " next_token = torch.multinomial(probs, num_samples=1)\n", " \n", " # Append the chosen token to the sequence\n", " current_sequence = torch.cat((current_sequence, next_token), dim=1)\n", " \n", " return current_sequence\n", "\n", "# Example usage\n", "def generate_text(model, tokenizer, start_text, max_length=50, temperature=0.7, top_k=50, top_p=0.9):\n", " device = next(model.parameters()).device\n", " start_sequence = torch.tensor([tokenizer.encode(start_text)], dtype=torch.long).to(device)\n", " generated_sequence = generate(model, start_sequence, max_length, temperature, top_k, top_p)\n", " generated_text = tokenizer.decode(generated_sequence[0].tolist())\n", " return generated_text\n", "\n", "# Assuming we have a trained model and a tokenizer\n", "\n", "start_text = \"Once upon a time\"\n", "generated_text = generate_text(model, tokenizer, start_text)\n", "print(generated_text)" ] } ], "metadata": { "kernelspec": { "display_name": "suno_env", "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": 2 }