code/plots.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "Plots: Notebook to generate all plots for the paper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "1. Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "2. Indirect indexing task accuracy plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = '/data/final-indirect-idx-logs/'\n",
    "rope_logs = ['skilled-spaceship-655.csv',\n",
    "             'fast-thunder-656.csv',\n",
    "             'bumbling-silence-657.csv'\n",
    "             ]\n",
    "pope_logs = ['wise-terrain-651.csv',\n",
    "            'lunar-fog-653.csv',\n",
    "            'gentle-microwave-654.csv',\n",
    "            ]\n",
    "# extract final accuracy from PoPE exp logs\n",
    "pope_acc = []\n",
    "for log in pope_logs:\n",
    "    log_path = base_dir + log\n",
    "    df = pd.read_csv(log_path)    \n",
    "    pope_acc.append(df[log[:-4] + ' - test/accuracy'].values[-1])\n",
    "# extract final accuracy from RoPE exp logs\n",
    "rope_acc = []\n",
    "for log in rope_logs:\n",
    "    log_path = base_dir + log\n",
    "    df = pd.read_csv(log_path)\n",
    "    rope_acc.append(df[log[:-4] + ' - test/accuracy'].values[-1])\n",
    "\n",
    "data = {'RoPE': rope_acc, 'PoPE': pope_acc}\n",
    "\n",
    "# Calculate means and standard deviations\n",
    "models = list(data.keys())\n",
    "means = [np.mean(data[model]) for model in models]\n",
    "stds = [np.std(data[model], ddof=1) for model in models]  # ddof=1 for sample std\n",
    "\n",
    "print(means, stds)\n",
    "# Create the bar plot\n",
    "fig, ax = plt.subplots(figsize=(2, 4))\n",
    "\n",
    "# Create bars with error bars\n",
    "bars = ax.bar(models, means, yerr=stds, capsize=5, \n",
    "              color=['lightcoral', 'lightgreen'],\n",
    "              alpha=0.8, edgecolor='black', linewidth=1)\n",
    "\n",
    "# Customize the plot\n",
    "ax.set_ylabel('Test Accuracy', fontsize=12)\n",
    "# ax.set_xlabel('Models', fontsize=12)\n",
    "# ax.set_title('Model Accuracy Comparison (mean ± std-dev over 3 seeds)', fontsize=14, fontweight='bold')\n",
    "ax.set_ylim(0, 1.0)  # Assuming accuracy is between 0 and 1\n",
    "ax.grid(axis='y', alpha=0.3)\n",
    "\n",
    "# Add value labels on top of bars\n",
    "# for i, (mean, std) in enumerate(zip(means, stds)):\n",
    "#     ax.text(i, mean + std + 0.01, f'{mean:.3f}±{std:.3f}', \n",
    "#             ha='center', va='bottom', fontsize=10)\n",
    "\n",
    "# Rotate x-axis labels if needed\n",
    "plt.xticks(rotation=0, ha='center')\n",
    "\n",
    "# Adjust layout to prevent label cutoff\n",
    "plt.tight_layout()\n",
    "\n",
    "# Show the plot\n",
    "# plt.show()\n",
    "\n",
    "# Optional: Save the plot\n",
    "plt.savefig(base_dir + 'indirect_idx_acc.png', dpi=300, bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "3. Length Generalization plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_fnames = [ # vanilla RoPE \n",
    "               'gpt2-124M-rope-loss-vs-length.pt',\n",
    "               'gpt2-253M-rope-loss-vs-length.pt',\n",
    "                'gpt2-774M-rope-loss-vs-length.pt',\n",
    "                # YaRN\n",
    "                'ft-yarn-gpt2-124M-rope-loss-vs-length-v3.pt',\n",
    "                'ft-yarn-gpt2-253M-rope-loss-vs-length-v3.pt',\n",
    "                'ft-yarn-gpt2-774M-rope-loss-vs-length-v3.pt',\n",
    "                # vanilla PoPE\n",
    "                'gpt2-124M-pope-loss-vs-length.pt',\n",
    "                'gpt2-253M-pope-loss-vs-length.pt',\n",
    "                'gpt2-774M-pope-loss-vs-length.pt',\n",
    "                # fine-tuned PoPE\n",
    "                'ft-gpt2-124M-pope-loss-vs-length-v3.pt',\n",
    "                'ft-gpt2-253M-pope-loss-vs-length-v3.pt',\n",
    "                'ft-gpt2-774M-pope-loss-vs-length-v3.pt',\n",
    "                ]\n",
    "colors = ['red', 'red', 'red', 'orange', 'orange', 'orange', 'green', 'green', 'green', 'blue', 'blue', 'blue']\n",
    "labels = ['RoPE-124M', 'RoPE-253M', 'RoPE-774M', 'YaRN-124M', 'YaRN-253M', 'YaRN-774M', 'PoPE-124M', 'PoPE-253M', 'PoPE-774M', 'PoPE+ft-124M', 'PoPE+ft-253M', 'PoPE+ft-774M']\n",
    "linestyles = [':', '--', '-', ':', '--', '-', ':', '--', '-', ':', '--', '-']\n",
    "min_overall_perplexity = float('inf')\n",
    "# --- Plotting ---\n",
    "plt.figure(figsize=(8, 5))\n",
    "# Plot perplexity vs length for each model\n",
    "for i in range(len(ckpt_fnames)):\n",
    "    loss_vs_length = torch.load(f'length_gen/final-ckpts-pg19/{ckpt_fnames[i]}', weights_only=False)\n",
    "    # plot perplexity vs length\n",
    "    sorted_items = sorted(loss_vs_length.items())\n",
    "    lengths = [item[0] for item in sorted_items]\n",
    "    losses = [item[1] for item in sorted_items]\n",
    "    perplexities = np.exp(losses)\n",
    "    if len(perplexities) > 0:\n",
    "        min_overall_perplexity = min(min_overall_perplexity, min(perplexities))\n",
    "    plt.plot(lengths, perplexities, label=labels[i], \n",
    "                linestyle=linestyles[i], marker='s', color=colors[i], markersize=5)\n",
    "model_series = ckpt_fnames[i].split('-')[1]\n",
    "plt.title(f'Length extrapolation for models trained on 1024 tokens', fontsize=16)\n",
    "plt.xlabel(\"Inference-time sequence length\", fontsize=12)\n",
    "plt.ylabel(\"Perplexity ($\\leftarrow$)\", fontsize=12)\n",
    "\n",
    "# Set the y-axis to start just below the minimum perplexity for better visibility\n",
    "# if min_overall_perplexity != float('inf'):\n",
    "#     plt.ylim(bottom=max(0, min_overall_perplexity - 5)) # Adjust the offset as needed\n",
    "\n",
    "plt.legend(fontsize=9)\n",
    "plt.grid(True, linestyle=':', alpha=0.6)\n",
    "# plt.show()\n",
    "plt.tight_layout()\n",
    "plt.savefig(f'length_gen/final-ckpts-pg19/length_extrapolate_all_models_yarn.png', dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a pandas DataFrame from the data in your image.\n",
    "# A DataFrame is a 2D labeled data structure, like a spreadsheet.\n",
    "data = {\n",
    "    'Model size': ['124M', '253M', '772M'],\n",
    "    'RoPE': [21.54883, 18.88341, 15.84698],\n",
    "    'PoPE': [21.32518, 18.55186, 15.45872]\n",
    "}\n",
    "df = pd.DataFrame(data)\n",
    "\n",
    "# Set the 'Model size' column as the index of the DataFrame.\n",
    "# This helps in plotting the x-axis labels correctly.\n",
    "df.set_index('Model size', inplace=True)\n",
    "\n",
    "# Define the number of groups and the width for each bar.\n",
    "n_groups = len(df.index)\n",
    "bar_width = 0.15\n",
    "\n",
    "# Set up the figure and axes for the plot.\n",
    "# 'fig' is the whole window or page, and 'ax' is the plot itself.\n",
    "fig, ax = plt.subplots(figsize=(4, 5))\n",
    "\n",
    "index = np.arange(n_groups)\n",
    "\n",
    "# Plot the bars for 'RoPE'.\n",
    "# We shift the bars slightly to the left.\n",
    "rects1 = ax.bar(index - bar_width/2, df['RoPE'], bar_width, label='RoPE', color='lightcoral')\n",
    "\n",
    "# Plot the bars for 'PoPE'.\n",
    "# We shift the bars slightly to the right.\n",
    "rects2 = ax.bar(index + bar_width/2, df['PoPE'], bar_width, label='PoPE', color='lightgreen')\n",
    "\n",
    "# Add labels for the x and y axes.\n",
    "ax.set_xlabel('Model size')\n",
    "ax.set_ylabel('Validation Perplexity')\n",
    "\n",
    "# Add a title to the plot.\n",
    "ax.set_title('Validation perplexity of RoPE v/s PoPE by model size')\n",
    "\n",
    "# Set the x-tick marks and labels to be the model sizes.\n",
    "ax.set_xticks(index)\n",
    "ax.set_xticklabels(df.index)\n",
    "\n",
    "ax.grid(axis='y', alpha=0.3)\n",
    "\n",
    "# Add a legend to distinguish the bars.\n",
    "ax.legend()\n",
    "\n",
    "# Add labels on top of each bar for better readability.\n",
    "# def autolabel(rects):\n",
    "#     \"\"\"Attach a text label above each bar in *rects*, displaying its height.\"\"\"\n",
    "#     for rect in rects:\n",
    "#         height = rect.get_height()\n",
    "#         ax.annotate(f'{height:.2f}',\n",
    "#                     xy=(rect.get_x() + rect.get_width() / 2, height),\n",
    "#                     xytext=(0, 3),  # 3 points vertical offset\n",
    "#                     textcoords=\"offset points\",\n",
    "#                     ha='center', va='bottom')\n",
    "\n",
    "# autolabel(rects1)\n",
    "# autolabel(rects2)\n",
    "\n",
    "# Ensure the layout is tight and clean.\n",
    "fig.tight_layout()\n",
    "\n",
    "# Display the plot.\n",
    "# plt.show()\n",
    "plt.savefig(f'/data/final-owt-ckpts/val_ppl_models.png', dpi=300)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}