transformer_v1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Language Model with LORA and QLORA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5, Loss: 7.073467694796049\n",
      "Epoch 2/5, Loss: 7.075631141662598\n",
      "Epoch 3/5, Loss: 7.072393857515776\n",
      "Epoch 4/5, Loss: 7.069617161383996\n",
      "Epoch 5/5, Loss: 7.0713474200322075\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.optim import AdamW\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "from torch.utils.data import random_split\n",
    "\n",
    "\n",
    "class LORALayer(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, rank, alpha=1):\n",
    "        super(LORALayer, self).__init__()\n",
    "        self.rank = rank\n",
    "        self.alpha = alpha\n",
    "\n",
    "        # Original weight and bias of the linear layer\n",
    "        self.weight = nn.Parameter(torch.Tensor(output_dim, input_dim))\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "        self.bias = nn.Parameter(torch.Tensor(output_dim))\n",
    "\n",
    "        # LORA specific parameters\n",
    "        self.A = nn.Parameter(torch.Tensor(input_dim, rank))\n",
    "        self.B = nn.Parameter(torch.Tensor(rank, output_dim))\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
    "        nn.init.zeros_(self.bias)\n",
    "        nn.init.normal_(self.A, 0, 0.02)\n",
    "        nn.init.normal_(self.B, 0, 0.02)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #print(\"LORALayer Input Shape:\", x.shape)\n",
    "        \n",
    "        original_size = x.size()\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "        x_flattened = x.reshape(-1, original_size[-1])\n",
    "\n",
    "        # Compute lora_adjustment for each input in the batch\n",
    "        lora_adjustment = self.alpha * (x_flattened @ self.A) @ self.B\n",
    "        lora_adjustment = lora_adjustment.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        #print(\"Adjusted lora_adjustment Shape:\", lora_adjustment.shape)\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "        \n",
    "        # Apply linear transformation to x_flattened\n",
    "        x_transformed = nn.functional.linear(x_flattened, self.weight, self.bias)\n",
    "        x_transformed = x_transformed.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        # Add lora_adjustment to the transformed x\n",
    "        x = x_transformed + lora_adjustment\n",
    "        #print(\"LORALayer Output Shape:\", x.shape)\n",
    "\n",
    "        return x\n",
    "\n",
    "class QLORALayer(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, rank, alpha=1, quantization_bits=8):\n",
    "        super(QLORALayer, self).__init__()\n",
    "        self.rank = rank\n",
    "        self.alpha = alpha\n",
    "        self.quantization_bits = quantization_bits\n",
    "\n",
    "        # Original weight and bias\n",
    "        self.weight = nn.Parameter(torch.Tensor(output_dim, input_dim))\n",
    "        self.bias = nn.Parameter(torch.Tensor(output_dim))\n",
    "\n",
    "        # QLORA specific parameters\n",
    "        self.A = nn.Parameter(torch.Tensor(input_dim, rank))\n",
    "        self.B = nn.Parameter(torch.Tensor(rank, output_dim))\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
    "        nn.init.zeros_(self.bias)\n",
    "        nn.init.normal_(self.A, 0, 0.02)\n",
    "        nn.init.normal_(self.B, 0, 0.02)\n",
    "\n",
    "    def quantize(self, x, num_bits):\n",
    "        # Implement a simple quantization method\n",
    "        scale = x.abs().max()\n",
    "        x_quantized = torch.round(x / scale * (2**num_bits - 1))\n",
    "        return x_quantized, scale\n",
    "\n",
    "    def forward(self, x):\n",
    "        #print(\"QLORALayer Input Shape:\", x.shape)\n",
    "        original_size = x.size()\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "        x_flattened = x.reshape(-1, original_size[-1])\n",
    "\n",
    "        A_quantized, scale_A = self.quantize(self.A, self.quantization_bits)\n",
    "        B_quantized, scale_B = self.quantize(self.B, self.quantization_bits)\n",
    "\n",
    "        # Compute lora_adjustment for each input in the batch\n",
    "        lora_adjustment = self.alpha * (x_flattened @ (A_quantized / scale_A)) @ (B_quantized / scale_B)\n",
    "        lora_adjustment = lora_adjustment.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        #print(\"Adjusted lora_adjustment Shape:\", lora_adjustment.shape)\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "\n",
    "        # Apply linear transformation to x_flattened\n",
    "        x_transformed = nn.functional.linear(x_flattened, self.weight, self.bias)\n",
    "        x_transformed = x_transformed.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        # Add lora_adjustment to the transformed x\n",
    "        x = x_transformed + lora_adjustment\n",
    "        #print(\"QLORALayer Output Shape:\", x.shape)\n",
    "\n",
    "        return x\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, embed_size, heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.embed_size = embed_size\n",
    "        self.heads = heads\n",
    "        self.head_dim = embed_size // heads\n",
    "\n",
    "        assert (\n",
    "            self.head_dim * heads == embed_size\n",
    "        ), \"Embedding size needs to be divisible by heads\"\n",
    "\n",
    "        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)\n",
    "\n",
    "    def forward(self, values, keys, query, mask):\n",
    "        N = query.shape[0]\n",
    "        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]\n",
    "\n",
    "        # Split the embedding into self.heads different pieces\n",
    "        values = values.reshape(N, value_len, self.heads, self.head_dim)\n",
    "        keys = keys.reshape(N, key_len, self.heads, self.head_dim)\n",
    "        queries = query.reshape(N, query_len, self.heads, self.head_dim)\n",
    "\n",
    "        values = self.values(values)\n",
    "        keys = self.keys(keys)\n",
    "        queries = self.queries(queries)\n",
    "\n",
    "        # Einsum does the matrix multiplication for query*keys for each training example\n",
    "        # with every other training example, then sum it up\n",
    "        attention = torch.einsum(\"nqhd,nkhd->nhqk\", [queries, keys])\n",
    "\n",
    "        if mask is not None:\n",
    "            attention = attention.masked_fill(mask == 0, float(\"-1e20\"))\n",
    "\n",
    "        attention = torch.softmax(attention / (self.embed_size ** (1 / 2)), dim=3)\n",
    "\n",
    "        out = torch.einsum(\"nhql,nlhd->nqhd\", [attention, values]).reshape(\n",
    "            N, query_len, self.heads * self.head_dim\n",
    "        )\n",
    "\n",
    "        out = self.fc_out(out)\n",
    "        return out\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, embed_size, heads, dropout, forward_expansion, rank):\n",
    "        super(TransformerBlock, self).__init__()\n",
    "        self.attention = MultiHeadAttention(embed_size, heads)\n",
    "        self.norm1 = nn.LayerNorm(embed_size)\n",
    "        self.norm2 = nn.LayerNorm(embed_size)\n",
    "\n",
    "        self.feed_forward = nn.Sequential(\n",
    "            LORALayer(embed_size, forward_expansion * embed_size, rank),\n",
    "            nn.ReLU(),\n",
    "            LORALayer(forward_expansion * embed_size, embed_size, rank),\n",
    "        )\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, value, key, query, mask):\n",
    "        attention = self.attention(value, key, query, mask)\n",
    "        x = self.dropout(self.norm1(attention + query))\n",
    "        forward = self.feed_forward(x)\n",
    "        out = self.dropout(self.norm2(forward + x))\n",
    "        return out\n",
    "\n",
    "class LanguageModelDecoder(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_size, num_layers, heads, forward_expansion, dropout, max_length, rank):\n",
    "        super(LanguageModelDecoder, self).__init__()\n",
    "        self.word_embedding = nn.Embedding(vocab_size, embed_size)\n",
    "        self.position_embedding = nn.Embedding(max_length, embed_size)\n",
    "\n",
    "        # Adding BatchNorm layers\n",
    "        self.bn1 = nn.BatchNorm1d(embed_size)\n",
    "        self.bn2 = nn.BatchNorm1d(embed_size)\n",
    "\n",
    "        self.layers = nn.ModuleList(\n",
    "            [\n",
    "                TransformerBlock(embed_size, heads, dropout, forward_expansion, rank)\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        # QLORA layers\n",
    "        self.qlora_feed_forward = nn.Sequential(\n",
    "            QLORALayer(embed_size, forward_expansion * embed_size, rank),\n",
    "            nn.ReLU(),\n",
    "            QLORALayer(forward_expansion * embed_size, embed_size, rank),\n",
    "        )\n",
    "        self.use_qlora = False  # Flag to toggle QLORA\n",
    "\n",
    "        self.fc_out = nn.Linear(embed_size, vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, trg_mask):\n",
    "        N, seq_length = x.shape\n",
    "        positions = torch.arange(0, seq_length).expand(N, seq_length).to(x.device)\n",
    "        x = self.dropout(self.word_embedding(x) + self.position_embedding(positions))\n",
    "\n",
    "        # Transpose for BatchNorm, apply batch normalization, and then transpose back\n",
    "        x = x.transpose(1, 2)\n",
    "        x = self.bn1(x)\n",
    "        x = x.transpose(1, 2)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, x, x, trg_mask)\n",
    "            if self.use_qlora:\n",
    "                x = self.qlora_feed_forward(x)\n",
    "\n",
    "        # Transpose for BatchNorm, apply batch normalization, and then transpose back\n",
    "        x = x.transpose(1, 2)\n",
    "        x = self.bn2(x)\n",
    "        x = x.transpose(1, 2)\n",
    "\n",
    "        out = self.fc_out(x)\n",
    "        return out\n",
    "\n",
    "    def toggle_qlora(self, use_qlora: bool):\n",
    "        self.use_qlora = use_qlora\n",
    "\n",
    "\n",
    " \n",
    "\n",
    "class LanguageModelTransformer(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_size=256, num_layers=6, forward_expansion=4, heads=8, dropout=0, max_length=100, rank=16):\n",
    "        super(LanguageModelTransformer, self).__init__()\n",
    "\n",
    "        self.decoder = LanguageModelDecoder(\n",
    "            vocab_size,\n",
    "            embed_size,\n",
    "            num_layers,\n",
    "            heads,\n",
    "            forward_expansion,\n",
    "            dropout,\n",
    "            max_length,\n",
    "            rank,\n",
    "        )\n",
    "\n",
    "    def forward(self, trg):\n",
    "        trg_mask = self.make_trg_mask(trg)\n",
    "        out = self.decoder(trg, trg_mask)\n",
    "        return out\n",
    "\n",
    "    def make_trg_mask(self, trg):\n",
    "        N, trg_len = trg.shape\n",
    "        trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(\n",
    "            N, 1, trg_len, trg_len\n",
    "        ).to(trg.device)\n",
    "\n",
    "        return trg_mask\n",
    "\n",
    "class LanguageDataset(Dataset):\n",
    "    def __init__(self, tokenized_texts, sequence_length):\n",
    "        self.tokenized_texts = tokenized_texts\n",
    "        self.sequence_length = sequence_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.tokenized_texts) - self.sequence_length\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return (\n",
    "            self.tokenized_texts[index:index+self.sequence_length].clone().detach(),\n",
    "            self.tokenized_texts[index+1:index+self.sequence_length+1].clone().detach()\n",
    "        )\n",
    "\n",
    "# Define vocabulary size and dummy data parameters\n",
    "NUM_WORDS = 1000  # Example vocabulary size\n",
    "sequence_length = 30  # Sequence length for the LanguageDataset\n",
    "dummy_data_size = 1000  # Total number of tokens in the dummy dataset\n",
    "\n",
    "# Generate random tokenized data\n",
    "tokenized_train_data = torch.randint(high=NUM_WORDS, size=(dummy_data_size,))\n",
    "\n",
    "# Create the complete dataset\n",
    "complete_dataset = LanguageDataset(tokenized_train_data, sequence_length)\n",
    "\n",
    "# Calculate the actual length of the dataset\n",
    "actual_dataset_length = len(complete_dataset)\n",
    "\n",
    "# Define the size of your validation set\n",
    "validation_size = int(0.2 * actual_dataset_length)  # 20% of the dataset\n",
    "training_size = actual_dataset_length - validation_size\n",
    "\n",
    "# Split the dataset into training and validation sets\n",
    "train_dataset, val_dataset = random_split(complete_dataset, [training_size, validation_size])\n",
    "\n",
    "# Create DataLoaders for the training and validation sets\n",
    "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)\n",
    "\n",
    "# Define device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Create the model instance\n",
    "model = LanguageModelTransformer(\n",
    "    vocab_size=NUM_WORDS,\n",
    "    embed_size=256,  # You can adjust these parameters as needed\n",
    "    num_layers=6,\n",
    "    forward_expansion=4,\n",
    "    heads=8,\n",
    "    dropout=0,\n",
    "    max_length=100,\n",
    "    rank=16\n",
    ").to(device)\n",
    "\n",
    "\n",
    "\n",
    "# Enable QLORA during training\n",
    "model.decoder.toggle_qlora(True)\n",
    "\n",
    "# Training loop\n",
    "# Assuming model is an instance of LanguageModelTransformer\n",
    "# Define loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = AdamW(model.parameters(), lr=0.0000001)  # Consider reducing this if still problematic\n",
    "scheduler = StepLR(optimizer, step_size=1, gamma=0.98)\n",
    "num_epochs = 5\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    model.decoder.toggle_qlora(True)\n",
    "    total_loss = 0\n",
    "\n",
    "    for batch in train_loader:\n",
    "        inputs, targets = batch\n",
    "        inputs, targets = inputs.to(device), targets.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs.view(-1, NUM_WORDS), targets.view(-1))\n",
    "        \n",
    "        # Check for NaN in loss\n",
    "        if math.isnan(loss.item()):\n",
    "            print(\"Encountered NaN loss, stopping training\")\n",
    "            break\n",
    "\n",
    "        loss.backward()\n",
    "        \n",
    "        # Gradient clipping\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1)\n",
    "\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "\n",
    "    scheduler.step()\n",
    "\n",
    "    # Check for NaN in total_loss\n",
    "    if math.isnan(total_loss):\n",
    "        print(f\"Epoch {epoch+1}/{num_epochs} stopped due to NaN loss\")\n",
    "        break\n",
    "    else:\n",
    "        print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss / len(train_loader)}\")\n",
    "\n",
    "    model.decoder.toggle_qlora(False)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Loss: 8.085800170898438, Perplexity: 3248.017765364175\n"
     ]
    }
   ],
   "source": [
    "def calculate_perplexity(loss):\n",
    "    return math.exp(loss)\n",
    "\n",
    "def evaluate(model, data_loader, device):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in data_loader:\n",
    "            inputs, targets = batch  # Adjust based on your dataset\n",
    "            inputs, targets = inputs.to(device), targets.to(device)  # Move inputs and targets to the same device as the model\n",
    "\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs.view(-1, NUM_WORDS), targets.view(-1))\n",
    "            total_loss += loss.item()\n",
    "\n",
    "    average_loss = total_loss / len(data_loader)\n",
    "    perplexity = calculate_perplexity(average_loss)\n",
    "\n",
    "    return average_loss, perplexity\n",
    "\n",
    "# Define device (if not already defined)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Calculate validation loss and perplexity\n",
    "val_loss, val_perplexity = evaluate(model, val_loader, device)\n",
    "print(f\"Validation Loss: {val_loss}, Perplexity: {val_perplexity}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TEST ON dATASET 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: datasets in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (2.14.7)\n",
      "Requirement already satisfied: numpy>=1.17 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (1.26.2)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (11.0.0)\n",
      "Requirement already satisfied: pyarrow-hotfix in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (0.6)\n",
      "Requirement already satisfied: dill<0.3.8,>=0.3.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (0.3.6)\n",
      "Requirement already satisfied: pandas in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (2.1.4)\n",
      "Requirement already satisfied: requests>=2.19.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (2.31.0)\n",
      "Requirement already satisfied: tqdm>=4.62.1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (4.65.0)\n",
      "Requirement already satisfied: xxhash in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (2.0.2)\n",
      "Requirement already satisfied: multiprocess in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (0.70.14)\n",
      "Requirement already satisfied: fsspec<=2023.10.0,>=2023.1.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from fsspec[http]<=2023.10.0,>=2023.1.0->datasets) (2023.10.0)\n",
      "Requirement already satisfied: aiohttp in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (3.9.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0.0,>=0.14.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (0.19.4)\n",
      "Requirement already satisfied: packaging in c:\\users\\robbi\\appdata\\roaming\\python\\python311\\site-packages (from datasets) (23.2)\n",
      "Requirement already satisfied: pyyaml>=5.1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from datasets) (6.0.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from aiohttp->datasets) (23.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from aiohttp->datasets) (6.0.4)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from aiohttp->datasets) (1.9.3)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from aiohttp->datasets) (1.4.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from aiohttp->datasets) (1.2.0)\n",
      "Requirement already satisfied: filelock in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from huggingface-hub<1.0.0,>=0.14.0->datasets) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from huggingface-hub<1.0.0,>=0.14.0->datasets) (4.7.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from requests>=2.19.0->datasets) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from requests>=2.19.0->datasets) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from requests>=2.19.0->datasets) (1.26.18)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from requests>=2.19.0->datasets) (2023.11.17)\n",
      "Requirement already satisfied: colorama in c:\\users\\robbi\\appdata\\roaming\\python\\python311\\site-packages (from tqdm>=4.62.1->datasets) (0.4.6)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in c:\\users\\robbi\\appdata\\roaming\\python\\python311\\site-packages (from pandas->datasets) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from pandas->datasets) (2023.3.post1)\n",
      "Requirement already satisfied: tzdata>=2022.1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from pandas->datasets) (2023.3)\n",
      "Requirement already satisfied: six>=1.5 in c:\\users\\robbi\\appdata\\roaming\\python\\python311\\site-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting apache_beam\n",
      "  Downloading apache_beam-2.53.0-cp311-cp311-win_amd64.whl.metadata (6.7 kB)\n",
      "Collecting crcmod<2.0,>=1.7 (from apache_beam)\n",
      "  Downloading crcmod-1.7.tar.gz (89 kB)\n",
      "     ---------------------------------------- 0.0/89.7 kB ? eta -:--:--\n",
      "     ---- ----------------------------------- 10.2/89.7 kB ? eta -:--:--\n",
      "     -------------------------- ----------- 61.4/89.7 kB 825.8 kB/s eta 0:00:01\n",
      "     -------------------------------------- 89.7/89.7 kB 854.1 kB/s eta 0:00:00\n",
      "  Preparing metadata (setup.py): started\n",
      "  Preparing metadata (setup.py): finished with status 'done'\n",
      "Collecting orjson<4,>=3.9.7 (from apache_beam)\n",
      "  Downloading orjson-3.9.10-cp311-none-win_amd64.whl.metadata (50 kB)\n",
      "     ---------------------------------------- 0.0/50.5 kB ? eta -:--:--\n",
      "     -------- ------------------------------- 10.2/50.5 kB ? eta -:--:--\n",
      "     -------------------------------------- 50.5/50.5 kB 649.0 kB/s eta 0:00:00\n",
      "Collecting dill<0.3.2,>=0.3.1.1 (from apache_beam)\n",
      "  Downloading dill-0.3.1.1.tar.gz (151 kB)\n",
      "     ---------------------------------------- 0.0/152.0 kB ? eta -:--:--\n",
      "     ------- ------------------------------- 30.7/152.0 kB 1.3 MB/s eta 0:00:01\n",
      "     --------- --------------------------- 41.0/152.0 kB 393.8 kB/s eta 0:00:01\n",
      "     ------------------- ----------------- 81.9/152.0 kB 919.0 kB/s eta 0:00:01\n",
      "     -------------------------- --------- 112.6/152.0 kB 652.2 kB/s eta 0:00:01\n",
      "     ------------------------------------ 152.0/152.0 kB 647.2 kB/s eta 0:00:00\n",
      "  Preparing metadata (setup.py): started\n",
      "  Preparing metadata (setup.py): finished with status 'done'\n",
      "Collecting cloudpickle~=2.2.1 (from apache_beam)\n",
      "  Downloading cloudpickle-2.2.1-py3-none-any.whl (25 kB)\n",
      "Collecting fastavro<2,>=0.23.6 (from apache_beam)\n",
      "  Downloading fastavro-1.9.3-cp311-cp311-win_amd64.whl.metadata (5.7 kB)\n",
      "Collecting fasteners<1.0,>=0.3 (from apache_beam)\n",
      "  Downloading fasteners-0.19-py3-none-any.whl.metadata (4.9 kB)\n",
      "Collecting grpcio!=1.48.0,<2,>=1.33.1 (from apache_beam)\n",
      "  Downloading grpcio-1.60.0-cp311-cp311-win_amd64.whl.metadata (4.2 kB)\n",
      "Collecting hdfs<3.0.0,>=2.1.0 (from apache_beam)\n",
      "  Downloading hdfs-2.7.3.tar.gz (43 kB)\n",
      "     ---------------------------------------- 0.0/43.5 kB ? eta -:--:--\n",
      "     ----------------------------------- -- 41.0/43.5 kB 960.0 kB/s eta 0:00:01\n",
      "     ---------------------------------------- 43.5/43.5 kB 1.0 MB/s eta 0:00:00\n",
      "  Preparing metadata (setup.py): started\n",
      "  Preparing metadata (setup.py): finished with status 'done'\n",
      "Collecting httplib2<0.23.0,>=0.8 (from apache_beam)\n",
      "  Downloading httplib2-0.22.0-py3-none-any.whl (96 kB)\n",
      "     ---------------------------------------- 0.0/96.9 kB ? eta -:--:--\n",
      "     ---------------- ----------------------- 41.0/96.9 kB 1.9 MB/s eta 0:00:01\n",
      "     -------------------------------------- - 92.2/96.9 kB 1.0 MB/s eta 0:00:01\n",
      "     ---------------------------------------- 96.9/96.9 kB 1.1 MB/s eta 0:00:00\n",
      "Collecting js2py<1,>=0.74 (from apache_beam)\n",
      "  Downloading Js2Py-0.74-py3-none-any.whl (1.0 MB)\n",
      "     ---------------------------------------- 0.0/1.0 MB ? eta -:--:--\n",
      "     - -------------------------------------- 0.0/1.0 MB 1.9 MB/s eta 0:00:01\n",
      "     -- ------------------------------------- 0.1/1.0 MB 1.6 MB/s eta 0:00:01\n",
      "     ----- ---------------------------------- 0.1/1.0 MB 1.2 MB/s eta 0:00:01\n",
      "     ------ --------------------------------- 0.2/1.0 MB 893.0 kB/s eta 0:00:01\n",
      "     ------- -------------------------------- 0.2/1.0 MB 953.7 kB/s eta 0:00:01\n",
      "     -------- ------------------------------- 0.2/1.0 MB 981.9 kB/s eta 0:00:01\n",
      "     ----------- ---------------------------- 0.3/1.0 MB 1.1 MB/s eta 0:00:01\n",
      "     ----------- ---------------------------- 0.3/1.0 MB 1.1 MB/s eta 0:00:01\n",
      "     -------------- ------------------------- 0.4/1.0 MB 955.7 kB/s eta 0:00:01\n",
      "     ----------------- ---------------------- 0.4/1.0 MB 981.9 kB/s eta 0:00:01\n",
      "     ----------------- ---------------------- 0.5/1.0 MB 880.6 kB/s eta 0:00:01\n",
      "     -------------------- ------------------- 0.5/1.0 MB 962.4 kB/s eta 0:00:01\n",
      "     ---------------------- ----------------- 0.6/1.0 MB 972.7 kB/s eta 0:00:01\n",
      "     ------------------------ --------------- 0.6/1.0 MB 999.4 kB/s eta 0:00:01\n",
      "     -------------------------- ------------- 0.7/1.0 MB 1.0 MB/s eta 0:00:01\n",
      "     ------------------------------ --------- 0.8/1.0 MB 1.1 MB/s eta 0:00:01\n",
      "     ------------------------------ --------- 0.8/1.0 MB 1.1 MB/s eta 0:00:01\n",
      "     --------------------------------- ------ 0.9/1.0 MB 1.0 MB/s eta 0:00:01\n",
      "     ---------------------------------- ----- 0.9/1.0 MB 1.0 MB/s eta 0:00:01\n",
      "     ------------------------------------- -- 1.0/1.0 MB 1.1 MB/s eta 0:00:01\n",
      "     -------------------------------------- - 1.0/1.0 MB 1.0 MB/s eta 0:00:01\n",
      "     ---------------------------------------- 1.0/1.0 MB 1.0 MB/s eta 0:00:00\n",
      "Requirement already satisfied: jsonschema<5.0.0,>=4.0.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from apache_beam) (4.19.2)\n",
      "Collecting jsonpickle<4.0.0,>=3.0.0 (from apache_beam)\n",
      "  Downloading jsonpickle-3.0.2-py3-none-any.whl.metadata (7.5 kB)\n",
      "Collecting numpy<1.25.0,>=1.14.3 (from apache_beam)\n",
      "  Downloading numpy-1.24.4-cp311-cp311-win_amd64.whl.metadata (5.6 kB)\n",
      "Collecting objsize<0.7.0,>=0.6.1 (from apache_beam)\n",
      "  Downloading objsize-0.6.1-py3-none-any.whl (9.3 kB)\n",
      "Requirement already satisfied: packaging>=22.0 in c:\\users\\robbi\\appdata\\roaming\\python\\python311\\site-packages (from apache_beam) (23.2)\n",
      "Collecting pymongo<5.0.0,>=3.8.0 (from apache_beam)\n",
      "  Downloading pymongo-4.6.1-cp311-cp311-win_amd64.whl.metadata (22 kB)\n",
      "Collecting proto-plus<2,>=1.7.1 (from apache_beam)\n",
      "  Downloading proto_plus-1.23.0-py3-none-any.whl.metadata (2.2 kB)\n",
      "Requirement already satisfied: protobuf!=4.0.*,!=4.21.*,!=4.22.0,!=4.23.*,!=4.24.*,<4.26.0,>=3.20.3 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from apache_beam) (3.20.3)\n",
      "Collecting pydot<2,>=1.2.0 (from apache_beam)\n",
      "  Downloading pydot-1.4.2-py2.py3-none-any.whl (21 kB)\n",
      "Requirement already satisfied: python-dateutil<3,>=2.8.0 in c:\\users\\robbi\\appdata\\roaming\\python\\python311\\site-packages (from apache_beam) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2018.3 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from apache_beam) (2023.3.post1)\n",
      "Requirement already satisfied: regex>=2020.6.8 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from apache_beam) (2023.10.3)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.24.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from apache_beam) (2.31.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from apache_beam) (4.7.1)\n",
      "Collecting zstandard<1,>=0.18.0 (from apache_beam)\n",
      "  Downloading zstandard-0.22.0-cp311-cp311-win_amd64.whl.metadata (3.0 kB)\n",
      "Requirement already satisfied: pyarrow<12.0.0,>=3.0.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from apache_beam) (11.0.0)\n",
      "Requirement already satisfied: pyarrow-hotfix<1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from apache_beam) (0.6)\n",
      "Collecting docopt (from hdfs<3.0.0,>=2.1.0->apache_beam)\n",
      "  Downloading docopt-0.6.2.tar.gz (25 kB)\n",
      "  Preparing metadata (setup.py): started\n",
      "  Preparing metadata (setup.py): finished with status 'done'\n",
      "Requirement already satisfied: six>=1.9.0 in c:\\users\\robbi\\appdata\\roaming\\python\\python311\\site-packages (from hdfs<3.0.0,>=2.1.0->apache_beam) (1.16.0)\n",
      "Collecting pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 (from httplib2<0.23.0,>=0.8->apache_beam)\n",
      "  Downloading pyparsing-3.1.1-py3-none-any.whl.metadata (5.1 kB)\n",
      "Collecting tzlocal>=1.2 (from js2py<1,>=0.74->apache_beam)\n",
      "  Downloading tzlocal-5.2-py3-none-any.whl.metadata (7.8 kB)\n",
      "Collecting pyjsparser>=2.5.1 (from js2py<1,>=0.74->apache_beam)\n",
      "  Downloading pyjsparser-2.7.1.tar.gz (24 kB)\n",
      "  Preparing metadata (setup.py): started\n",
      "  Preparing metadata (setup.py): finished with status 'done'\n",
      "Requirement already satisfied: attrs>=22.2.0 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from jsonschema<5.0.0,>=4.0.0->apache_beam) (23.1.0)\n",
      "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from jsonschema<5.0.0,>=4.0.0->apache_beam) (2023.7.1)\n",
      "Requirement already satisfied: referencing>=0.28.4 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from jsonschema<5.0.0,>=4.0.0->apache_beam) (0.30.2)\n",
      "Requirement already satisfied: rpds-py>=0.7.1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from jsonschema<5.0.0,>=4.0.0->apache_beam) (0.10.6)\n",
      "Collecting dnspython<3.0.0,>=1.16.0 (from pymongo<5.0.0,>=3.8.0->apache_beam)\n",
      "  Downloading dnspython-2.4.2-py3-none-any.whl.metadata (4.9 kB)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from requests<3.0.0,>=2.24.0->apache_beam) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from requests<3.0.0,>=2.24.0->apache_beam) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from requests<3.0.0,>=2.24.0->apache_beam) (1.26.18)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from requests<3.0.0,>=2.24.0->apache_beam) (2023.11.17)\n",
      "Requirement already satisfied: tzdata in c:\\users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\lib\\site-packages (from tzlocal>=1.2->js2py<1,>=0.74->apache_beam) (2023.3)\n",
      "Downloading apache_beam-2.53.0-cp311-cp311-win_amd64.whl (5.0 MB)\n",
      "   ---------------------------------------- 0.0/5.0 MB ? eta -:--:--\n",
      "   ---------------------------------------- 0.1/5.0 MB 1.7 MB/s eta 0:00:03\n",
      "   - -------------------------------------- 0.1/5.0 MB 1.7 MB/s eta 0:00:03\n",
      "   - -------------------------------------- 0.2/5.0 MB 1.7 MB/s eta 0:00:03\n",
      "   -- ------------------------------------- 0.3/5.0 MB 1.7 MB/s eta 0:00:03\n",
      "   --- ------------------------------------ 0.4/5.0 MB 1.7 MB/s eta 0:00:03\n",
      "   --- ------------------------------------ 0.5/5.0 MB 1.7 MB/s eta 0:00:03\n",
      "   ---- ----------------------------------- 0.5/5.0 MB 1.5 MB/s eta 0:00:03\n",
      "   ---- ----------------------------------- 0.6/5.0 MB 1.6 MB/s eta 0:00:03\n",
      "   ----- ---------------------------------- 0.7/5.0 MB 1.6 MB/s eta 0:00:03\n",
      "   ----- ---------------------------------- 0.7/5.0 MB 1.6 MB/s eta 0:00:03\n",
      "   ------ --------------------------------- 0.8/5.0 MB 1.6 MB/s eta 0:00:03\n",
      "   ------- -------------------------------- 0.9/5.0 MB 1.6 MB/s eta 0:00:03\n",
      "   ------- -------------------------------- 0.9/5.0 MB 1.6 MB/s eta 0:00:03\n",
      "   ------- -------------------------------- 1.0/5.0 MB 1.5 MB/s eta 0:00:03\n",
      "   -------- ------------------------------- 1.0/5.0 MB 1.5 MB/s eta 0:00:03\n",
      "   -------- ------------------------------- 1.1/5.0 MB 1.5 MB/s eta 0:00:03\n",
      "   -------- ------------------------------- 1.1/5.0 MB 1.5 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 1.4 MB/s eta 0:00:03\n",
      "   --------- ------------------------------ 1.2/5.0 MB 529.1 kB/s eta 0:00:08\n",
      "   --------- ------------------------------ 1.2/5.0 MB 529.1 kB/s eta 0:00:08\n",
      "   --------- ------------------------------ 1.2/5.0 MB 523.4 kB/s eta 0:00:08\n",
      "   --------- ------------------------------ 1.2/5.0 MB 518.2 kB/s eta 0:00:08\n",
      "   --------- ------------------------------ 1.2/5.0 MB 508.8 kB/s eta 0:00:08\n",
      "   --------- ------------------------------ 1.2/5.0 MB 508.8 kB/s eta 0:00:08\n",
      "   --------- ------------------------------ 1.2/5.0 MB 504.2 kB/s eta 0:00:08\n",
      "   --------- ------------------------------ 1.2/5.0 MB 495.6 kB/s eta 0:00:08\n",
      "   ---------- ----------------------------- 1.3/5.0 MB 491.6 kB/s eta 0:00:08\n",
      "   ---------- ----------------------------- 1.3/5.0 MB 493.4 kB/s eta 0:00:08\n",
      "   ---------- ----------------------------- 1.3/5.0 MB 497.3 kB/s eta 0:00:08\n",
      "   ----------- ---------------------------- 1.4/5.0 MB 507.6 kB/s eta 0:00:08\n",
      "   ----------- ---------------------------- 1.4/5.0 MB 517.6 kB/s eta 0:00:07\n",
      "   ----------- ---------------------------- 1.5/5.0 MB 530.9 kB/s eta 0:00:07\n",
      "   ------------ --------------------------- 1.5/5.0 MB 543.1 kB/s eta 0:00:07\n",
      "   ------------ --------------------------- 1.6/5.0 MB 549.0 kB/s eta 0:00:07\n",
      "   ------------- -------------------------- 1.7/5.0 MB 561.3 kB/s eta 0:00:06\n",
      "   ------------- -------------------------- 1.7/5.0 MB 569.1 kB/s eta 0:00:06\n",
      "   -------------- ------------------------- 1.8/5.0 MB 581.1 kB/s eta 0:00:06\n",
      "   -------------- ------------------------- 1.8/5.0 MB 592.2 kB/s eta 0:00:06\n",
      "   --------------- ------------------------ 1.9/5.0 MB 606.2 kB/s eta 0:00:06\n",
      "   --------------- ------------------------ 2.0/5.0 MB 616.6 kB/s eta 0:00:05\n",
      "   ---------------- ----------------------- 2.0/5.0 MB 633.0 kB/s eta 0:00:05\n",
      "   ---------------- ----------------------- 2.1/5.0 MB 645.9 kB/s eta 0:00:05\n",
      "   ----------------- ---------------------- 2.2/5.0 MB 655.5 kB/s eta 0:00:05\n",
      "   ----------------- ---------------------- 2.2/5.0 MB 664.5 kB/s eta 0:00:05\n",
      "   ------------------ --------------------- 2.3/5.0 MB 676.5 kB/s eta 0:00:04\n",
      "   ------------------ --------------------- 2.4/5.0 MB 682.1 kB/s eta 0:00:04\n",
      "   ------------------- -------------------- 2.5/5.0 MB 699.2 kB/s eta 0:00:04\n",
      "   -------------------- ------------------- 2.5/5.0 MB 713.1 kB/s eta 0:00:04\n",
      "   -------------------- ------------------- 2.6/5.0 MB 726.9 kB/s eta 0:00:04\n",
      "   --------------------- ------------------ 2.7/5.0 MB 742.5 kB/s eta 0:00:04\n",
      "   ---------------------- ----------------- 2.8/5.0 MB 755.4 kB/s eta 0:00:03\n",
      "   ----------------------- ---------------- 2.9/5.0 MB 771.0 kB/s eta 0:00:03\n",
      "   ----------------------- ---------------- 3.0/5.0 MB 783.1 kB/s eta 0:00:03\n",
      "   ------------------------ --------------- 3.0/5.0 MB 795.1 kB/s eta 0:00:03\n",
      "   ------------------------ --------------- 3.1/5.0 MB 803.3 kB/s eta 0:00:03\n",
      "   ------------------------- -------------- 3.1/5.0 MB 805.3 kB/s eta 0:00:03\n",
      "   ------------------------- -------------- 3.1/5.0 MB 805.3 kB/s eta 0:00:03\n",
      "   ------------------------- -------------- 3.2/5.0 MB 803.8 kB/s eta 0:00:03\n",
      "   -------------------------- ------------- 3.3/5.0 MB 811.7 kB/s eta 0:00:03\n",
      "   --------------------------- ------------ 3.4/5.0 MB 822.4 kB/s eta 0:00:02\n",
      "   --------------------------- ------------ 3.5/5.0 MB 830.9 kB/s eta 0:00:02\n",
      "   ---------------------------- ----------- 3.5/5.0 MB 830.8 kB/s eta 0:00:02\n",
      "   ---------------------------- ----------- 3.6/5.0 MB 836.1 kB/s eta 0:00:02\n",
      "   ----------------------------- ---------- 3.6/5.0 MB 846.0 kB/s eta 0:00:02\n",
      "   ----------------------------- ---------- 3.7/5.0 MB 855.7 kB/s eta 0:00:02\n",
      "   ------------------------------ --------- 3.8/5.0 MB 866.0 kB/s eta 0:00:02\n",
      "   ------------------------------- -------- 3.9/5.0 MB 874.7 kB/s eta 0:00:02\n",
      "   ------------------------------- -------- 4.0/5.0 MB 886.1 kB/s eta 0:00:02\n",
      "   -------------------------------- ------- 4.1/5.0 MB 895.0 kB/s eta 0:00:02\n",
      "   --------------------------------- ------ 4.1/5.0 MB 899.2 kB/s eta 0:00:01\n",
      "   --------------------------------- ------ 4.2/5.0 MB 907.8 kB/s eta 0:00:01\n",
      "   ---------------------------------- ----- 4.3/5.0 MB 916.2 kB/s eta 0:00:01\n",
      "   ----------------------------------- ---- 4.4/5.0 MB 928.7 kB/s eta 0:00:01\n",
      "   ----------------------------------- ---- 4.5/5.0 MB 938.9 kB/s eta 0:00:01\n",
      "   ------------------------------------ --- 4.6/5.0 MB 949.9 kB/s eta 0:00:01\n",
      "   ------------------------------------- -- 4.7/5.0 MB 956.8 kB/s eta 0:00:01\n",
      "   ------------------------------------- -- 4.8/5.0 MB 966.4 kB/s eta 0:00:01\n",
      "   -------------------------------------- - 4.9/5.0 MB 977.9 kB/s eta 0:00:01\n",
      "   ---------------------------------------  4.9/5.0 MB 987.1 kB/s eta 0:00:01\n",
      "   ---------------------------------------- 5.0/5.0 MB 992.1 kB/s eta 0:00:00\n",
      "Downloading fastavro-1.9.3-cp311-cp311-win_amd64.whl (499 kB)\n",
      "   ---------------------------------------- 0.0/499.3 kB ? eta -:--:--\n",
      "   ------- -------------------------------- 92.2/499.3 kB 2.6 MB/s eta 0:00:01\n",
      "   ---------------- ----------------------- 204.8/499.3 kB 2.1 MB/s eta 0:00:01\n",
      "   ------------------------ --------------- 307.2/499.3 kB 2.1 MB/s eta 0:00:01\n",
      "   ------------------------------- -------- 399.4/499.3 kB 2.3 MB/s eta 0:00:01\n",
      "   ---------------------------------------- 499.3/499.3 kB 2.1 MB/s eta 0:00:00\n",
      "Downloading fasteners-0.19-py3-none-any.whl (18 kB)\n",
      "Downloading grpcio-1.60.0-cp311-cp311-win_amd64.whl (3.7 MB)\n",
      "   ---------------------------------------- 0.0/3.7 MB ? eta -:--:--\n",
      "    --------------------------------------- 0.1/3.7 MB 1.9 MB/s eta 0:00:02\n",
      "   - -------------------------------------- 0.2/3.7 MB 2.1 MB/s eta 0:00:02\n",
      "   -- ------------------------------------- 0.3/3.7 MB 2.1 MB/s eta 0:00:02\n",
      "   --- ------------------------------------ 0.4/3.7 MB 2.1 MB/s eta 0:00:02\n",
      "   ----- ---------------------------------- 0.5/3.7 MB 2.2 MB/s eta 0:00:02\n",
      "   ------ --------------------------------- 0.6/3.7 MB 2.1 MB/s eta 0:00:02\n",
      "   ------- -------------------------------- 0.7/3.7 MB 2.1 MB/s eta 0:00:02\n",
      "   ------- -------------------------------- 0.7/3.7 MB 2.0 MB/s eta 0:00:02\n",
      "   -------- ------------------------------- 0.8/3.7 MB 1.9 MB/s eta 0:00:02\n",
      "   --------- ------------------------------ 0.9/3.7 MB 2.0 MB/s eta 0:00:02\n",
      "   ---------- ----------------------------- 1.0/3.7 MB 2.0 MB/s eta 0:00:02\n",
      "   ----------- ---------------------------- 1.1/3.7 MB 1.9 MB/s eta 0:00:02\n",
      "   ------------ --------------------------- 1.2/3.7 MB 1.9 MB/s eta 0:00:02\n",
      "   ------------- -------------------------- 1.2/3.7 MB 1.9 MB/s eta 0:00:02\n",
      "   -------------- ------------------------- 1.3/3.7 MB 1.9 MB/s eta 0:00:02\n",
      "   -------------- ------------------------- 1.4/3.7 MB 1.8 MB/s eta 0:00:02\n",
      "   --------------- ------------------------ 1.4/3.7 MB 1.8 MB/s eta 0:00:02\n",
      "   --------------- ------------------------ 1.5/3.7 MB 1.8 MB/s eta 0:00:02\n",
      "   ---------------- ----------------------- 1.6/3.7 MB 1.8 MB/s eta 0:00:02\n",
      "   ----------------- ---------------------- 1.6/3.7 MB 1.8 MB/s eta 0:00:02\n",
      "   ------------------ --------------------- 1.8/3.7 MB 1.8 MB/s eta 0:00:02\n",
      "   -------------------- ------------------- 1.9/3.7 MB 1.8 MB/s eta 0:00:02\n",
      "   --------------------- ------------------ 2.0/3.7 MB 1.8 MB/s eta 0:00:01\n",
      "   ---------------------- ----------------- 2.1/3.7 MB 1.9 MB/s eta 0:00:01\n",
      "   ----------------------- ---------------- 2.2/3.7 MB 1.9 MB/s eta 0:00:01\n",
      "   ------------------------ --------------- 2.3/3.7 MB 1.9 MB/s eta 0:00:01\n",
      "   ------------------------- -------------- 2.4/3.7 MB 1.9 MB/s eta 0:00:01\n",
      "   -------------------------- ------------- 2.5/3.7 MB 1.9 MB/s eta 0:00:01\n",
      "   --------------------------- ------------ 2.6/3.7 MB 1.9 MB/s eta 0:00:01\n",
      "   ---------------------------- ----------- 2.7/3.7 MB 1.9 MB/s eta 0:00:01\n",
      "   ------------------------------ --------- 2.8/3.7 MB 1.9 MB/s eta 0:00:01\n",
      "   ------------------------------- -------- 2.9/3.7 MB 2.0 MB/s eta 0:00:01\n",
      "   -------------------------------- ------- 3.0/3.7 MB 2.0 MB/s eta 0:00:01\n",
      "   --------------------------------- ------ 3.1/3.7 MB 2.0 MB/s eta 0:00:01\n",
      "   ----------------------------------- ---- 3.2/3.7 MB 2.0 MB/s eta 0:00:01\n",
      "   ------------------------------------ --- 3.4/3.7 MB 2.0 MB/s eta 0:00:01\n",
      "   ------------------------------------- -- 3.5/3.7 MB 2.0 MB/s eta 0:00:01\n",
      "   -------------------------------------- - 3.6/3.7 MB 2.0 MB/s eta 0:00:01\n",
      "   ---------------------------------------  3.7/3.7 MB 2.0 MB/s eta 0:00:01\n",
      "   ---------------------------------------- 3.7/3.7 MB 2.0 MB/s eta 0:00:00\n",
      "Downloading jsonpickle-3.0.2-py3-none-any.whl (40 kB)\n",
      "   ---------------------------------------- 0.0/40.7 kB ? eta -:--:--\n",
      "   ---------------------------------------- 40.7/40.7 kB 2.0 MB/s eta 0:00:00\n",
      "Downloading numpy-1.24.4-cp311-cp311-win_amd64.whl (14.8 MB)\n",
      "   ---------------------------------------- 0.0/14.8 MB ? eta -:--:--\n",
      "   ---------------------------------------- 0.1/14.8 MB 2.4 MB/s eta 0:00:07\n",
      "    --------------------------------------- 0.3/14.8 MB 2.6 MB/s eta 0:00:06\n",
      "   - -------------------------------------- 0.4/14.8 MB 2.7 MB/s eta 0:00:06\n",
      "   - -------------------------------------- 0.5/14.8 MB 2.6 MB/s eta 0:00:06\n",
      "   - -------------------------------------- 0.6/14.8 MB 2.7 MB/s eta 0:00:06\n",
      "   -- ------------------------------------- 0.8/14.8 MB 2.7 MB/s eta 0:00:06\n",
      "   -- ------------------------------------- 0.9/14.8 MB 2.7 MB/s eta 0:00:06\n",
      "   -- ------------------------------------- 1.0/14.8 MB 2.7 MB/s eta 0:00:06\n",
      "   --- ------------------------------------ 1.1/14.8 MB 2.7 MB/s eta 0:00:06\n",
      "   --- ------------------------------------ 1.3/14.8 MB 2.8 MB/s eta 0:00:05\n",
      "   --- ------------------------------------ 1.4/14.8 MB 2.7 MB/s eta 0:00:05\n",
      "   ---- ----------------------------------- 1.5/14.8 MB 2.7 MB/s eta 0:00:05\n",
      "   ---- ----------------------------------- 1.7/14.8 MB 2.8 MB/s eta 0:00:05\n",
      "   ---- ----------------------------------- 1.8/14.8 MB 2.7 MB/s eta 0:00:05\n",
      "   ----- ---------------------------------- 1.9/14.8 MB 2.7 MB/s eta 0:00:05\n",
      "   ----- ---------------------------------- 2.0/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   ----- ---------------------------------- 2.1/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   ----- ---------------------------------- 2.2/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   ------ --------------------------------- 2.3/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   ------ --------------------------------- 2.5/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   ------ --------------------------------- 2.6/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   ------- -------------------------------- 2.7/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   ------- -------------------------------- 2.8/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   ------- -------------------------------- 2.9/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   -------- ------------------------------- 3.1/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   -------- ------------------------------- 3.2/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   -------- ------------------------------- 3.3/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   --------- ------------------------------ 3.3/14.8 MB 2.6 MB/s eta 0:00:05\n",
      "   --------- ------------------------------ 3.5/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   --------- ------------------------------ 3.6/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   --------- ------------------------------ 3.7/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ---------- ----------------------------- 3.8/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ---------- ----------------------------- 3.9/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ---------- ----------------------------- 4.0/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ---------- ----------------------------- 4.1/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ----------- ---------------------------- 4.2/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ----------- ---------------------------- 4.3/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ------------ --------------------------- 4.5/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ------------ --------------------------- 4.6/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ------------ --------------------------- 4.7/14.8 MB 2.5 MB/s eta 0:00:05\n",
      "   ------------- -------------------------- 4.9/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ------------- -------------------------- 5.0/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ------------- -------------------------- 5.2/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   -------------- ------------------------- 5.3/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   -------------- ------------------------- 5.5/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   -------------- ------------------------- 5.6/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   --------------- ------------------------ 5.7/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   --------------- ------------------------ 5.8/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   --------------- ------------------------ 5.9/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   ---------------- ----------------------- 6.0/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   ---------------- ----------------------- 6.1/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ---------------- ----------------------- 6.2/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ---------------- ----------------------- 6.3/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ----------------- ---------------------- 6.4/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ----------------- ---------------------- 6.5/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ----------------- ---------------------- 6.6/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ------------------ --------------------- 6.7/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ------------------ --------------------- 6.9/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ------------------ --------------------- 7.0/14.8 MB 2.6 MB/s eta 0:00:04\n",
      "   ------------------- -------------------- 7.1/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ------------------- -------------------- 7.2/14.8 MB 2.5 MB/s eta 0:00:04\n",
      "   ------------------- -------------------- 7.4/14.8 MB 2.5 MB/s eta 0:00:03\n",
      "   -------------------- ------------------- 7.5/14.8 MB 2.5 MB/s eta 0:00:03\n",
      "   -------------------- ------------------- 7.6/14.8 MB 2.5 MB/s eta 0:00:03\n",
      "   -------------------- ------------------- 7.7/14.8 MB 2.5 MB/s eta 0:00:03\n",
      "   --------------------- ------------------ 7.9/14.8 MB 2.5 MB/s eta 0:00:03\n",
      "   --------------------- ------------------ 8.0/14.8 MB 2.6 MB/s eta 0:00:03\n",
      "   ---------------------- ----------------- 8.2/14.8 MB 2.6 MB/s eta 0:00:03\n",
      "   ---------------------- ----------------- 8.3/14.8 MB 2.6 MB/s eta 0:00:03\n",
      "   ---------------------- ----------------- 8.5/14.8 MB 2.6 MB/s eta 0:00:03\n",
      "   ----------------------- ---------------- 8.7/14.8 MB 2.6 MB/s eta 0:00:03\n",
      "   ----------------------- ---------------- 8.9/14.8 MB 2.6 MB/s eta 0:00:03\n",
      "   ------------------------ --------------- 9.0/14.8 MB 2.7 MB/s eta 0:00:03\n",
      "   ------------------------ --------------- 9.2/14.8 MB 2.7 MB/s eta 0:00:03\n",
      "   ------------------------- -------------- 9.4/14.8 MB 2.7 MB/s eta 0:00:03\n",
      "   ------------------------- -------------- 9.6/14.8 MB 2.7 MB/s eta 0:00:02\n",
      "   -------------------------- ------------- 9.8/14.8 MB 2.7 MB/s eta 0:00:02\n",
      "   -------------------------- ------------- 9.9/14.8 MB 2.7 MB/s eta 0:00:02\n",
      "   --------------------------- ------------ 10.1/14.8 MB 2.7 MB/s eta 0:00:02\n",
      "   --------------------------- ------------ 10.2/14.8 MB 2.7 MB/s eta 0:00:02\n",
      "   ---------------------------- ----------- 10.4/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ---------------------------- ----------- 10.6/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ---------------------------- ----------- 10.7/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ----------------------------- ---------- 10.9/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ----------------------------- ---------- 11.0/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ------------------------------ --------- 11.2/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ------------------------------ --------- 11.4/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ------------------------------- -------- 11.6/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ------------------------------- -------- 11.7/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   ------------------------------- -------- 11.8/14.8 MB 2.8 MB/s eta 0:00:02\n",
      "   -------------------------------- ------- 12.0/14.8 MB 2.8 MB/s eta 0:00:01\n",
      "   -------------------------------- ------- 12.2/14.8 MB 2.8 MB/s eta 0:00:01\n",
      "   --------------------------------- ------ 12.3/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   --------------------------------- ------ 12.4/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   --------------------------------- ------ 12.6/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ---------------------------------- ----- 12.7/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ---------------------------------- ----- 12.8/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ---------------------------------- ----- 13.0/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ----------------------------------- ---- 13.1/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ----------------------------------- ---- 13.2/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ----------------------------------- ---- 13.3/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ------------------------------------ --- 13.5/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ------------------------------------ --- 13.6/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ------------------------------------- -- 13.8/14.8 MB 2.9 MB/s eta 0:00:01\n",
      "   ------------------------------------- -- 14.0/14.8 MB 3.0 MB/s eta 0:00:01\n",
      "   -------------------------------------- - 14.1/14.8 MB 3.0 MB/s eta 0:00:01\n",
      "   -------------------------------------- - 14.3/14.8 MB 3.0 MB/s eta 0:00:01\n",
      "   -------------------------------------- - 14.4/14.8 MB 3.0 MB/s eta 0:00:01\n",
      "   ---------------------------------------  14.6/14.8 MB 3.1 MB/s eta 0:00:01\n",
      "   ---------------------------------------  14.8/14.8 MB 3.1 MB/s eta 0:00:01\n",
      "   ---------------------------------------- 14.8/14.8 MB 3.1 MB/s eta 0:00:00\n",
      "Downloading orjson-3.9.10-cp311-none-win_amd64.whl (135 kB)\n",
      "   ---------------------------------------- 0.0/135.0 kB ? eta -:--:--\n",
      "   ---------------------------------------- 135.0/135.0 kB 4.0 MB/s eta 0:00:00\n",
      "Downloading proto_plus-1.23.0-py3-none-any.whl (48 kB)\n",
      "   ---------------------------------------- 0.0/48.8 kB ? eta -:--:--\n",
      "   ---------------------------------------- 48.8/48.8 kB 2.6 MB/s eta 0:00:00\n",
      "Downloading pymongo-4.6.1-cp311-cp311-win_amd64.whl (472 kB)\n",
      "   ---------------------------------------- 0.0/472.7 kB ? eta -:--:--\n",
      "   ----------------- ---------------------- 204.8/472.7 kB 6.3 MB/s eta 0:00:01\n",
      "   ----------------------------------- ---- 419.8/472.7 kB 5.3 MB/s eta 0:00:01\n",
      "   ---------------------------------------- 472.7/472.7 kB 4.2 MB/s eta 0:00:00\n",
      "Downloading zstandard-0.22.0-cp311-cp311-win_amd64.whl (511 kB)\n",
      "   ---------------------------------------- 0.0/511.6 kB ? eta -:--:--\n",
      "   ---------------- ----------------------- 204.8/511.6 kB 6.3 MB/s eta 0:00:01\n",
      "   ---------------------------------- ----- 440.3/511.6 kB 4.6 MB/s eta 0:00:01\n",
      "   ---------------------------------------- 511.6/511.6 kB 4.6 MB/s eta 0:00:00\n",
      "Downloading dnspython-2.4.2-py3-none-any.whl (300 kB)\n",
      "   ---------------------------------------- 0.0/300.4 kB ? eta -:--:--\n",
      "   ------------------------------ --------- 225.3/300.4 kB 4.6 MB/s eta 0:00:01\n",
      "   ---------------------------------------- 300.4/300.4 kB 4.7 MB/s eta 0:00:00\n",
      "Downloading pyparsing-3.1.1-py3-none-any.whl (103 kB)\n",
      "   ---------------------------------------- 0.0/103.1 kB ? eta -:--:--\n",
      "   ---------------------------------------- 103.1/103.1 kB 5.8 MB/s eta 0:00:00\n",
      "Downloading tzlocal-5.2-py3-none-any.whl (17 kB)\n",
      "Building wheels for collected packages: crcmod, dill, hdfs, pyjsparser, docopt\n",
      "  Building wheel for crcmod (setup.py): started\n",
      "  Building wheel for crcmod (setup.py): finished with status 'done'\n",
      "  Created wheel for crcmod: filename=crcmod-1.7-cp311-cp311-win_amd64.whl size=25183 sha256=f68444f492d3c618f8be030d024b06f14352c85a4e8e6211e2fe16c1e506a926\n",
      "  Stored in directory: c:\\users\\robbi\\appdata\\local\\pip\\cache\\wheels\\23\\94\\7a\\8cb7d14597e6395ce969933f01aed9ea8fa5f5b4d4c8a61e99\n",
      "  Building wheel for dill (setup.py): started\n",
      "  Building wheel for dill (setup.py): finished with status 'done'\n",
      "  Created wheel for dill: filename=dill-0.3.1.1-py3-none-any.whl size=78667 sha256=8aafca05246b64573283bb698e8da850deb8a4483eac94eeeef00983de12c3aa\n",
      "  Stored in directory: c:\\users\\robbi\\appdata\\local\\pip\\cache\\wheels\\01\\60\\80\\1622338bcecce31a5664ef01c203cc5a7b09f59588d9c07376\n",
      "  Building wheel for hdfs (setup.py): started\n",
      "  Building wheel for hdfs (setup.py): finished with status 'done'\n",
      "  Created wheel for hdfs: filename=hdfs-2.7.3-py3-none-any.whl size=34350 sha256=f07c418d0213731b8353ce97a6ec99b975f8f034baa2a49184d81ad9a872e708\n",
      "  Stored in directory: c:\\users\\robbi\\appdata\\local\\pip\\cache\\wheels\\b9\\1d\\dc\\eb0833be25464c359903d356c4204721c6a672c26ff164cdc3\n",
      "  Building wheel for pyjsparser (setup.py): started\n",
      "  Building wheel for pyjsparser (setup.py): finished with status 'done'\n",
      "  Created wheel for pyjsparser: filename=pyjsparser-2.7.1-py3-none-any.whl size=25988 sha256=2e4af5837743c6427cd76452ec1f5dec2b763f813c34c025644dc16b20edfa90\n",
      "  Stored in directory: c:\\users\\robbi\\appdata\\local\\pip\\cache\\wheels\\a5\\9a\\30\\1003e89ab4555b81840ca46d361bf184f1e6ad880cae3b62a9\n",
      "  Building wheel for docopt (setup.py): started\n",
      "  Building wheel for docopt (setup.py): finished with status 'done'\n",
      "  Created wheel for docopt: filename=docopt-0.6.2-py2.py3-none-any.whl size=13773 sha256=ce9af9da1becdbf2b82e79817de64a830bdaf33f6cd58ebb6bb953894ace0fc2\n",
      "  Stored in directory: c:\\users\\robbi\\appdata\\local\\pip\\cache\\wheels\\1a\\b0\\8c\\4b75c4116c31f83c8f9f047231251e13cc74481cca4a78a9ce\n",
      "Successfully built crcmod dill hdfs pyjsparser docopt\n",
      "Installing collected packages: pyjsparser, docopt, crcmod, zstandard, tzlocal, pyparsing, proto-plus, orjson, objsize, numpy, jsonpickle, grpcio, fasteners, fastavro, dnspython, dill, cloudpickle, pymongo, pydot, js2py, httplib2, hdfs, apache_beam\n",
      "  Attempting uninstall: numpy\n",
      "    Found existing installation: numpy 1.26.2\n",
      "    Uninstalling numpy-1.26.2:\n",
      "      Successfully uninstalled numpy-1.26.2\n",
      "  Attempting uninstall: dill\n",
      "    Found existing installation: dill 0.3.6\n",
      "    Uninstalling dill-0.3.6:\n",
      "      Successfully uninstalled dill-0.3.6\n",
      "  Attempting uninstall: cloudpickle\n",
      "    Found existing installation: cloudpickle 3.0.0\n",
      "    Uninstalling cloudpickle-3.0.0:\n",
      "      Successfully uninstalled cloudpickle-3.0.0\n",
      "Successfully installed apache_beam-2.53.0 cloudpickle-2.2.1 crcmod-1.7 dill-0.3.1.1 dnspython-2.4.2 docopt-0.6.2 fastavro-1.9.3 fasteners-0.19 grpcio-1.60.0 hdfs-2.7.3 httplib2-0.22.0 js2py-0.74 jsonpickle-3.0.2 numpy-1.24.4 objsize-0.6.1 orjson-3.9.10 proto-plus-1.23.0 pydot-1.4.2 pyjsparser-2.7.1 pymongo-4.6.1 pyparsing-3.1.1 tzlocal-5.2 zstandard-0.22.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "multiprocess 0.70.14 requires dill>=0.3.6, but you have dill 0.3.1.1 which is incompatible.\n"
     ]
    }
   ],
   "source": [
    "!pip install apache_beam\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5, Loss: 10.499467606929352\n",
      "Epoch 2/5, Loss: 10.49840387818077\n",
      "Epoch 3/5, Loss: 10.496170871612177\n",
      "Epoch 4/5, Loss: 10.49417147156456\n",
      "Epoch 5/5, Loss: 10.492957106852094\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.optim import AdamW\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "from torch.utils.data import random_split\n",
    "from transformers import BertTokenizer\n",
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "class LORALayer(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, rank, alpha=1):\n",
    "        super(LORALayer, self).__init__()\n",
    "        self.rank = rank\n",
    "        self.alpha = alpha\n",
    "\n",
    "        # Original weight and bias of the linear layer\n",
    "        self.weight = nn.Parameter(torch.Tensor(output_dim, input_dim))\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "        self.bias = nn.Parameter(torch.Tensor(output_dim))\n",
    "\n",
    "        # LORA specific parameters\n",
    "        self.A = nn.Parameter(torch.Tensor(input_dim, rank))\n",
    "        self.B = nn.Parameter(torch.Tensor(rank, output_dim))\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
    "        nn.init.zeros_(self.bias)\n",
    "        nn.init.normal_(self.A, 0, 0.02)\n",
    "        nn.init.normal_(self.B, 0, 0.02)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #print(\"LORALayer Input Shape:\", x.shape)\n",
    "        \n",
    "        original_size = x.size()\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "        x_flattened = x.reshape(-1, original_size[-1])\n",
    "\n",
    "        # Compute lora_adjustment for each input in the batch\n",
    "        lora_adjustment = self.alpha * (x_flattened @ self.A) @ self.B\n",
    "        lora_adjustment = lora_adjustment.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        #print(\"Adjusted lora_adjustment Shape:\", lora_adjustment.shape)\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "        \n",
    "        # Apply linear transformation to x_flattened\n",
    "        x_transformed = nn.functional.linear(x_flattened, self.weight, self.bias)\n",
    "        x_transformed = x_transformed.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        # Add lora_adjustment to the transformed x\n",
    "        x = x_transformed + lora_adjustment\n",
    "        #print(\"LORALayer Output Shape:\", x.shape)\n",
    "\n",
    "        return x\n",
    "\n",
    "class QLORALayer(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, rank, alpha=1, quantization_bits=4):\n",
    "        super(QLORALayer, self).__init__()\n",
    "        self.rank = rank\n",
    "        self.alpha = alpha\n",
    "        self.quantization_bits = quantization_bits\n",
    "\n",
    "        # Original weight and bias\n",
    "        self.weight = nn.Parameter(torch.Tensor(output_dim, input_dim))\n",
    "        self.bias = nn.Parameter(torch.Tensor(output_dim))\n",
    "\n",
    "        # QLORA specific parameters\n",
    "        self.A = nn.Parameter(torch.Tensor(input_dim, rank))\n",
    "        self.B = nn.Parameter(torch.Tensor(rank, output_dim))\n",
    "\n",
    "        self.dropout = nn.Dropout(0.1)\n",
    "        self.layer_norm = nn.LayerNorm(output_dim)\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
    "        nn.init.zeros_(self.bias)\n",
    "        nn.init.normal_(self.A, 0, 0.02)\n",
    "        nn.init.normal_(self.B, 0, 0.02)\n",
    "\n",
    "    def quantize(self, x, num_bits):\n",
    "        # Implement a simple quantization method\n",
    "        scale = x.abs().max()\n",
    "        x_quantized = torch.round(x / scale * (2**num_bits - 1))\n",
    "        return x_quantized, scale\n",
    "\n",
    "    def forward(self, x):\n",
    "        #print(\"QLORALayer Input Shape:\", x.shape)\n",
    "        original_size = x.size()\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "        x_flattened = x.reshape(-1, original_size[-1])\n",
    "\n",
    "        A_quantized, scale_A = self.quantize(self.A, self.quantization_bits)\n",
    "        B_quantized, scale_B = self.quantize(self.B, self.quantization_bits)\n",
    "\n",
    "        # Compute lora_adjustment for each input in the batch\n",
    "        lora_adjustment = self.alpha * (x_flattened @ (A_quantized / scale_A)) @ (B_quantized / scale_B)\n",
    "        lora_adjustment = lora_adjustment.reshape(batch_size, seq_len, -1)\n",
    "        lora_adjustment = self.dropout(lora_adjustment)\n",
    "        #print(\"Adjusted lora_adjustment Shape:\", lora_adjustment.shape)\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "\n",
    "        # Apply linear transformation to x_flattened\n",
    "        x_transformed = nn.functional.linear(x_flattened, self.weight, self.bias)\n",
    "        x_transformed = x_transformed.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        # Add lora_adjustment to the transformed x\n",
    "        x = x_transformed + lora_adjustment\n",
    "        x = self.layer_norm(x)\n",
    "\n",
    "        #print(\"QLORALayer Output Shape:\", x.shape)\n",
    "\n",
    "        return x\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, embed_size, heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.embed_size = embed_size\n",
    "        self.heads = heads\n",
    "        self.head_dim = embed_size // heads\n",
    "\n",
    "        assert (\n",
    "            self.head_dim * heads == embed_size\n",
    "        ), \"Embedding size needs to be divisible by heads\"\n",
    "\n",
    "        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)\n",
    "\n",
    "    def forward(self, values, keys, query, mask):\n",
    "        N = query.shape[0]\n",
    "        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]\n",
    "\n",
    "        # Split the embedding into self.heads different pieces\n",
    "        values = values.reshape(N, value_len, self.heads, self.head_dim)\n",
    "        keys = keys.reshape(N, key_len, self.heads, self.head_dim)\n",
    "        queries = query.reshape(N, query_len, self.heads, self.head_dim)\n",
    "\n",
    "        values = self.values(values)\n",
    "        keys = self.keys(keys)\n",
    "        queries = self.queries(queries)\n",
    "\n",
    "        # Einsum does the matrix multiplication for query*keys for each training example\n",
    "        # with every other training example, then sum it up\n",
    "        attention = torch.einsum(\"nqhd,nkhd->nhqk\", [queries, keys])\n",
    "\n",
    "        if mask is not None:\n",
    "            attention = attention.masked_fill(mask == 0, float(\"-1e20\"))\n",
    "\n",
    "        attention = torch.softmax(attention / (self.embed_size ** (1 / 2)), dim=3)\n",
    "\n",
    "        out = torch.einsum(\"nhql,nlhd->nqhd\", [attention, values]).reshape(\n",
    "            N, query_len, self.heads * self.head_dim\n",
    "        )\n",
    "\n",
    "        out = self.fc_out(out)\n",
    "        return out\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, embed_size, heads, dropout, forward_expansion, rank):\n",
    "        super(TransformerBlock, self).__init__()\n",
    "        self.attention = MultiHeadAttention(embed_size, heads)\n",
    "        self.norm1 = nn.LayerNorm(embed_size)\n",
    "        self.norm2 = nn.LayerNorm(embed_size)\n",
    "\n",
    "        self.feed_forward = nn.Sequential(\n",
    "            LORALayer(embed_size, forward_expansion * embed_size, rank),\n",
    "            nn.ReLU(),\n",
    "            LORALayer(forward_expansion * embed_size, embed_size, rank),\n",
    "        )\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, value, key, query, mask):\n",
    "        attention = self.attention(value, key, query, mask)\n",
    "        x = self.dropout(self.norm1(attention + query))\n",
    "        forward = self.feed_forward(x)\n",
    "        out = self.dropout(self.norm2(forward + x))\n",
    "        return out\n",
    "\n",
    "class LanguageModelDecoder(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_size, num_layers, heads, forward_expansion, dropout, max_length, rank):\n",
    "        super(LanguageModelDecoder, self).__init__()\n",
    "        self.word_embedding = nn.Embedding(vocab_size, embed_size)\n",
    "        self.position_embedding = nn.Embedding(max_length, embed_size)\n",
    "\n",
    "        # Adding BatchNorm layers\n",
    "        self.bn1 = nn.BatchNorm1d(embed_size)\n",
    "        self.bn2 = nn.BatchNorm1d(embed_size)\n",
    "\n",
    "        self.layers = nn.ModuleList(\n",
    "            [\n",
    "                TransformerBlock(embed_size, heads, dropout, forward_expansion, rank)\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        # QLORA layers\n",
    "        self.qlora_feed_forward = nn.Sequential(\n",
    "            QLORALayer(embed_size, forward_expansion * embed_size, rank),\n",
    "            nn.ReLU(),\n",
    "            QLORALayer(forward_expansion * embed_size, embed_size, rank),\n",
    "        )\n",
    "        self.use_qlora = False  # Flag to toggle QLORA\n",
    "\n",
    "        self.fc_out = nn.Linear(embed_size, vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, trg_mask):\n",
    "        N, seq_length = x.shape\n",
    "        positions = torch.arange(0, seq_length).expand(N, seq_length).to(x.device)\n",
    "        x = self.dropout(self.word_embedding(x) + self.position_embedding(positions))\n",
    "\n",
    "        # Transpose for BatchNorm, apply batch normalization, and then transpose back\n",
    "        x = x.transpose(1, 2)\n",
    "        x = self.bn1(x)\n",
    "        x = x.transpose(1, 2)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, x, x, trg_mask)\n",
    "            if self.use_qlora:\n",
    "                x = self.qlora_feed_forward(x)\n",
    "\n",
    "        # Transpose for BatchNorm, apply batch normalization, and then transpose back\n",
    "        x = x.transpose(1, 2)\n",
    "        x = self.bn2(x)\n",
    "        x = x.transpose(1, 2)\n",
    "\n",
    "        out = self.fc_out(x)\n",
    "        return out\n",
    "\n",
    "    def toggle_qlora(self, use_qlora: bool):\n",
    "        self.use_qlora = use_qlora\n",
    "\n",
    "\n",
    " \n",
    "\n",
    "class LanguageModelTransformer(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_size=256, num_layers=6, forward_expansion=4, heads=8, dropout=0, max_length=100, rank=16):\n",
    "        super(LanguageModelTransformer, self).__init__()\n",
    "\n",
    "        self.decoder = LanguageModelDecoder(\n",
    "            vocab_size,\n",
    "            embed_size,\n",
    "            num_layers,\n",
    "            heads,\n",
    "            forward_expansion,\n",
    "            dropout,\n",
    "            max_length,\n",
    "            rank,\n",
    "        )\n",
    "\n",
    "    def forward(self, trg):\n",
    "        trg_mask = self.make_trg_mask(trg)\n",
    "        out = self.decoder(trg, trg_mask)\n",
    "        return out\n",
    "\n",
    "    def make_trg_mask(self, trg):\n",
    "        N, trg_len = trg.shape\n",
    "        trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(\n",
    "            N, 1, trg_len, trg_len\n",
    "        ).to(trg.device)\n",
    "\n",
    "        return trg_mask\n",
    "\n",
    "# Define vocabulary size and dummy data parameters\n",
    "NUM_WORDS = 1000  # Example vocabulary size\n",
    "sequence_length = 30  # Sequence length for the LanguageDataset\n",
    "dummy_data_size = 1000  # Total number of tokens in the dummy dataset\n",
    "\n",
    "\n",
    "\n",
    "# Load dataset\n",
    "dataset = load_dataset('wikipedia', '20220301.simple')\n",
    "\n",
    "# Initialize tokenizer\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "vocab_size = tokenizer.vocab_size\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    # Tokenize the text\n",
    "    tokenized_output = tokenizer(examples['text'], padding='max_length', truncation=True, max_length=sequence_length)\n",
    "    \n",
    "    # Shift input_ids to create labels and truncate the last token\n",
    "    labels = [seq[1:] + [tokenizer.pad_token_id] for seq in tokenized_output['input_ids']]\n",
    "    tokenized_output['labels'] = labels\n",
    "    \n",
    "    return tokenized_output\n",
    "\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)\n",
    "\n",
    "# Set format for PyTorch\n",
    "tokenized_datasets.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])\n",
    "\n",
    "train_loader = DataLoader(tokenized_datasets['train'], batch_size=64, shuffle=True)\n",
    "\n",
    "# Define device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Create the model instance\n",
    "model = LanguageModelTransformer(\n",
    "    vocab_size=vocab_size,  # Use the vocab size from the tokenizer\n",
    "    embed_size=256,\n",
    "    num_layers=6,\n",
    "    forward_expansion=4,\n",
    "    heads=8,\n",
    "    dropout=0,\n",
    "    max_length=100,\n",
    "    rank=16\n",
    ").to(device)\n",
    "\n",
    "\n",
    "\n",
    "# Enable QLORA during training\n",
    "model.decoder.toggle_qlora(True)\n",
    "\n",
    "# Training loop\n",
    "# Assuming model is an instance of LanguageModelTransformer\n",
    "# Define loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = AdamW(model.parameters(), lr=1e-8, weight_decay=1e-4)\n",
    "scheduler = StepLR(optimizer, step_size=4, gamma=0.98)\n",
    "num_epochs = 5\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    model.decoder.toggle_qlora(True)\n",
    "    total_loss = 0\n",
    "\n",
    "    for batch in train_loader:\n",
    "        inputs = batch['input_ids'].to(device)\n",
    "        targets = batch['labels'].to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        #print(\"Output shape:\", outputs.shape)\n",
    "        loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))\n",
    "        \n",
    "        # Check for NaN in loss\n",
    "        if math.isnan(loss.item()):\n",
    "            print(\"Encountered NaN loss, stopping training\")\n",
    "            break\n",
    "\n",
    "        loss.backward()\n",
    "        \n",
    "        # Gradient clipping\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)\n",
    "\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "\n",
    "    scheduler.step()\n",
    "\n",
    "    # Check for NaN in total_loss\n",
    "    if math.isnan(total_loss):\n",
    "        print(f\"Epoch {epoch+1}/{num_epochs} stopped due to NaN loss\")\n",
    "        break\n",
    "    else:\n",
    "        print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss / len(train_loader)}\")\n",
    "\n",
    "    model.decoder.toggle_qlora(False)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# v3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc703bf93d2f4a0eb0316345b622dac7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/205328 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5, Loss: 10.485922775375542\n",
      "Epoch 2/5, Loss: 10.48466368396173\n",
      "Epoch 3/5, Loss: 10.483034437042265\n",
      "Epoch 4/5, Loss: 10.481329586246593\n",
      "Epoch 5/5, Loss: 10.47995846856231\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.optim import AdamW\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "from torch.utils.data import random_split\n",
    "from transformers import BertTokenizer\n",
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "class LORALayer(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, rank, alpha=1):\n",
    "        super(LORALayer, self).__init__()\n",
    "        self.rank = rank\n",
    "        self.alpha = alpha\n",
    "\n",
    "        # Original weight and bias of the linear layer\n",
    "        self.weight = nn.Parameter(torch.Tensor(output_dim, input_dim))\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "        self.bias = nn.Parameter(torch.Tensor(output_dim))\n",
    "\n",
    "        # LORA specific parameters\n",
    "        self.A = nn.Parameter(torch.Tensor(input_dim, rank))\n",
    "        self.B = nn.Parameter(torch.Tensor(rank, output_dim))\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
    "        nn.init.zeros_(self.bias)\n",
    "        nn.init.normal_(self.A, 0, 0.02)\n",
    "        nn.init.normal_(self.B, 0, 0.02)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #print(\"LORALayer Input Shape:\", x.shape)\n",
    "        \n",
    "        original_size = x.size()\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "        x_flattened = x.reshape(-1, original_size[-1])\n",
    "\n",
    "        # Compute lora_adjustment for each input in the batch\n",
    "        lora_adjustment = self.alpha * (x_flattened @ self.A) @ self.B\n",
    "        lora_adjustment = lora_adjustment.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        #print(\"Adjusted lora_adjustment Shape:\", lora_adjustment.shape)\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "        \n",
    "        # Apply linear transformation to x_flattened\n",
    "        x_transformed = nn.functional.linear(x_flattened, self.weight, self.bias)\n",
    "        x_transformed = x_transformed.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        # Add lora_adjustment to the transformed x\n",
    "        x = x_transformed + lora_adjustment\n",
    "        #print(\"LORALayer Output Shape:\", x.shape)\n",
    "\n",
    "        return x\n",
    "\n",
    "class QLORALayer(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, rank, alpha=1, quantization_bits=8):\n",
    "        super(QLORALayer, self).__init__()\n",
    "        self.rank = rank\n",
    "        self.alpha = alpha\n",
    "        self.quantization_bits = quantization_bits\n",
    "\n",
    "        # Original weight and bias\n",
    "        self.weight = nn.Parameter(torch.Tensor(output_dim, input_dim))\n",
    "        self.bias = nn.Parameter(torch.Tensor(output_dim))\n",
    "\n",
    "        # QLORA specific parameters\n",
    "        self.A = nn.Parameter(torch.Tensor(input_dim, rank))\n",
    "        self.B = nn.Parameter(torch.Tensor(rank, output_dim))\n",
    "\n",
    "        self.dropout = nn.Dropout(0.1)\n",
    "        self.layer_norm = nn.LayerNorm(output_dim)\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
    "        nn.init.zeros_(self.bias)\n",
    "        nn.init.normal_(self.A, 0, 0.02)\n",
    "        nn.init.normal_(self.B, 0, 0.02)\n",
    "\n",
    "    def quantize(self, x, num_bits):\n",
    "        # Implement a simple quantization method\n",
    "        scale = x.abs().max()\n",
    "        x_quantized = torch.round(x / scale * (2**num_bits - 1))\n",
    "        return x_quantized, scale\n",
    "\n",
    "    def forward(self, x):\n",
    "        #print(\"QLORALayer Input Shape:\", x.shape)\n",
    "        original_size = x.size()\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "        x_flattened = x.reshape(-1, original_size[-1])\n",
    "\n",
    "        A_quantized, scale_A = self.quantize(self.A, self.quantization_bits)\n",
    "        B_quantized, scale_B = self.quantize(self.B, self.quantization_bits)\n",
    "\n",
    "        # Compute lora_adjustment for each input in the batch\n",
    "        lora_adjustment = self.alpha * (x_flattened @ (A_quantized / scale_A)) @ (B_quantized / scale_B)\n",
    "        lora_adjustment = lora_adjustment.reshape(batch_size, seq_len, -1)\n",
    "        lora_adjustment = self.dropout(lora_adjustment)\n",
    "        #print(\"Adjusted lora_adjustment Shape:\", lora_adjustment.shape)\n",
    "        #print(\"self.weight Shape:\", self.weight.shape)\n",
    "\n",
    "        # Apply linear transformation to x_flattened\n",
    "        x_transformed = nn.functional.linear(x_flattened, self.weight, self.bias)\n",
    "        x_transformed = x_transformed.reshape(batch_size, seq_len, -1)\n",
    "\n",
    "        # Add lora_adjustment to the transformed x\n",
    "        x = x_transformed + lora_adjustment\n",
    "        x = self.layer_norm(x)\n",
    "\n",
    "        #print(\"QLORALayer Output Shape:\", x.shape)\n",
    "\n",
    "        return x\n",
    "    \n",
    "    def update_alpha(self, new_alpha):\n",
    "        \"\"\"\n",
    "        Update the alpha scaling factor.\n",
    "        \"\"\"\n",
    "        self.alpha = new_alpha\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, embed_size, heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.embed_size = embed_size\n",
    "        self.heads = heads\n",
    "        self.head_dim = embed_size // heads\n",
    "\n",
    "        assert (\n",
    "            self.head_dim * heads == embed_size\n",
    "        ), \"Embedding size needs to be divisible by heads\"\n",
    "\n",
    "        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)\n",
    "        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)\n",
    "\n",
    "    def forward(self, values, keys, query, mask):\n",
    "        N = query.shape[0]\n",
    "        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]\n",
    "\n",
    "        # Split the embedding into self.heads different pieces\n",
    "        values = values.reshape(N, value_len, self.heads, self.head_dim)\n",
    "        keys = keys.reshape(N, key_len, self.heads, self.head_dim)\n",
    "        queries = query.reshape(N, query_len, self.heads, self.head_dim)\n",
    "\n",
    "        values = self.values(values)\n",
    "        keys = self.keys(keys)\n",
    "        queries = self.queries(queries)\n",
    "\n",
    "        # Einsum does the matrix multiplication for query*keys for each training example\n",
    "        # with every other training example, then sum it up\n",
    "        attention = torch.einsum(\"nqhd,nkhd->nhqk\", [queries, keys])\n",
    "\n",
    "        if mask is not None:\n",
    "            attention = attention.masked_fill(mask == 0, float(\"-1e20\"))\n",
    "\n",
    "        attention = torch.softmax(attention / (self.embed_size ** (1 / 2)), dim=3)\n",
    "\n",
    "        out = torch.einsum(\"nhql,nlhd->nqhd\", [attention, values]).reshape(\n",
    "            N, query_len, self.heads * self.head_dim\n",
    "        )\n",
    "\n",
    "        out = self.fc_out(out)\n",
    "        return out\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, embed_size, heads, dropout, forward_expansion, rank):\n",
    "        super(TransformerBlock, self).__init__()\n",
    "        self.attention = MultiHeadAttention(embed_size, heads)\n",
    "        self.norm1 = nn.LayerNorm(embed_size)\n",
    "        self.norm2 = nn.LayerNorm(embed_size)\n",
    "\n",
    "        self.feed_forward = nn.Sequential(\n",
    "            LORALayer(embed_size, forward_expansion * embed_size, rank),\n",
    "            nn.ReLU(),\n",
    "            LORALayer(forward_expansion * embed_size, embed_size, rank),\n",
    "        )\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, value, key, query, mask):\n",
    "        attention = self.attention(value, key, query, mask)\n",
    "        x = self.dropout(self.norm1(attention + query))\n",
    "        forward = self.feed_forward(x)\n",
    "        out = self.dropout(self.norm2(forward + x))\n",
    "        return out\n",
    "\n",
    "class LanguageModelDecoder(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_size, num_layers, heads, forward_expansion, dropout, max_length, rank):\n",
    "        super(LanguageModelDecoder, self).__init__()\n",
    "        self.word_embedding = nn.Embedding(vocab_size, embed_size)\n",
    "        self.position_embedding = nn.Embedding(max_length, embed_size)\n",
    "\n",
    "        # Adding BatchNorm layers\n",
    "        self.bn1 = nn.BatchNorm1d(embed_size)\n",
    "        self.bn2 = nn.BatchNorm1d(embed_size)\n",
    "\n",
    "        self.layers = nn.ModuleList(\n",
    "            [\n",
    "                TransformerBlock(embed_size, heads, dropout, forward_expansion, rank)\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        # QLORA layers\n",
    "        self.qlora_feed_forward = nn.Sequential(\n",
    "            QLORALayer(embed_size, forward_expansion * embed_size, rank),\n",
    "            nn.ReLU(),\n",
    "            QLORALayer(forward_expansion * embed_size, embed_size, rank),\n",
    "        )\n",
    "        self.use_qlora = False  # Flag to toggle QLORA\n",
    "\n",
    "        self.fc_out = nn.Linear(embed_size, vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, trg_mask):\n",
    "        N, seq_length = x.shape\n",
    "        positions = torch.arange(0, seq_length).expand(N, seq_length).to(x.device)\n",
    "        x = self.dropout(self.word_embedding(x) + self.position_embedding(positions))\n",
    "\n",
    "        # Transpose for BatchNorm, apply batch normalization, and then transpose back\n",
    "        x = x.transpose(1, 2)\n",
    "        x = self.bn1(x)\n",
    "        x = x.transpose(1, 2)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, x, x, trg_mask)\n",
    "            if self.use_qlora:\n",
    "                x = self.qlora_feed_forward(x)\n",
    "\n",
    "        # Transpose for BatchNorm, apply batch normalization, and then transpose back\n",
    "        x = x.transpose(1, 2)\n",
    "        x = self.bn2(x)\n",
    "        x = x.transpose(1, 2)\n",
    "\n",
    "        out = self.fc_out(x)\n",
    "        return out\n",
    "\n",
    "    def toggle_qlora(self, use_qlora: bool):\n",
    "        self.use_qlora = use_qlora\n",
    "\n",
    "\n",
    "class LanguageModelTransformer(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_size=256, num_layers=6, forward_expansion=4, heads=8, dropout=0, max_length=100, rank=16):\n",
    "        super(LanguageModelTransformer, self).__init__()\n",
    "\n",
    "        self.decoder = LanguageModelDecoder(\n",
    "            vocab_size,\n",
    "            embed_size,\n",
    "            num_layers,\n",
    "            heads,\n",
    "            forward_expansion,\n",
    "            dropout,\n",
    "            max_length,\n",
    "            rank,\n",
    "        )\n",
    "\n",
    "    def forward(self, trg):\n",
    "        trg_mask = self.make_trg_mask(trg)\n",
    "        out = self.decoder(trg, trg_mask)\n",
    "        return out\n",
    "\n",
    "    def make_trg_mask(self, trg):\n",
    "        N, trg_len = trg.shape\n",
    "        trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(\n",
    "            N, 1, trg_len, trg_len\n",
    "        ).to(trg.device)\n",
    "\n",
    "        return trg_mask\n",
    "\n",
    "# Define vocabulary size and dummy data parameters\n",
    "NUM_WORDS = 1000  # Example vocabulary size\n",
    "sequence_length = 30  # Sequence length for the LanguageDataset\n",
    "dummy_data_size = 1000  # Total number of tokens in the dummy dataset\n",
    "\n",
    "\n",
    "\n",
    "# Load dataset\n",
    "dataset = load_dataset('wikipedia', '20220301.simple')\n",
    "\n",
    "# Initialize tokenizer\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "vocab_size = tokenizer.vocab_size\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    # Tokenize the text\n",
    "    tokenized_output = tokenizer(examples['text'], padding='max_length', truncation=True, max_length=sequence_length)\n",
    "    \n",
    "    # Shift input_ids to create labels and truncate the last token\n",
    "    labels = [seq[1:] + [tokenizer.pad_token_id] for seq in tokenized_output['input_ids']]\n",
    "    tokenized_output['labels'] = labels\n",
    "    \n",
    "    return tokenized_output\n",
    "\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)\n",
    "\n",
    "# Set format for PyTorch\n",
    "tokenized_datasets.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])\n",
    "\n",
    "train_loader = DataLoader(tokenized_datasets['train'], batch_size=64, shuffle=True)\n",
    "\n",
    "# Define device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Create the model instance\n",
    "model = LanguageModelTransformer(\n",
    "    vocab_size=vocab_size,  # Use the vocab size from the tokenizer\n",
    "    embed_size=256,\n",
    "    num_layers=6,\n",
    "    forward_expansion=4,\n",
    "    heads=8,\n",
    "    dropout=0,\n",
    "    max_length=100,\n",
    "    rank=16\n",
    ").to(device)\n",
    "\n",
    "\n",
    "def calculate_new_alpha(current_loss, initial_loss, initial_alpha=1.0, final_alpha=0.1):\n",
    "    \"\"\"\n",
    "    Calculate a new alpha value based on the current loss.\n",
    "    \"\"\"\n",
    "    if current_loss >= initial_loss:\n",
    "        return initial_alpha  # Keep initial alpha if loss isn't decreasing\n",
    "\n",
    "    loss_ratio = current_loss / initial_loss\n",
    "    alpha_range = initial_alpha - final_alpha\n",
    "    new_alpha = final_alpha + (alpha_range * loss_ratio)\n",
    "    return new_alpha\n",
    "\n",
    "# Enable QLORA during training\n",
    "model.decoder.toggle_qlora(True)\n",
    "\n",
    "initial_loss = None\n",
    "# Training loop\n",
    "# Assuming model is an instance of LanguageModelTransformer\n",
    "# Define loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = AdamW(model.parameters(), lr=1e-8, weight_decay=1e-4)\n",
    "scheduler = StepLR(optimizer, step_size=4, gamma=0.98)\n",
    "num_epochs = 5\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    model.decoder.toggle_qlora(True)\n",
    "    total_loss = 0\n",
    "\n",
    "    for batch_idx, batch in enumerate(train_loader):\n",
    "        inputs = batch['input_ids'].to(device)\n",
    "        targets = batch['labels'].to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))\n",
    "\n",
    "        # Check for NaN in loss\n",
    "        if math.isnan(loss.item()):\n",
    "            print(\"Encountered NaN loss, stopping training\")\n",
    "            break\n",
    "\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)\n",
    "        optimizer.step()\n",
    "\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        # Set the initial_loss after the first batch of the first epoch\n",
    "        if initial_loss is None and batch_idx == 0:\n",
    "            initial_loss = loss.item()\n",
    "\n",
    "    scheduler.step()\n",
    "\n",
    "    # Check for NaN in total_loss\n",
    "    if math.isnan(total_loss):\n",
    "        print(f\"Epoch {epoch+1}/{num_epochs} stopped due to NaN loss\")\n",
    "        break\n",
    "    else:\n",
    "        print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss / len(train_loader)}\")\n",
    "\n",
    "    # Average loss for the epoch\n",
    "    average_loss = total_loss / len(train_loader)\n",
    "\n",
    "    # Update alpha at the end of each epoch based on the average loss\n",
    "    new_alpha = calculate_new_alpha(average_loss, initial_loss)\n",
    "    for layer in model.modules():\n",
    "        if isinstance(layer, QLORALayer):\n",
    "            layer.update_alpha(new_alpha)\n",
    "\n",
    "    #model.decoder.toggle_qlora(False)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "241539925f2442bbb688f8e41af47df6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6f53b5e8be03487dbf2568633bd5416b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/8040549 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "OSError",
     "evalue": "[Errno 28] No space left on device",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mOSError\u001b[0m                                   Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[3], line 15\u001b[0m\n\u001b[0;32m     13\u001b[0m subset_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mint\u001b[39m(\u001b[38;5;241m0.3\u001b[39m \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mlen\u001b[39m(dataset[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m'\u001b[39m]))  \u001b[38;5;66;03m# 30% of the dataset\u001b[39;00m\n\u001b[0;32m     14\u001b[0m subset_indices \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mrandperm(\u001b[38;5;28mlen\u001b[39m(dataset[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m'\u001b[39m]))\u001b[38;5;241m.\u001b[39mtolist()[:subset_size]  \u001b[38;5;66;03m# Randomly select indices\u001b[39;00m\n\u001b[1;32m---> 15\u001b[0m formatted_dataset \u001b[38;5;241m=\u001b[39m \u001b[43mdataset\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mtrain\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mselect\u001b[49m\u001b[43m(\u001b[49m\u001b[43msubset_indices\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mformat_stackexchange_dpo\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatched\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mload_from_cache_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[0;32m     17\u001b[0m \u001b[38;5;66;03m# Convert formatted dataset to DataLoader for batch processing\u001b[39;00m\n\u001b[0;32m     18\u001b[0m dpo_dataloader \u001b[38;5;241m=\u001b[39m DataLoader(formatted_dataset, batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m64\u001b[39m, shuffle\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:592\u001b[0m, in \u001b[0;36mtransmit_tasks.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    590\u001b[0m     \u001b[38;5;28mself\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mself\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    591\u001b[0m \u001b[38;5;66;03m# apply actual function\u001b[39;00m\n\u001b[1;32m--> 592\u001b[0m out: Union[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDatasetDict\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\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[0;32m    593\u001b[0m datasets: List[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(out\u001b[38;5;241m.\u001b[39mvalues()) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(out, \u001b[38;5;28mdict\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m [out]\n\u001b[0;32m    594\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m dataset \u001b[38;5;129;01min\u001b[39;00m datasets:\n\u001b[0;32m    595\u001b[0m     \u001b[38;5;66;03m# Remove task templates if a column mapping of the template is no longer valid\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:557\u001b[0m, in \u001b[0;36mtransmit_format.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    550\u001b[0m self_format \u001b[38;5;241m=\u001b[39m {\n\u001b[0;32m    551\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtype\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_type,\n\u001b[0;32m    552\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mformat_kwargs\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_kwargs,\n\u001b[0;32m    553\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcolumns\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_columns,\n\u001b[0;32m    554\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moutput_all_columns\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_output_all_columns,\n\u001b[0;32m    555\u001b[0m }\n\u001b[0;32m    556\u001b[0m \u001b[38;5;66;03m# apply actual function\u001b[39;00m\n\u001b[1;32m--> 557\u001b[0m out: Union[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDatasetDict\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\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[0;32m    558\u001b[0m datasets: List[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(out\u001b[38;5;241m.\u001b[39mvalues()) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(out, \u001b[38;5;28mdict\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m [out]\n\u001b[0;32m    559\u001b[0m \u001b[38;5;66;03m# re-apply format to the output\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:3097\u001b[0m, in \u001b[0;36mDataset.map\u001b[1;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)\u001b[0m\n\u001b[0;32m   3090\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m transformed_dataset \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m   3091\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m logging\u001b[38;5;241m.\u001b[39mtqdm(\n\u001b[0;32m   3092\u001b[0m         disable\u001b[38;5;241m=\u001b[39m\u001b[38;5;129;01mnot\u001b[39;00m logging\u001b[38;5;241m.\u001b[39mis_progress_bar_enabled(),\n\u001b[0;32m   3093\u001b[0m         unit\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m examples\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   3094\u001b[0m         total\u001b[38;5;241m=\u001b[39mpbar_total,\n\u001b[0;32m   3095\u001b[0m         desc\u001b[38;5;241m=\u001b[39mdesc \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMap\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   3096\u001b[0m     ) \u001b[38;5;28;01mas\u001b[39;00m pbar:\n\u001b[1;32m-> 3097\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mrank\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdone\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontent\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mDataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_map_single\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mdataset_kwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m   3098\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mdone\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m   3099\u001b[0m \u001b[43m                \u001b[49m\u001b[43mshards_done\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:3493\u001b[0m, in \u001b[0;36mDataset._map_single\u001b[1;34m(shard, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset)\u001b[0m\n\u001b[0;32m   3491\u001b[0m         writer\u001b[38;5;241m.\u001b[39mwrite_table(pa\u001b[38;5;241m.\u001b[39mTable\u001b[38;5;241m.\u001b[39mfrom_pandas(batch))\n\u001b[0;32m   3492\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 3493\u001b[0m         \u001b[43mwriter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_batch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   3494\u001b[0m num_examples_progress_update \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m num_examples_in_batch\n\u001b[0;32m   3495\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m time\u001b[38;5;241m.\u001b[39mtime() \u001b[38;5;241m>\u001b[39m _time \u001b[38;5;241m+\u001b[39m config\u001b[38;5;241m.\u001b[39mPBAR_REFRESH_TIME_INTERVAL:\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_writer.py:559\u001b[0m, in \u001b[0;36mArrowWriter.write_batch\u001b[1;34m(self, batch_examples, writer_batch_size)\u001b[0m\n\u001b[0;32m    557\u001b[0m schema \u001b[38;5;241m=\u001b[39m inferred_features\u001b[38;5;241m.\u001b[39marrow_schema \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpa_writer \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mschema\n\u001b[0;32m    558\u001b[0m pa_table \u001b[38;5;241m=\u001b[39m pa\u001b[38;5;241m.\u001b[39mTable\u001b[38;5;241m.\u001b[39mfrom_arrays(arrays, schema\u001b[38;5;241m=\u001b[39mschema)\n\u001b[1;32m--> 559\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_table\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpa_table\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_writer.py:577\u001b[0m, in \u001b[0;36mArrowWriter.write_table\u001b[1;34m(self, pa_table, writer_batch_size)\u001b[0m\n\u001b[0;32m    575\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_bytes \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m pa_table\u001b[38;5;241m.\u001b[39mnbytes\n\u001b[0;32m    576\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_examples \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m pa_table\u001b[38;5;241m.\u001b[39mnum_rows\n\u001b[1;32m--> 577\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpa_writer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_table\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpa_table\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\pyarrow\\ipc.pxi:525\u001b[0m, in \u001b[0;36mpyarrow.lib._CRecordBatchWriter.write_table\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\fsspec\\implementations\\local.py:369\u001b[0m, in \u001b[0;36mLocalFileOpener.write\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    368\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrite\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m--> 369\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[43mf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite\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;31mOSError\u001b[0m: [Errno 28] No space left on device"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from datasets import load_dataset\n",
    "\n",
    "def format_stackexchange_dpo(samples):\n",
    "    return {\n",
    "        \"prompt\": [\"Question: \" + question + \"\\n\\nAnswer: \" for question in samples[\"question\"]],\n",
    "        \"chosen\": samples[\"response_j\"],   # Rated better than k\n",
    "        \"rejected\": samples[\"response_k\"], # Rated worse than j\n",
    "    }\n",
    "\n",
    "# Load and format a subset (30%) of the StackExchange DPO dataset\n",
    "dataset = load_dataset(\"lvwerra/stack-exchange-paired\")\n",
    "subset_size = int(0.3 * len(dataset['train']))  # 30% of the dataset\n",
    "subset_indices = torch.randperm(len(dataset['train'])).tolist()[:subset_size]  # Randomly select indices\n",
    "formatted_dataset = dataset['train'].select(subset_indices).map(format_stackexchange_dpo, batched=True, load_from_cache_file=False)\n",
    "\n",
    "# Convert formatted dataset to DataLoader for batch processing\n",
    "dpo_dataloader = DataLoader(formatted_dataset, batch_size=64, shuffle=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import MarginRankingLoss\n",
    "\n",
    "# Define DPO-specific loss function\n",
    "dpo_loss_function = MarginRankingLoss(margin=1.0)\n",
    "dpo_num_epochs = 2  # Define the number of epochs for DPO training\n",
    "\n",
    "# DPO Training loop\n",
    "for epoch in range(dpo_num_epochs):\n",
    "    model.train()  # Ensure the model is in training mode\n",
    "    total_dpo_loss = 0\n",
    "\n",
    "    for batch in dpo_dataloader:\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Prepare the input for the model\n",
    "        prompts = batch['prompt'].to(device)\n",
    "        preferred_responses = batch['chosen'].to(device)\n",
    "        less_preferred_responses = batch['rejected'].to(device)\n",
    "\n",
    "        # Forward pass and model's scoring mechanism for responses\n",
    "        # The model should output scores for the preferred and less-preferred responses\n",
    "        output_preferred = model(preferred_responses)\n",
    "        output_less_preferred = model(less_preferred_responses)\n",
    "\n",
    "        # Compute DPO loss\n",
    "        dpo_loss = dpo_loss_function(output_preferred, output_less_preferred, torch.ones(output_preferred.size(0)).to(device))\n",
    "        total_dpo_loss += dpo_loss.item()\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        dpo_loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)\n",
    "        optimizer.step()\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{dpo_num_epochs}, DPO Loss: {total_dpo_loss / len(dpo_dataloader)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9fea535cbdee4130a7d781b89b49306f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27fb83651ab742229337be6e2791f636",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8adffbc8fcc94c98b3d56ec57aecf2d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "72e92636ea084e43a5b8be22748295b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/315M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90a27fad206e4ceaae78bf4a0753a05a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "47f9ecacbf1e41e39658a44aa260a0c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38b764a9b41e4cffb9d81f6d1673b483",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/315M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a3db09d349f645b797816ccb4fff9d02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39da8049d30644c7bb8e4645a8193051",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/315M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9fded1f7ed4457e8c34cbe23447eeef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "72d35ce6943945f18511c82e4be24441",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/315M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "74f8790ac422461abc924e6f180b1930",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/315M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "032680e87ed04781ae913d9224dff235",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc4e7cfb38694e429fe638af7b64ef17",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f364e8969cf4e87aa9e03ff1fb2310a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f330912806e74ccbacc6a65762b2d8d8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b500bff5ef30499e90ba7a63573f380b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2345471103e64529ad1c7cbbf6742bb5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ccd54d3a16c64bbe86b135210aa1a3db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9511d1d871b141bdb35b979c1cfcdd61",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83900e1172164d0c9c608fb68654af07",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "278ffd8341a14a739d7357df67f46337",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d38343b2bf134b8aa99dec89a1b5ce5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "324cd976ec444c0f8bcc81e430729fa4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9d8b1855f2fd43d2bf7a4062aecc34d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/312M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3df35f478098429899a6b38f974afa1d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78b16c98b1fb4036958a8f7a8c0e8d67",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5dde1e0e2a084d47b059f3803ab1f3a1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "417685210bee47999b663b78317cc759",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bbdc5c495cfa46968fa58ab45e94a804",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "15222ea75d47455cbe280cd345f30681",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/311M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1be121a4c10a4c32a7733c6696311604",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/312M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b653e482bc3f40dd9b7b37246d7a62f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c3560da56e74b51a64479a9e28329f2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5c6290bd041a4174a5256215f34f565e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/315M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d22bbb4478844a381011bab951b27d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/312M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "58ebcdfdcc2b4843bcd9e4213990ba8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "709bb33a0bbd47bc972bbaa7dc97aa78",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e08eb94a16114cd594b97441854a21cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4bc5c920e0bf43a1a221a6ae612d7313",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/312M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a55e624eeece44c1a8a0ff9f126cb966",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/312M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "32eba012d1e44e70bcacc32e3a9b78b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d906d934b47a45f595af1d4f0b4c288b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9cf26a7366e241d6a9f7ca94637a59bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34f7adf7abd24c05834c4a1a4ca82040",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9060f6efb36542e5a65928198af07a15",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "060c6c9b6267401db1a800fd22bd4870",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "715a59dcd74e4c2289cc9dc9eb2f522c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c1270810588548a8b96c9584a144a6c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "86eb0490a966463d8c4547d2a0b30921",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dfec2a70c0594bb0a6498a63466394bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/312M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "19b5989a88224d0c9cb63f3fec2350ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f14e3bd4ea234adfb718193f8cce3102",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/313M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2417012de86441449b73b6aad3955e76",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/314M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ddd6285621d9475b9d74740863652bb5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "260d6cbbf40345a3be893e4b959099f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d072eed005b84f948056043ec13dde6b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16b1b11d8bf4416c96dfc561b87e0c34",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/26801833 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "OSError",
     "evalue": "[Errno 28] No space left on device",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mOSError\u001b[0m                                   Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[6], line 20\u001b[0m\n\u001b[0;32m     17\u001b[0m dataset \u001b[38;5;241m=\u001b[39m load_dataset(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlvwerra/stack-exchange-paired\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     19\u001b[0m \u001b[38;5;66;03m# Format the dataset for DPO\u001b[39;00m\n\u001b[1;32m---> 20\u001b[0m formatted_dataset \u001b[38;5;241m=\u001b[39m \u001b[43mdataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mformat_stackexchange_dpo\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatched\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[0;32m     22\u001b[0m \u001b[38;5;66;03m# Display the first few formatted examples (optional)\u001b[39;00m\n\u001b[0;32m     23\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m3\u001b[39m):\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\dataset_dict.py:853\u001b[0m, in \u001b[0;36mDatasetDict.map\u001b[1;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, desc)\u001b[0m\n\u001b[0;32m    850\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cache_file_names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    851\u001b[0m     cache_file_names \u001b[38;5;241m=\u001b[39m {k: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m}\n\u001b[0;32m    852\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m DatasetDict(\n\u001b[1;32m--> 853\u001b[0m     \u001b[43m{\u001b[49m\n\u001b[0;32m    854\u001b[0m \u001b[43m        \u001b[49m\u001b[43mk\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    855\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfunction\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfunction\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    856\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwith_indices\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwith_indices\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    857\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwith_rank\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwith_rank\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    858\u001b[0m \u001b[43m            \u001b[49m\u001b[43minput_columns\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_columns\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    859\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatched\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatched\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    860\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    861\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdrop_last_batch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdrop_last_batch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    862\u001b[0m \u001b[43m            \u001b[49m\u001b[43mremove_columns\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mremove_columns\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    863\u001b[0m \u001b[43m            \u001b[49m\u001b[43mkeep_in_memory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeep_in_memory\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    864\u001b[0m \u001b[43m            \u001b[49m\u001b[43mload_from_cache_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mload_from_cache_file\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    865\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcache_file_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_file_names\u001b[49m\u001b[43m[\u001b[49m\u001b[43mk\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    866\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    867\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfeatures\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfeatures\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    868\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdisable_nullable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdisable_nullable\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    869\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfn_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfn_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    870\u001b[0m \u001b[43m            \u001b[49m\u001b[43mnum_proc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnum_proc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    871\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdesc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdesc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    872\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    873\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mk\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    874\u001b[0m \u001b[43m    \u001b[49m\u001b[43m}\u001b[49m\n\u001b[0;32m    875\u001b[0m )\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\dataset_dict.py:854\u001b[0m, in \u001b[0;36m<dictcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m    850\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cache_file_names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    851\u001b[0m     cache_file_names \u001b[38;5;241m=\u001b[39m {k: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m}\n\u001b[0;32m    852\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m DatasetDict(\n\u001b[0;32m    853\u001b[0m     {\n\u001b[1;32m--> 854\u001b[0m         k: \u001b[43mdataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    855\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfunction\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfunction\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    856\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwith_indices\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwith_indices\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    857\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwith_rank\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwith_rank\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    858\u001b[0m \u001b[43m            \u001b[49m\u001b[43minput_columns\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_columns\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    859\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatched\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatched\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    860\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    861\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdrop_last_batch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdrop_last_batch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    862\u001b[0m \u001b[43m            \u001b[49m\u001b[43mremove_columns\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mremove_columns\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    863\u001b[0m \u001b[43m            \u001b[49m\u001b[43mkeep_in_memory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeep_in_memory\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    864\u001b[0m \u001b[43m            \u001b[49m\u001b[43mload_from_cache_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mload_from_cache_file\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    865\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcache_file_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_file_names\u001b[49m\u001b[43m[\u001b[49m\u001b[43mk\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    866\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    867\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfeatures\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfeatures\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    868\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdisable_nullable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdisable_nullable\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    869\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfn_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfn_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    870\u001b[0m \u001b[43m            \u001b[49m\u001b[43mnum_proc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnum_proc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    871\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdesc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdesc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    872\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    873\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m k, dataset \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mitems()\n\u001b[0;32m    874\u001b[0m     }\n\u001b[0;32m    875\u001b[0m )\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:592\u001b[0m, in \u001b[0;36mtransmit_tasks.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    590\u001b[0m     \u001b[38;5;28mself\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mself\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    591\u001b[0m \u001b[38;5;66;03m# apply actual function\u001b[39;00m\n\u001b[1;32m--> 592\u001b[0m out: Union[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDatasetDict\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\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[0;32m    593\u001b[0m datasets: List[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(out\u001b[38;5;241m.\u001b[39mvalues()) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(out, \u001b[38;5;28mdict\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m [out]\n\u001b[0;32m    594\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m dataset \u001b[38;5;129;01min\u001b[39;00m datasets:\n\u001b[0;32m    595\u001b[0m     \u001b[38;5;66;03m# Remove task templates if a column mapping of the template is no longer valid\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:557\u001b[0m, in \u001b[0;36mtransmit_format.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    550\u001b[0m self_format \u001b[38;5;241m=\u001b[39m {\n\u001b[0;32m    551\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtype\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_type,\n\u001b[0;32m    552\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mformat_kwargs\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_kwargs,\n\u001b[0;32m    553\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcolumns\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_columns,\n\u001b[0;32m    554\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moutput_all_columns\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_output_all_columns,\n\u001b[0;32m    555\u001b[0m }\n\u001b[0;32m    556\u001b[0m \u001b[38;5;66;03m# apply actual function\u001b[39;00m\n\u001b[1;32m--> 557\u001b[0m out: Union[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDatasetDict\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\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[0;32m    558\u001b[0m datasets: List[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(out\u001b[38;5;241m.\u001b[39mvalues()) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(out, \u001b[38;5;28mdict\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m [out]\n\u001b[0;32m    559\u001b[0m \u001b[38;5;66;03m# re-apply format to the output\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:3097\u001b[0m, in \u001b[0;36mDataset.map\u001b[1;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)\u001b[0m\n\u001b[0;32m   3090\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m transformed_dataset \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m   3091\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m logging\u001b[38;5;241m.\u001b[39mtqdm(\n\u001b[0;32m   3092\u001b[0m         disable\u001b[38;5;241m=\u001b[39m\u001b[38;5;129;01mnot\u001b[39;00m logging\u001b[38;5;241m.\u001b[39mis_progress_bar_enabled(),\n\u001b[0;32m   3093\u001b[0m         unit\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m examples\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   3094\u001b[0m         total\u001b[38;5;241m=\u001b[39mpbar_total,\n\u001b[0;32m   3095\u001b[0m         desc\u001b[38;5;241m=\u001b[39mdesc \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMap\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   3096\u001b[0m     ) \u001b[38;5;28;01mas\u001b[39;00m pbar:\n\u001b[1;32m-> 3097\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mrank\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdone\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontent\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mDataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_map_single\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mdataset_kwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m   3098\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mdone\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m   3099\u001b[0m \u001b[43m                \u001b[49m\u001b[43mshards_done\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:3493\u001b[0m, in \u001b[0;36mDataset._map_single\u001b[1;34m(shard, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset)\u001b[0m\n\u001b[0;32m   3491\u001b[0m         writer\u001b[38;5;241m.\u001b[39mwrite_table(pa\u001b[38;5;241m.\u001b[39mTable\u001b[38;5;241m.\u001b[39mfrom_pandas(batch))\n\u001b[0;32m   3492\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 3493\u001b[0m         \u001b[43mwriter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_batch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   3494\u001b[0m num_examples_progress_update \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m num_examples_in_batch\n\u001b[0;32m   3495\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m time\u001b[38;5;241m.\u001b[39mtime() \u001b[38;5;241m>\u001b[39m _time \u001b[38;5;241m+\u001b[39m config\u001b[38;5;241m.\u001b[39mPBAR_REFRESH_TIME_INTERVAL:\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_writer.py:559\u001b[0m, in \u001b[0;36mArrowWriter.write_batch\u001b[1;34m(self, batch_examples, writer_batch_size)\u001b[0m\n\u001b[0;32m    557\u001b[0m schema \u001b[38;5;241m=\u001b[39m inferred_features\u001b[38;5;241m.\u001b[39marrow_schema \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpa_writer \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mschema\n\u001b[0;32m    558\u001b[0m pa_table \u001b[38;5;241m=\u001b[39m pa\u001b[38;5;241m.\u001b[39mTable\u001b[38;5;241m.\u001b[39mfrom_arrays(arrays, schema\u001b[38;5;241m=\u001b[39mschema)\n\u001b[1;32m--> 559\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_table\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpa_table\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_writer.py:577\u001b[0m, in \u001b[0;36mArrowWriter.write_table\u001b[1;34m(self, pa_table, writer_batch_size)\u001b[0m\n\u001b[0;32m    575\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_bytes \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m pa_table\u001b[38;5;241m.\u001b[39mnbytes\n\u001b[0;32m    576\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_examples \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m pa_table\u001b[38;5;241m.\u001b[39mnum_rows\n\u001b[1;32m--> 577\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpa_writer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_table\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpa_table\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\pyarrow\\ipc.pxi:525\u001b[0m, in \u001b[0;36mpyarrow.lib._CRecordBatchWriter.write_table\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\fsspec\\implementations\\local.py:369\u001b[0m, in \u001b[0;36mLocalFileOpener.write\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    368\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrite\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m--> 369\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[43mf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite\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;31mOSError\u001b[0m: [Errno 28] No space left on device"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from torch.nn import MarginRankingLoss\n",
    "from datasets import load_dataset\n",
    "\n",
    "def format_stackexchange_dpo(samples):\n",
    "    \"\"\"Format StackExchange dataset for DPO.\"\"\"\n",
    "    return {\n",
    "        \"prompt\": [\n",
    "            \"Question: \" + question + \"\\n\\nAnswer: \"\n",
    "            for question in samples[\"question\"]\n",
    "        ],\n",
    "        \"chosen\": samples[\"response_j\"],   # Rated better than k\n",
    "        \"rejected\": samples[\"response_k\"], # Rated worse than j\n",
    "    }\n",
    "\n",
    "# Load the StackExchange DPO dataset\n",
    "dataset = load_dataset(\"lvwerra/stack-exchange-paired\")\n",
    "\n",
    "# Format the dataset for DPO\n",
    "formatted_dataset = dataset.map(format_stackexchange_dpo, batched=True)\n",
    "\n",
    "# Display the first few formatted examples (optional)\n",
    "for i in range(3):\n",
    "    print(f\"Sample {i+1}:\")\n",
    "    print(\"Prompt:\", formatted_dataset['train'][i]['prompt'])\n",
    "    print(\"Chosen Answer:\", formatted_dataset['train'][i]['chosen'])\n",
    "    print(\"Rejected Answer:\", formatted_dataset['train'][i]['rejected'])\n",
    "    print()\n",
    "\n",
    "\n",
    "# Convert formatted dataset to DataLoader for batch processing\n",
    "dpo_dataloader = DataLoader(formatted_dataset['train'], batch_size=64, shuffle=True)\n",
    "\n",
    "# Define DPO-specific loss function\n",
    "dpo_loss_function = MarginRankingLoss(margin=1.0)\n",
    "\n",
    "dpo_num_epochs = 2\n",
    "\n",
    "# DPO Training loop\n",
    "for epoch in range(dpo_num_epochs):\n",
    "    model.train()\n",
    "    total_dpo_loss = 0\n",
    "\n",
    "    for batch in dpo_dataloader:\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        prompts = batch['prompt'].to(device)\n",
    "        preferred_responses = batch['chosen'].to(device)\n",
    "        less_preferred_responses = batch['rejected'].to(device)\n",
    "\n",
    "        # Forward pass and model's scoring mechanism for responses\n",
    "        # Modify according to how your model outputs scores\n",
    "        output_preferred = model(preferred_responses)  # Needs specific implementation\n",
    "        output_less_preferred = model(less_preferred_responses)  # Needs specific implementation\n",
    "\n",
    "        # Compute DPO loss\n",
    "        dpo_loss = dpo_loss_function(output_preferred, output_less_preferred, torch.ones(output_preferred.size(0)).to(device))\n",
    "        total_dpo_loss += dpo_loss.item()\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        dpo_loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)\n",
    "        optimizer.step()\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{dpo_num_epochs}, DPO Loss: {total_dpo_loss / len(dpo_dataloader)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\robbi\\AppData\\Local\\Temp\\ipykernel_20456\\1180072541.py:8: FutureWarning: set_caching_enabled is deprecated and will be removed in the next major version of datasets. Use datasets.enable_caching() or datasets.disable_caching() instead. This function will be removed in a future version of datasets.\n",
      "  set_caching_enabled(False)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7b35d31ffb8b42b89edf79ea74035fd4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "57aec96b0c534cbe8e23688f9daa84f1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/26801833 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "OSError",
     "evalue": "[Errno 28] No space left on device",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mOSError\u001b[0m                                   Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[7], line 23\u001b[0m\n\u001b[0;32m     21\u001b[0m \u001b[38;5;66;03m# Load and format the StackExchange DPO dataset\u001b[39;00m\n\u001b[0;32m     22\u001b[0m dataset \u001b[38;5;241m=\u001b[39m load_dataset(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlvwerra/stack-exchange-paired\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m---> 23\u001b[0m formatted_dataset \u001b[38;5;241m=\u001b[39m \u001b[43mdataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mformat_stackexchange_dpo\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatched\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mload_from_cache_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[0;32m     25\u001b[0m \u001b[38;5;66;03m# Convert formatted dataset to DataLoader for batch processing\u001b[39;00m\n\u001b[0;32m     26\u001b[0m dpo_dataloader \u001b[38;5;241m=\u001b[39m DataLoader(formatted_dataset[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m'\u001b[39m], batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m64\u001b[39m, shuffle\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\dataset_dict.py:853\u001b[0m, in \u001b[0;36mDatasetDict.map\u001b[1;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, desc)\u001b[0m\n\u001b[0;32m    850\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cache_file_names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    851\u001b[0m     cache_file_names \u001b[38;5;241m=\u001b[39m {k: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m}\n\u001b[0;32m    852\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m DatasetDict(\n\u001b[1;32m--> 853\u001b[0m     \u001b[43m{\u001b[49m\n\u001b[0;32m    854\u001b[0m \u001b[43m        \u001b[49m\u001b[43mk\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    855\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfunction\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfunction\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    856\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwith_indices\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwith_indices\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    857\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwith_rank\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwith_rank\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    858\u001b[0m \u001b[43m            \u001b[49m\u001b[43minput_columns\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_columns\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    859\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatched\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatched\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    860\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    861\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdrop_last_batch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdrop_last_batch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    862\u001b[0m \u001b[43m            \u001b[49m\u001b[43mremove_columns\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mremove_columns\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    863\u001b[0m \u001b[43m            \u001b[49m\u001b[43mkeep_in_memory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeep_in_memory\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    864\u001b[0m \u001b[43m            \u001b[49m\u001b[43mload_from_cache_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mload_from_cache_file\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    865\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcache_file_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_file_names\u001b[49m\u001b[43m[\u001b[49m\u001b[43mk\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    866\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    867\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfeatures\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfeatures\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    868\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdisable_nullable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdisable_nullable\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    869\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfn_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfn_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    870\u001b[0m \u001b[43m            \u001b[49m\u001b[43mnum_proc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnum_proc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    871\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdesc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdesc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    872\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    873\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mk\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    874\u001b[0m \u001b[43m    \u001b[49m\u001b[43m}\u001b[49m\n\u001b[0;32m    875\u001b[0m )\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\dataset_dict.py:854\u001b[0m, in \u001b[0;36m<dictcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m    850\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cache_file_names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    851\u001b[0m     cache_file_names \u001b[38;5;241m=\u001b[39m {k: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m}\n\u001b[0;32m    852\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m DatasetDict(\n\u001b[0;32m    853\u001b[0m     {\n\u001b[1;32m--> 854\u001b[0m         k: \u001b[43mdataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    855\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfunction\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfunction\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    856\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwith_indices\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwith_indices\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    857\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwith_rank\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwith_rank\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    858\u001b[0m \u001b[43m            \u001b[49m\u001b[43minput_columns\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_columns\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    859\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatched\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatched\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    860\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    861\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdrop_last_batch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdrop_last_batch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    862\u001b[0m \u001b[43m            \u001b[49m\u001b[43mremove_columns\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mremove_columns\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    863\u001b[0m \u001b[43m            \u001b[49m\u001b[43mkeep_in_memory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeep_in_memory\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    864\u001b[0m \u001b[43m            \u001b[49m\u001b[43mload_from_cache_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mload_from_cache_file\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    865\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcache_file_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_file_names\u001b[49m\u001b[43m[\u001b[49m\u001b[43mk\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    866\u001b[0m \u001b[43m            \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    867\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfeatures\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfeatures\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    868\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdisable_nullable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdisable_nullable\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    869\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfn_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfn_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    870\u001b[0m \u001b[43m            \u001b[49m\u001b[43mnum_proc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnum_proc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    871\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdesc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdesc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    872\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    873\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m k, dataset \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mitems()\n\u001b[0;32m    874\u001b[0m     }\n\u001b[0;32m    875\u001b[0m )\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:592\u001b[0m, in \u001b[0;36mtransmit_tasks.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    590\u001b[0m     \u001b[38;5;28mself\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mself\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    591\u001b[0m \u001b[38;5;66;03m# apply actual function\u001b[39;00m\n\u001b[1;32m--> 592\u001b[0m out: Union[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDatasetDict\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\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[0;32m    593\u001b[0m datasets: List[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(out\u001b[38;5;241m.\u001b[39mvalues()) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(out, \u001b[38;5;28mdict\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m [out]\n\u001b[0;32m    594\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m dataset \u001b[38;5;129;01min\u001b[39;00m datasets:\n\u001b[0;32m    595\u001b[0m     \u001b[38;5;66;03m# Remove task templates if a column mapping of the template is no longer valid\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:557\u001b[0m, in \u001b[0;36mtransmit_format.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    550\u001b[0m self_format \u001b[38;5;241m=\u001b[39m {\n\u001b[0;32m    551\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtype\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_type,\n\u001b[0;32m    552\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mformat_kwargs\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_kwargs,\n\u001b[0;32m    553\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcolumns\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_format_columns,\n\u001b[0;32m    554\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moutput_all_columns\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_output_all_columns,\n\u001b[0;32m    555\u001b[0m }\n\u001b[0;32m    556\u001b[0m \u001b[38;5;66;03m# apply actual function\u001b[39;00m\n\u001b[1;32m--> 557\u001b[0m out: Union[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDatasetDict\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\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[0;32m    558\u001b[0m datasets: List[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDataset\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(out\u001b[38;5;241m.\u001b[39mvalues()) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(out, \u001b[38;5;28mdict\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m [out]\n\u001b[0;32m    559\u001b[0m \u001b[38;5;66;03m# re-apply format to the output\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:3097\u001b[0m, in \u001b[0;36mDataset.map\u001b[1;34m(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)\u001b[0m\n\u001b[0;32m   3090\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m transformed_dataset \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m   3091\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m logging\u001b[38;5;241m.\u001b[39mtqdm(\n\u001b[0;32m   3092\u001b[0m         disable\u001b[38;5;241m=\u001b[39m\u001b[38;5;129;01mnot\u001b[39;00m logging\u001b[38;5;241m.\u001b[39mis_progress_bar_enabled(),\n\u001b[0;32m   3093\u001b[0m         unit\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m examples\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   3094\u001b[0m         total\u001b[38;5;241m=\u001b[39mpbar_total,\n\u001b[0;32m   3095\u001b[0m         desc\u001b[38;5;241m=\u001b[39mdesc \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMap\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   3096\u001b[0m     ) \u001b[38;5;28;01mas\u001b[39;00m pbar:\n\u001b[1;32m-> 3097\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mrank\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdone\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontent\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mDataset\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_map_single\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mdataset_kwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m   3098\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mdone\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m   3099\u001b[0m \u001b[43m                \u001b[49m\u001b[43mshards_done\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_dataset.py:3493\u001b[0m, in \u001b[0;36mDataset._map_single\u001b[1;34m(shard, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset)\u001b[0m\n\u001b[0;32m   3491\u001b[0m         writer\u001b[38;5;241m.\u001b[39mwrite_table(pa\u001b[38;5;241m.\u001b[39mTable\u001b[38;5;241m.\u001b[39mfrom_pandas(batch))\n\u001b[0;32m   3492\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 3493\u001b[0m         \u001b[43mwriter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_batch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   3494\u001b[0m num_examples_progress_update \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m num_examples_in_batch\n\u001b[0;32m   3495\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m time\u001b[38;5;241m.\u001b[39mtime() \u001b[38;5;241m>\u001b[39m _time \u001b[38;5;241m+\u001b[39m config\u001b[38;5;241m.\u001b[39mPBAR_REFRESH_TIME_INTERVAL:\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_writer.py:559\u001b[0m, in \u001b[0;36mArrowWriter.write_batch\u001b[1;34m(self, batch_examples, writer_batch_size)\u001b[0m\n\u001b[0;32m    557\u001b[0m schema \u001b[38;5;241m=\u001b[39m inferred_features\u001b[38;5;241m.\u001b[39marrow_schema \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpa_writer \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mschema\n\u001b[0;32m    558\u001b[0m pa_table \u001b[38;5;241m=\u001b[39m pa\u001b[38;5;241m.\u001b[39mTable\u001b[38;5;241m.\u001b[39mfrom_arrays(arrays, schema\u001b[38;5;241m=\u001b[39mschema)\n\u001b[1;32m--> 559\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_table\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpa_table\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\datasets\\arrow_writer.py:577\u001b[0m, in \u001b[0;36mArrowWriter.write_table\u001b[1;34m(self, pa_table, writer_batch_size)\u001b[0m\n\u001b[0;32m    575\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_bytes \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m pa_table\u001b[38;5;241m.\u001b[39mnbytes\n\u001b[0;32m    576\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_examples \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m pa_table\u001b[38;5;241m.\u001b[39mnum_rows\n\u001b[1;32m--> 577\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpa_writer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite_table\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpa_table\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwriter_batch_size\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\pyarrow\\ipc.pxi:525\u001b[0m, in \u001b[0;36mpyarrow.lib._CRecordBatchWriter.write_table\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\fsspec\\implementations\\local.py:369\u001b[0m, in \u001b[0;36mLocalFileOpener.write\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    368\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrite\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m--> 369\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[43mf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite\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;31mOSError\u001b[0m: [Errno 28] No space left on device"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset, set_caching_enabled\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.nn import MarginRankingLoss\n",
    "from transformers import BertTokenizer\n",
    "\n",
    "# Set a custom cache directory or disable caching\n",
    "set_caching_enabled(False)\n",
    "\n",
    "def format_stackexchange_dpo(samples):\n",
    "    \"\"\"Format StackExchange dataset for DPO.\"\"\"\n",
    "    return {\n",
    "        \"prompt\": [\n",
    "            \"Question: \" + question + \"\\n\\nAnswer: \"\n",
    "            for question in samples[\"question\"]\n",
    "        ],\n",
    "        \"chosen\": samples[\"response_j\"],   # Rated better than k\n",
    "        \"rejected\": samples[\"response_k\"], # Rated worse than j\n",
    "    }\n",
    "\n",
    "# Load and format the StackExchange DPO dataset\n",
    "dataset = load_dataset(\"lvwerra/stack-exchange-paired\")\n",
    "formatted_dataset = dataset.map(format_stackexchange_dpo, batched=True, load_from_cache_file=False)\n",
    "\n",
    "# Convert formatted dataset to DataLoader for batch processing\n",
    "dpo_dataloader = DataLoader(formatted_dataset['train'], batch_size=64, shuffle=True)\n",
    "\n",
    "# Initialize your model here (ensure it's already fine-tuned)\n",
    "# Example: model = YourFineTunedModel()\n",
    "\n",
    "# Define DPO-specific loss function\n",
    "dpo_loss_function = MarginRankingLoss(margin=1.0)\n",
    "\n",
    "# Define optimizer for DPO training\n",
    "# Example: optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)\n",
    "\n",
    "dpo_num_epochs = 2\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# DPO Training loop\n",
    "for epoch in range(dpo_num_epochs):\n",
    "    model.train()\n",
    "    total_dpo_loss = 0\n",
    "\n",
    "    for batch in dpo_dataloader:\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        prompts = batch['prompt'].to(device)\n",
    "        preferred_responses = batch['chosen'].to(device)\n",
    "        less_preferred_responses = batch['rejected'].to(device)\n",
    "\n",
    "        # Forward pass and model's scoring mechanism for responses\n",
    "        # Modify according to how your model outputs scores\n",
    "        output_preferred = model(preferred_responses)  # Needs specific implementation\n",
    "        output_less_preferred = model(less_preferred_responses)  # Needs specific implementation\n",
    "\n",
    "        # Compute DPO loss\n",
    "        dpo_loss = dpo_loss_function(output_preferred, output_less_preferred, torch.ones(output_preferred.size(0)).to(device))\n",
    "        total_dpo_loss += dpo_loss.item()\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        dpo_loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)\n",
    "        optimizer.step()\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{dpo_num_epochs}, DPO Loss: {total_dpo_loss / len(dpo_dataloader)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.nn import MarginRankingLoss\n",
    "\n",
    "# Define the formatting function\n",
    "def format_stackexchange_dpo(samples):\n",
    "    \"\"\"Format StackExchange dataset for DPO.\"\"\"\n",
    "    return {\n",
    "        \"prompt\": [\"Question: \" + question + \"\\n\\nAnswer: \" for question in samples[\"question\"]],\n",
    "        \"chosen\": samples[\"response_j\"],   # Rated better than k\n",
    "        \"rejected\": samples[\"response_k\"], # Rated worse than j\n",
    "    }\n",
    "\n",
    "# Load and format a subset (30%) of the StackExchange DPO dataset\n",
    "dataset = load_dataset(\"lvwerra/stack-exchange-paired\")\n",
    "subset_size = int(0.3 * len(dataset['train']))  # 30% of the dataset\n",
    "subset_indices = torch.randperm(len(dataset['train'])).tolist()[:subset_size]  # Randomly select indices\n",
    "formatted_dataset = dataset['train'].select(subset_indices).map(format_stackexchange_dpo, batched=True, load_from_cache_file=False)\n",
    "\n",
    "# Convert the formatted dataset to DataLoader for batch processing\n",
    "dpo_dataloader = DataLoader(formatted_dataset, batch_size=64, shuffle=True)\n",
    "\n",
    "# Define DPO-specific loss function\n",
    "dpo_loss_function = MarginRankingLoss(margin=1.0)\n",
    "dpo_num_epochs = 2\n",
    "\n",
    "# Initialize the model and optimizer here\n",
    "\n",
    "# DPO Training loop\n",
    "for epoch in range(dpo_num_epochs):\n",
    "    model.train()\n",
    "    total_dpo_loss = 0\n",
    "\n",
    "    for batch in dpo_dataloader:\n",
    "        optimizer.zero_grad()\n",
    "        # Your forward pass and loss calculation here\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        dpo_loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)\n",
    "        optimizer.step()\n",
    "\n",
    "        total_dpo_loss += dpo_loss.item()\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{dpo_num_epochs}, DPO Loss: {total_dpo_loss / len(dpo_dataloader)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assuming DPO dataset is loaded and formatted as dpo_dataset\n",
    "# and model is already fine-tuned\n",
    "from datasets import load_dataset\n",
    "\n",
    "def format_stackexchange_dpo(samples):\n",
    "    \"\"\"Format StackExchange dataset for DPO.\"\"\"\n",
    "    return {\n",
    "        \"prompt\": [\n",
    "            \"Question: \" + question + \"\\n\\nAnswer: \"\n",
    "            for question in samples[\"question\"]\n",
    "        ],\n",
    "        \"chosen\": samples[\"response_j\"],   # Rated better than k\n",
    "        \"rejected\": samples[\"response_k\"], # Rated worse than j\n",
    "    }\n",
    "\n",
    "# Load the StackExchange DPO dataset\n",
    "dataset = load_dataset(\"lvwerra/stack-exchange-paired\")\n",
    "\n",
    "# Format the dataset for DPO\n",
    "formatted_dataset = dataset.map(format_stackexchange_dpo, batched=True)\n",
    "\n",
    "# Display the first few formatted examples (optional)\n",
    "for i in range(3):\n",
    "    print(f\"Sample {i+1}:\")\n",
    "    print(\"Prompt:\", formatted_dataset['train'][i]['prompt'])\n",
    "    print(\"Chosen Answer:\", formatted_dataset['train'][i]['chosen'])\n",
    "    print(\"Rejected Answer:\", formatted_dataset['train'][i]['rejected'])\n",
    "    print()\n",
    "\n",
    "# Define DPO-specific loss function\n",
    "def dpo_loss_function(model_output_preferred, model_output_less_preferred):\n",
    "    # Implement a loss function that penalizes the model\n",
    "    # when preferred response is scored lower than the less-preferred\n",
    "    # For example, you could use margin ranking loss\n",
    "    return loss\n",
    "\n",
    "# DPO Training loop\n",
    "for epoch in range(dpo_num_epochs):\n",
    "    model.train()\n",
    "    total_dpo_loss = 0\n",
    "\n",
    "    for batch in dpo_dataset:\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        prompts = batch['prompt'].to(device)\n",
    "        preferred_responses = batch['chosen'].to(device)\n",
    "        less_preferred_responses = batch['rejected'].to(device)\n",
    "\n",
    "        # Forward pass for both response types\n",
    "        output_preferred = model(preferred_responses)\n",
    "        output_less_preferred = model(less_preferred_responses)\n",
    "\n",
    "        # Compute DPO loss\n",
    "        dpo_loss = dpo_loss_function(output_preferred, output_less_preferred)\n",
    "        total_dpo_loss += dpo_loss.item()\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        dpo_loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)\n",
    "        optimizer.step()\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{dpo_num_epochs}, DPO Loss: {total_dpo_loss / len(dpo_dataset)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_gpu_env_llm",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}