main_fig_parsing.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hoVTDs7JPELG"
      },
      "outputs": [],
      "source": [
        "from wandb import Api\n",
        "from collections import defaultdict\n",
        "from tqdm.auto import tqdm\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "from itertools import chain\n",
        "\n",
        "api = Api(api_key=\"API_KEY\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "e5b881e02cee40beaa8675e0147aa034",
            "1c90198c360b48f48f3eae7122274a95",
            "8b7b921acb67421aa5c033a5520ecf6b",
            "ea163fc1bac844d3896a353e23d4d44b",
            "d196e8a2a3e54e7dbaa368afb30284a3",
            "ea9468e4a9694e56ab9bc8b9eeb50f0e",
            "8721d11076a348fc9f78d362b2ccae8d",
            "b2a62c6569dc4037898c1547461731eb",
            "24015665b5be43c885379657e3ffca1e",
            "0e34f25a677f42d0b1a4afd4193041d8",
            "40b23ea880554c26bd91ef357c66c12d"
          ]
        },
        "id": "-D_LUXOaPR19",
        "outputId": "b12452e5-8cf6-44c5-c5b3-b0275f707c17"
      },
      "outputs": [],
      "source": [
        "def get_best_metric_from_run(run, metric: str = \"valid/accuracy\"):\n",
        "    if run.state == \"finished\":\n",
        "      return run.summary[\"max_accuracy\"]\n",
        "    else:\n",
        "      print(f\"Skipping unfinished run {run.id}\")\n",
        "      return None\n",
        "\n",
        "sweep_main_fig = api.sweep(\"entity/project/sweep\")\n",
        "\n",
        "all_runs = list(chain(sweep_main_fig.runs))\n",
        "metrics = defaultdict(list)\n",
        "\n",
        "for run in tqdm(all_runs):\n",
        "    accuracy = get_best_metric_from_run(run)\n",
        "    if accuracy is not None:\n",
        "      metrics['model'].append(run.config['model.sequence_mixer.name'])\n",
        "      metrics['lr'].append(run.config['learning_rate'])\n",
        "      metrics['d_model'].append(run.config['model']['d_model'])\n",
        "      metrics['seq_len'].append(run.config['data']['input_seq_len'])\n",
        "      metrics['accuracy'].append(accuracy)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wYYaoI7N9vmq"
      },
      "outputs": [],
      "source": [
        "data = pd.DataFrame.from_records(metrics)\n",
        "df_with_acc = data.groupby([\"model\", \"d_model\", \"seq_len\", \"lr\"]).mean().reset_index()\n",
        "model_results = df_with_acc.sort_values([\"model\", \"d_model\", \"seq_len\", \"accuracy\"]).drop_duplicates([\"model\", \"d_model\", \"seq_len\"], keep='last')\n",
        "\n",
        "df_with_acc_std = data.groupby([\"model\", \"d_model\", \"seq_len\", \"lr\"]).std().reset_index()\n",
        "model_results_std = df_with_acc_std.iloc[model_results.index]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 211
        },
        "id": "SNqDyeNO-t6O",
        "outputId": "14ab1ce7-ab1d-4bfd-a565-3c04da3a9fba"
      },
      "outputs": [],
      "source": [
        "sns.set_style(\"whitegrid\")\n",
        "\n",
        "fig, axs = plt.subplots(nrows=1, ncols=5, figsize=(12, 2.5), sharey=True)\n",
        "%config InlineBackend.figure_format='retina'\n",
        "\n",
        "\n",
        "axs[0].title.set_text('Sequence Length: 128')\n",
        "axs[0].set_ylabel('Accuracy')\n",
        "\n",
        "axs[1].title.set_text('Sequence Length: 256')\n",
        "axs[2].title.set_text('Sequence Length: 512')\n",
        "axs[3].title.set_text('Sequence Length: 1024')\n",
        "axs[4].title.set_text('Sequence Length: 2048')\n",
        "\n",
        "\n",
        "for i in range(5):\n",
        "  axs[i].set_xticks([0, 1, 2, 3])\n",
        "  axs[i].set_xticklabels([64, 128, 256, 512])\n",
        "\n",
        "markers = [\"o\", \"s\", \"^\", \"+\", \"*\", \"d\", 'H']\n",
        "model_to_name = {\n",
        "    \"based\": \"Based\",\n",
        "    \"rebased\": \"ReBased\",\n",
        "    \"attention\": \"Attention\",\n",
        "    \"conv_attention\": \"СonvAttention\",\n",
        "    \"rwkv\": \"RWKV\",\n",
        "    \"mamba\": \"Mamba\",\n",
        "    \"conv_rwkv\": \"ConvRWKV\"\n",
        "}\n",
        "\n",
        "for i, seq_len in enumerate([128, 256, 512, 1024, 2048]):\n",
        "  for j, model in enumerate([\"rebased\", \"based\", \"rwkv\", \"mamba\", \"conv_attention\"]):\n",
        "    line = model_results[(model_results.seq_len == seq_len) & (model_results.model == model)].accuracy.tolist()\n",
        "    std = model_results_std[(model_results.seq_len == seq_len) & (model_results.model == model)].accuracy.tolist()\n",
        "    up_std = [s if s+l < 1 else 1 - l for l, s in zip(line, std)]\n",
        "    low_std = [s if l-s>0 else l for l, s in zip(line, std)]\n",
        "\n",
        "    x = list(range(len(line)))\n",
        "    axs[i].errorbar(x, line, yerr=(low_std, up_std), marker=markers[j], label=model_to_name[model], capsize=5)\n",
        "\n",
        "\n",
        "fig.tight_layout()\n",
        "handles, labels = plt.gca().get_legend_handles_labels()\n",
        "fig.legend(handles, labels, loc='lower center', ncol=5, bbox_to_anchor=(0.5, -0.15))\n",
        "fig.text(0.5, 0.0, 'Model Dim', ha='center')\n",
        "plt.savefig(\"main_fig.pdf\", bbox_inches=\"tight\")\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "0e34f25a677f42d0b1a4afd4193041d8": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1c90198c360b48f48f3eae7122274a95": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ea9468e4a9694e56ab9bc8b9eeb50f0e",
            "placeholder": "​",
            "style": "IPY_MODEL_8721d11076a348fc9f78d362b2ccae8d",
            "value": "100%"
          }
        },
        "24015665b5be43c885379657e3ffca1e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "40b23ea880554c26bd91ef357c66c12d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "8721d11076a348fc9f78d362b2ccae8d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "8b7b921acb67421aa5c033a5520ecf6b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b2a62c6569dc4037898c1547461731eb",
            "max": 2800,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_24015665b5be43c885379657e3ffca1e",
            "value": 2800
          }
        },
        "b2a62c6569dc4037898c1547461731eb": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d196e8a2a3e54e7dbaa368afb30284a3": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e5b881e02cee40beaa8675e0147aa034": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_1c90198c360b48f48f3eae7122274a95",
              "IPY_MODEL_8b7b921acb67421aa5c033a5520ecf6b",
              "IPY_MODEL_ea163fc1bac844d3896a353e23d4d44b"
            ],
            "layout": "IPY_MODEL_d196e8a2a3e54e7dbaa368afb30284a3"
          }
        },
        "ea163fc1bac844d3896a353e23d4d44b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0e34f25a677f42d0b1a4afd4193041d8",
            "placeholder": "​",
            "style": "IPY_MODEL_40b23ea880554c26bd91ef357c66c12d",
            "value": " 2800/2800 [00:00&lt;00:00, 32016.10it/s]"
          }
        },
        "ea9468e4a9694e56ab9bc8b9eeb50f0e": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}