RadioGPT_3_Finally_RadioGPT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "33EWi1Bq4kM6"
   },
   "source": [
    "![RadioGPT Banner](https://openfileserver.chloelavrat.com/workshops/RadioGPT/assets/radiogpt-banner.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "99Nn_cKD4g_G"
   },
   "source": [
    "> 💡 **PLEASE CONNECT USING A GPU SESSION FOR MORE COMPUTE POWER** :\n",
    ">\n",
    "> `Runtime > Change runtime type > T4 GPU > Save`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "cellView": "form",
    "id": "jU92rWwZ8re3"
   },
   "outputs": [],
   "source": [
    "#@title Initialize the notebook\n",
    "!git clone https://github.com/chloelavrat/RadioGPT.git > /dev/null 2>&1\n",
    "!cd RadioGPT && git checkout clavrat/first-version > /dev/null 2>&1\n",
    "!pip install torch datasets tqdm transformers tiktoken > /dev/null 2>&1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_zehV3Qz5yjo"
   },
   "source": [
    "# 📀 Dataset Overview"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vHMxt75R6JXc",
    "outputId": "bb23afb6-6842-48a4-a8aa-db2f4cde326e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading dataset...\n",
      "Dataset downloaded!\n",
      "Loaded 110368 conversations\n",
      "Maximum sequence length: 64\n",
      "Dataset loaded !\n"
     ]
    }
   ],
   "source": [
    "#@title Select your radio dataset\n",
    "#@markdown Please select your favorit Radio Station ! Then run the cell to load the dataset 🌱\n",
    "import os\n",
    "import subprocess\n",
    "from RadioGPT.gptmodel.core.dataset import AlpacaDataset\n",
    "from RadioGPT.gptmodel.core.utils import download_dataset\n",
    "\n",
    "radio_station = 'France Inter' # @param [\"France Inter\", \"Mouv’\", \"France Culture\"]\n",
    "\n",
    "base_url = \"https://openfileserver.chloelavrat.com/workshops/RadioGPT/dataset/\"\n",
    "block_size = 64\n",
    "\n",
    "if radio_station == 'France Inter':\n",
    "  file = \"inter.json\"\n",
    "  destination = \"dataset/inter.json\"\n",
    "\n",
    "if radio_station == 'Mouv’':\n",
    "  file = \"mouv.json\"\n",
    "  destination = \"dataset/mouv.json\"\n",
    "\n",
    "if radio_station == 'France Culture':\n",
    "  file = \"culture.json\"\n",
    "  destination = \"dataset/culture.json\"\n",
    "\n",
    "download_dataset(base_url+file, destination)\n",
    "dataset = AlpacaDataset(destination, block_size)\n",
    "print(\"Dataset loaded !\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zND_ANv351e3"
   },
   "source": [
    "# 🧠 Model definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-FE4E1o9-cW9",
    "outputId": "859d278b-fa4b-4bb8-a1c4-d1bc5551ef79"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading model...\n",
      "Loading model...\n",
      "Model loaded !\n",
      "You will use your cpu\n",
      "Total number of parameters   83.1M\n"
     ]
    }
   ],
   "source": [
    "#@title Loading ...\n",
    "# Load RadioGPT's checkpoint :)\n",
    "from RadioGPT.gptmodel.core.model import GPTlite\n",
    "from RadioGPT.gptmodel.core.utils import load_model\n",
    "import os, subprocess, torch\n",
    "# get device\n",
    "device = (\n",
    "    torch.device(\"mps\") if torch.backends.mps.is_available() else\n",
    "    torch.device(\"cuda\") if torch.cuda.is_available() else\n",
    "    torch.device(\"cpu\")\n",
    ")\n",
    "\n",
    "# downloading model\n",
    "print(\"Downloading model...\")\n",
    "os.makedirs(\"models\", exist_ok=True)\n",
    "subprocess.run([\"wget\", \"https://openfileserver.chloelavrat.com/workshops/RadioGPT/models/model_gpt_chat_best.pth\", \"-O\", \"models/model_gpt_chat_best.pth\"])\n",
    "\n",
    "def load_model(model_path, device, config):\n",
    "    # Load the model checkpoint\n",
    "    checkpoint = torch.load(\n",
    "        model_path, map_location=device, weights_only=False)\n",
    "\n",
    "    model = GPTlite(config)\n",
    "    model.load_state_dict(checkpoint)\n",
    "\n",
    "    return model\n",
    "\n",
    "# Hyperparameters\n",
    "block_size = 64\n",
    "n_embd = 512\n",
    "n_head = 8\n",
    "n_layer = 10\n",
    "dropout = 0.2\n",
    "\n",
    "config = {\n",
    "    'context_size': block_size,\n",
    "    'vocab_size': dataset.vocab_size,\n",
    "    'embedding_dim': n_embd,\n",
    "    'num_heads': n_head,\n",
    "    'num_layers': n_layer,\n",
    "    'dropout': dropout\n",
    "}\n",
    "\n",
    "# Loading model in memory\n",
    "print(\"Loading model...\")\n",
    "model = load_model(\"models/model_gpt_chat_best.pth\", device, config)\n",
    "model = model.to(device)\n",
    "print(\"Model loaded !\")\n",
    "\n",
    "# Get the total number of parameters in the model\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "\n",
    "print(f\"You will use your {device}\")\n",
    "print(f\"Total number of parameters   {total_params / 1e6:.1f}M\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8OD2BWvy566z"
   },
   "source": [
    "# ⛳️ RadioGPT Training\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "hAf1VVuc-n2U"
   },
   "outputs": [],
   "source": [
    "# training parameters\n",
    "learning_rate = 1e-3\n",
    "epochs = 1000\n",
    "batch_size = 128\n",
    "grad_clip = 0.5\n",
    "eval_every = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "86YPl1zD-oP0"
   },
   "outputs": [],
   "source": [
    "#@title Casual Training Loop\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import OneCycleLR\n",
    "from tqdm import tqdm\n",
    "\n",
    "# --- Setup --- #\n",
    "print(\"Let's fry some eggs!! (your loss should be less than 5, restart cell if not...)\")\n",
    "print(\"-\" * 5)\n",
    "\n",
    "# Initialize scaler for mixed precision training\n",
    "scaler = torch.amp.GradScaler()\n",
    "\n",
    "# Split dataset into training and validation sets\n",
    "train_size = int(0.9 * len(dataset))\n",
    "val_size = len(dataset) - train_size\n",
    "train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n",
    "\n",
    "# Create data loaders\n",
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    num_workers=0,\n",
    "    pin_memory=True\n",
    ")\n",
    "\n",
    "val_loader = DataLoader(\n",
    "    val_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    "    num_workers=0,\n",
    "    pin_memory=True\n",
    ")\n",
    "\n",
    "# Initialize optimizer and learning rate scheduler\n",
    "optimizer = optim.AdamW(\n",
    "    model.parameters(),\n",
    "    lr=learning_rate,\n",
    "    weight_decay=0.1,\n",
    "    betas=(0.9, 0.999)\n",
    ")\n",
    "\n",
    "lr_scheduler = OneCycleLR(\n",
    "    optimizer,\n",
    "    max_lr=learning_rate,\n",
    "    epochs=epochs,\n",
    "    steps_per_epoch=len(train_loader),\n",
    "    pct_start=0.1,\n",
    "    div_factor=25,\n",
    "    final_div_factor=1e4\n",
    ")\n",
    "\n",
    "# Evaluation function\n",
    "def evaluate(model, dataloader):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for x, y in dataloader:\n",
    "            x, y = x.to(device), y.to(device)\n",
    "            with torch.cuda.amp.autocast():\n",
    "                _, loss = model(x, y)\n",
    "            total_loss += loss.item()\n",
    "    return total_loss / len(dataloader)\n",
    "\n",
    "# --- Training Loop --- #\n",
    "best_validation_loss = float('inf')\n",
    "no_improvement_count = 0\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    model.train()\n",
    "    total_train_loss = 0\n",
    "    train_progress_bar = tqdm(train_loader, desc=f\"Epoch {epoch + 1}\")\n",
    "\n",
    "    # Training phase\n",
    "    for batch_idx, (x, y) in enumerate(train_progress_bar):\n",
    "        x, y = x.to(device), y.to(device)\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Forward pass with mixed precision\n",
    "        with torch.cuda.amp.autocast():\n",
    "            logits, loss = model(x, y)\n",
    "\n",
    "        # Backward pass\n",
    "        scaler.scale(loss).backward()\n",
    "\n",
    "        if grad_clip is not None:\n",
    "            scaler.unscale_(optimizer)\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n",
    "\n",
    "        scaler.step(optimizer)\n",
    "        scaler.update()\n",
    "\n",
    "        # Update learning rate scheduler\n",
    "        lr_scheduler.step()\n",
    "\n",
    "        total_train_loss += loss.item()\n",
    "\n",
    "        # Update progress bar\n",
    "        train_progress_bar.set_postfix({\n",
    "            'loss': f'{loss.item():.3f}',\n",
    "            'lr': f'{optimizer.param_groups[0][\"lr\"]:.2e}'\n",
    "        })\n",
    "\n",
    "    # Validation phase every `eval_every` epochs\n",
    "    if (epoch + 1) % eval_every == 0:\n",
    "        avg_train_loss = total_train_loss / len(train_loader)\n",
    "        val_loss = evaluate(model, val_loader)\n",
    "\n",
    "        print(\"-\" * 30)\n",
    "        print(f\"Epoch {epoch + 1}/{epochs}\")\n",
    "        print(f\"Average Train Loss: {avg_train_loss:.4f}\")\n",
    "        print(f\"Validation Loss: {val_loss:.4f}\")\n",
    "        print(f\"Learning Rate: {optimizer.param_groups[0]['lr']:.2e}\")\n",
    "        print(\"-\" * 30)\n",
    "\n",
    "        # Save the best model based on validation loss\n",
    "        if val_loss < best_validation_loss:\n",
    "            best_validation_loss = val_loss\n",
    "            no_improvement_count = 0\n",
    "            torch.save(model.state_dict(), f\"{config['Training']['save_dir']}/best_model.pth\")\n",
    "            print(f\"New best model saved! Loss: {best_validation_loss:.4f}\")\n",
    "        else:\n",
    "            no_improvement_count += 1\n",
    "\n",
    "        # Early stopping condition\n",
    "        if no_improvement_count >= config[\"Training\"][\"patience\"]:\n",
    "            print(f\"No improvement in {config['Training']['patience']} evaluations. Early stopping.\")\n",
    "            break\n",
    "\n",
    "        model.train()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aSlusiv-6FlK"
   },
   "source": [
    "# 🎨 Let's Play"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "TWhgbWHa56Xc"
   },
   "outputs": [],
   "source": [
    "#@title Selected prompt\n",
    "\n",
    "# play with the model\n",
    "def generate_response(model, dataset, prompt, device, max_new_tokens):\n",
    "    model.eval()\n",
    "    # Encode the prompt\n",
    "    input_tensor = dataset.encode(prompt).unsqueeze(0).to(device)\n",
    "\n",
    "    # Generate text using the model's generate method\n",
    "    with torch.no_grad():\n",
    "        generated_indices = model.generate(input_tensor, max_new_tokens)\n",
    "        generated_text = dataset.decode(generated_indices[0].tolist())\n",
    "\n",
    "    # Return only the newly generated part (after the prompt)\n",
    "    return generated_text[len(prompt):]\n",
    "\n",
    "\n",
    "prompt = 'Prépare une recette de pâtes à la carbonara.' # @param [\"Prépare une recette de pâtes à la carbonara.\", \"Quel est l'élément chimique avec le numéro atomique 29 ?\", \"Rédige un court paragraphe sur le thème de l'amitié et de la confiance.\"]\n",
    "max_new_tokens = 106 # @param {type:\"slider\", min:5, max:500, step:1}\n",
    "\n",
    "\n",
    "prompt = f\"Question: {prompt}\\nAnswer:\"\n",
    "prompt = prompt + generate_response(model, dataset, prompt, device, max_new_tokens)\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "Ja7k1zxP_Su9"
   },
   "outputs": [],
   "source": [
    "#@title Let's prompt it!\n",
    "\n",
    "# play with the model\n",
    "def generate_response(model, dataset, prompt, device, max_new_tokens):\n",
    "    model.eval()\n",
    "    # Encode the prompt\n",
    "    input_tensor = dataset.encode(prompt).unsqueeze(0).to(device)\n",
    "\n",
    "    # Generate text using the model's generate method\n",
    "    with torch.no_grad():\n",
    "        generated_indices = model.generate(input_tensor, max_new_tokens)\n",
    "        generated_text = dataset.decode(generated_indices[0].tolist())\n",
    "\n",
    "    # Return only the newly generated part (after the prompt)\n",
    "    return generated_text[len(prompt):]\n",
    "\n",
    "\n",
    "prompt = 'Décris les différences entre le modèle GPT-2 et le modèle GPT-3.' # @param {type:\"string\"}\n",
    "max_new_tokens = 180 # @param {type:\"slider\", min:5, max:500, step:1}\n",
    "\n",
    "\n",
    "prompt = f\"Question: {prompt}\\nAnswer:\"\n",
    "prompt = prompt + generate_response(model, dataset, prompt, device, max_new_tokens)\n",
    "print(prompt)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}