numeraitransformerera_epochs_15_ALL.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "gpuClass": "premium",
      "name": "numeraitransformerera_epochs_15_ALL.ipynb",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "premium"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mglowacki100/AndrewNG-Coursera/blob/master/numeraitransformerera_epochs_15_ALL.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "An attempt to feed entire era as sequence to a Transformer with some tricks around memory efficiency on limited GPU compute of Colab.\n",
        "\n",
        "I tried playing around with the models, not sure if this is an acceptable implementation of Transformers but seem to be working.\n",
        "\n",
        "Twitter: [@parmarsuraj99](https://twitter.com/parmarsuraj99), RC: surajp"
      ],
      "metadata": {
        "id": "U9SwapBsrQEP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "sO7ajUxOsDLE",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "99ffec6c-f23b-42b6-b722-2b8fdc71cad6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/bin/bash: nvidia-smi: command not found\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fr-lyN8PrB1E",
        "outputId": "37f655ba-1d2c-44be-8b20-7c765fb4aa84",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting numerapi\n",
            "  Downloading numerapi-2.14.0-py3-none-any.whl (26 kB)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from numerapi) (2.27.1)\n",
            "Requirement already satisfied: pytz in /usr/local/lib/python3.10/dist-packages (from numerapi) (2022.7.1)\n",
            "Requirement already satisfied: python-dateutil in /usr/local/lib/python3.10/dist-packages (from numerapi) (2.8.2)\n",
            "Requirement already satisfied: tqdm>=4.29.1 in /usr/local/lib/python3.10/dist-packages (from numerapi) (4.65.0)\n",
            "Requirement already satisfied: click>=7.0 in /usr/local/lib/python3.10/dist-packages (from numerapi) (8.1.3)\n",
            "Requirement already satisfied: pandas>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from numerapi) (1.5.3)\n",
            "Requirement already satisfied: numpy>=1.21.0 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.0->numerapi) (1.22.4)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil->numerapi) (1.16.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->numerapi) (1.26.15)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->numerapi) (2022.12.7)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->numerapi) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->numerapi) (3.4)\n",
            "Installing collected packages: numerapi\n",
            "Successfully installed numerapi-2.14.0\n"
          ]
        }
      ],
      "source": [
        "!pip install numerapi"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from numerapi import NumerAPI\n",
        "napi = NumerAPI()\n",
        "\n",
        "napi.download_dataset(\"v4.1/train_int8.parquet\", \"../data/train.parquet\")\n",
        "napi.download_dataset(\"v4.1/validation_int8.parquet\", \"../data/validation.parquet\")\n",
        "napi.download_dataset(\"v4.1/live_int8.parquet\", \"../data/live.parquet\")\n",
        "napi.download_dataset(\"v4.1/live_example_preds.parquet\", \"../data/live_example_preds.parquet\")\n",
        "napi.download_dataset(\"v4.1/validation_example_preds.parquet\", \"../data/validation_example_preds.parquet\")\n",
        "napi.download_dataset(\"v4.1/features.json\", \"../data/features.json\")\n",
        "napi.download_dataset(\"v4.1/meta_model.parquet\", \"../data/meta_model.parquet\")"
      ],
      "metadata": {
        "id": "rfyyhPODrPsK",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 165
        },
        "outputId": "7196ac91-271e-45b7-c210-682cc66aec20"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "../data/train.parquet: 1.45GB [00:55, 26.3MB/s]                            \n",
            "../data/validation.parquet: 1.57GB [00:59, 26.6MB/s]                            \n",
            "../data/live.parquet: 4.47MB [00:00, 7.17MB/s]                            \n",
            "../data/live_example_preds.parquet: 132kB [00:00, 520kB/s]                            \n",
            "../data/validation_example_preds.parquet: 59.1MB [00:02, 21.2MB/s]                            \n",
            "../data/features.json: 703kB [00:00, 1.67MB/s]                          \n",
            "../data/meta_model.parquet: 21.2MB [00:01, 16.9MB/s]                            \n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'../data/meta_model.parquet'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run all"
      ],
      "metadata": {
        "id": "r_rDekOHHXwN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import gc, os\n",
        "import pandas as pd\n",
        "import random\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import json\n",
        "\n",
        "import joblib\n",
        "from joblib import Parallel, delayed\n",
        "\n",
        "from tqdm import tqdm\n",
        "from sklearn.linear_model import LinearRegression\n",
        "from sklearn.decomposition import PCA\n",
        "from sklearn.cluster import KMeans\n",
        "from sklearn.preprocessing import scale, normalize\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "import torch.optim as optim\n",
        "import math\n",
        "\n",
        "from numerapi import NumerAPI\n",
        "\n",
        "tqdm.pandas()"
      ],
      "metadata": {
        "id": "Uwlq0cmvzv4I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# read data from parquet files\n",
        "# to reduce the memory usage on low memory Colab instance, read live data for col info\n",
        "\n",
        "with open(\"../data/features.json\") as f:\n",
        "    features_meta = json.load(f)\n",
        "\n",
        "features_to_load = features_meta[\"feature_sets\"][\"small\"] # medium or all features\n",
        "\n",
        "live = pd.read_parquet(\"../data/live.parquet\")\n",
        "\n",
        "cols_to_ignore = [\n",
        "    c for c in live.columns if (c not in features_to_load and \"feature_\" in c)\n",
        "]\n",
        "cols_to_load = [c for c in live.columns if c not in cols_to_ignore]\n",
        "\n",
        "train = pd.read_parquet(\"../data/train.parquet\", columns=cols_to_load)\n",
        "validation = pd.read_parquet(\"../data/validation.parquet\", columns=cols_to_load)\n",
        "live = pd.read_parquet(\"../data/live.parquet\", columns=cols_to_load)\n",
        "\n",
        "live_example_preds = pd.read_parquet(\"../data/live_example_preds.parquet\")\n",
        "validation_example_preds = pd.read_parquet(\"../data/validation_example_preds.parquet\")\n",
        "meta_model = pd.read_parquet(\"../data/meta_model.parquet\")\n",
        "\n",
        "feature_names = [f for f in train.columns if \"feature_\" in f]\n",
        "feature_names = [f for f in feature_names if f in features_to_load]\n",
        "target_names = [t for t in train.columns if \"target_\" in t]#[:5]\n",
        "\n",
        "TARGET_NAME = target_names[0]\n",
        "PREDICTION_NAME = \"prediction\"\n",
        "\n",
        "# all eras of TARGET_NAME must be present in diagnostics\n",
        "validation = validation.dropna(subset=[TARGET_NAME], axis=0).copy()\n",
        "gc.collect()\n",
        "\n",
        "train[\"era_int\"] = train[\"era\"].astype(int)\n",
        "validation[\"era_int\"] = validation[\"era\"].astype(int)\n",
        "gc.collect()\n",
        "\n",
        "train[feature_names] = train[feature_names].fillna(2)\n",
        "validation[feature_names] = validation[feature_names].fillna(2)\n",
        "gc.collect()\n",
        "\n",
        "# use a better method to handle NaN in targets\n",
        "train[target_names] = train[target_names].fillna(0.5)\n",
        "validation[target_names] = validation[target_names].fillna(0.5)\n",
        "gc.collect()\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "# device = \"cpu\"\n"
      ],
      "metadata": {
        "id": "qTw5N8p30zq4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train.columns"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PejDRLKQ3PyA",
        "outputId": "3a9a2d6f-a377-49a8-c7d1-62cd63ce75eb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Index(['era', 'data_type', 'feature_froggier_unlearned_underworkman',\n",
              "       'feature_fribble_gusseted_stickjaw',\n",
              "       'feature_unswaddled_inenarrable_goody', 'feature_esculent_erotic_epoxy',\n",
              "       'feature_coraciiform_sciurine_reef',\n",
              "       'feature_burning_phrygian_axinomancy',\n",
              "       'feature_bijou_penetrant_syringa', 'feature_subfusc_furriest_nervule',\n",
              "       'feature_ugrian_schizocarpic_skulk',\n",
              "       'feature_unpainted_censual_pinacoid',\n",
              "       'feature_tragical_rainbowy_seafarer', 'feature_sodding_choosy_eruption',\n",
              "       'feature_massive_demisable_spouse', 'feature_unventilated_sollar_bason',\n",
              "       'feature_godliest_consistorian_woodpecker',\n",
              "       'feature_undrilled_wheezier_countermand',\n",
              "       'feature_corporatist_seborrheic_hopi',\n",
              "       'feature_undisguised_unenviable_stamen',\n",
              "       'feature_fearsome_merry_bluewing',\n",
              "       'feature_septuple_bonapartean_sanbenito',\n",
              "       'feature_guardian_frore_rolling',\n",
              "       'feature_distressed_bloated_disquietude',\n",
              "       'feature_unsizable_ancestral_collocutor',\n",
              "       'feature_ecstatic_foundational_crinoidea',\n",
              "       'feature_entopic_interpreted_subsidiary',\n",
              "       'feature_just_flavescent_draff', 'feature_elaborate_intimate_bor',\n",
              "       'feature_iffy_pretty_gumming', 'feature_piping_geotactic_cusp',\n",
              "       'feature_cyclopedic_maestoso_daguerreotypist',\n",
              "       'feature_unreproved_cultish_glioma',\n",
              "       'feature_mancunian_stalky_charmeuse', 'target', 'target_nomi_v4_20',\n",
              "       'target_nomi_v4_60', 'target_tyler_v4_20', 'target_tyler_v4_60',\n",
              "       'target_victor_v4_20', 'target_victor_v4_60', 'target_ralph_v4_20',\n",
              "       'target_ralph_v4_60', 'target_waldo_v4_20', 'target_waldo_v4_60',\n",
              "       'target_jerome_v4_20', 'target_jerome_v4_60', 'target_janet_v4_20',\n",
              "       'target_janet_v4_60', 'target_ben_v4_20', 'target_ben_v4_60',\n",
              "       'target_alan_v4_20', 'target_alan_v4_60', 'target_paul_v4_20',\n",
              "       'target_paul_v4_60', 'target_george_v4_20', 'target_george_v4_60',\n",
              "       'target_william_v4_20', 'target_william_v4_60', 'target_arthur_v4_20',\n",
              "       'target_arthur_v4_60', 'target_thomas_v4_20', 'target_thomas_v4_60',\n",
              "       'target_cyrus_v4_20', 'target_cyrus_v4_60', 'target_caroline_v4_20',\n",
              "       'target_caroline_v4_60', 'target_sam_v4_20', 'target_sam_v4_60',\n",
              "       'target_xerxes_v4_20', 'target_xerxes_v4_60', 'era_int'],\n",
              "      dtype='object')"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np"
      ],
      "metadata": {
        "id": "A-gU-gAZQkSU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for df in [train, validation, live]:\n",
        "  feature_columns = df.filter(regex='^feature_').columns.tolist()\n",
        "  for f in feature_columns:\n",
        "    df[f\"{f}_A\"] = np.where(df[f] == 2,  2, np.where(df[f] > 2, 4, 0))\n",
        "    df[f\"{f}_B\"] = np.where(df[f] == 2,  2, np.where(df[f].isin([0,4]), 4, 0))\n",
        "  #df = df.drop(columns=feature_columns)\n"
      ],
      "metadata": {
        "id": "InOLrd_xOdPi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#train = train.drop(columns=feature_columns)"
      ],
      "metadata": {
        "id": "toUHv12mSvQW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#validation = validation.drop(columns=feature_columns)\n",
        "#live = live.drop(columns=feature_columns)"
      ],
      "metadata": {
        "id": "ExKTc5tzS3oe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "list(train.columns)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WNTHfvM9OdUl",
        "outputId": "bac79b46-54f1-488c-dc0f-f5a60656d4cc"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['era',\n",
              " 'data_type',\n",
              " 'target',\n",
              " 'target_nomi_v4_20',\n",
              " 'target_nomi_v4_60',\n",
              " 'target_tyler_v4_20',\n",
              " 'target_tyler_v4_60',\n",
              " 'target_victor_v4_20',\n",
              " 'target_victor_v4_60',\n",
              " 'target_ralph_v4_20',\n",
              " 'target_ralph_v4_60',\n",
              " 'target_waldo_v4_20',\n",
              " 'target_waldo_v4_60',\n",
              " 'target_jerome_v4_20',\n",
              " 'target_jerome_v4_60',\n",
              " 'target_janet_v4_20',\n",
              " 'target_janet_v4_60',\n",
              " 'target_ben_v4_20',\n",
              " 'target_ben_v4_60',\n",
              " 'target_alan_v4_20',\n",
              " 'target_alan_v4_60',\n",
              " 'target_paul_v4_20',\n",
              " 'target_paul_v4_60',\n",
              " 'target_george_v4_20',\n",
              " 'target_george_v4_60',\n",
              " 'target_william_v4_20',\n",
              " 'target_william_v4_60',\n",
              " 'target_arthur_v4_20',\n",
              " 'target_arthur_v4_60',\n",
              " 'target_thomas_v4_20',\n",
              " 'target_thomas_v4_60',\n",
              " 'target_cyrus_v4_20',\n",
              " 'target_cyrus_v4_60',\n",
              " 'target_caroline_v4_20',\n",
              " 'target_caroline_v4_60',\n",
              " 'target_sam_v4_20',\n",
              " 'target_sam_v4_60',\n",
              " 'target_xerxes_v4_20',\n",
              " 'target_xerxes_v4_60',\n",
              " 'era_int',\n",
              " 'feature_froggier_unlearned_underworkman_A',\n",
              " 'feature_froggier_unlearned_underworkman_B',\n",
              " 'feature_fribble_gusseted_stickjaw_A',\n",
              " 'feature_fribble_gusseted_stickjaw_B',\n",
              " 'feature_unswaddled_inenarrable_goody_A',\n",
              " 'feature_unswaddled_inenarrable_goody_B',\n",
              " 'feature_esculent_erotic_epoxy_A',\n",
              " 'feature_esculent_erotic_epoxy_B',\n",
              " 'feature_coraciiform_sciurine_reef_A',\n",
              " 'feature_coraciiform_sciurine_reef_B',\n",
              " 'feature_burning_phrygian_axinomancy_A',\n",
              " 'feature_burning_phrygian_axinomancy_B',\n",
              " 'feature_bijou_penetrant_syringa_A',\n",
              " 'feature_bijou_penetrant_syringa_B',\n",
              " 'feature_subfusc_furriest_nervule_A',\n",
              " 'feature_subfusc_furriest_nervule_B',\n",
              " 'feature_ugrian_schizocarpic_skulk_A',\n",
              " 'feature_ugrian_schizocarpic_skulk_B',\n",
              " 'feature_unpainted_censual_pinacoid_A',\n",
              " 'feature_unpainted_censual_pinacoid_B',\n",
              " 'feature_tragical_rainbowy_seafarer_A',\n",
              " 'feature_tragical_rainbowy_seafarer_B',\n",
              " 'feature_sodding_choosy_eruption_A',\n",
              " 'feature_sodding_choosy_eruption_B',\n",
              " 'feature_massive_demisable_spouse_A',\n",
              " 'feature_massive_demisable_spouse_B',\n",
              " 'feature_unventilated_sollar_bason_A',\n",
              " 'feature_unventilated_sollar_bason_B',\n",
              " 'feature_godliest_consistorian_woodpecker_A',\n",
              " 'feature_godliest_consistorian_woodpecker_B',\n",
              " 'feature_undrilled_wheezier_countermand_A',\n",
              " 'feature_undrilled_wheezier_countermand_B',\n",
              " 'feature_corporatist_seborrheic_hopi_A',\n",
              " 'feature_corporatist_seborrheic_hopi_B',\n",
              " 'feature_undisguised_unenviable_stamen_A',\n",
              " 'feature_undisguised_unenviable_stamen_B',\n",
              " 'feature_fearsome_merry_bluewing_A',\n",
              " 'feature_fearsome_merry_bluewing_B',\n",
              " 'feature_septuple_bonapartean_sanbenito_A',\n",
              " 'feature_septuple_bonapartean_sanbenito_B',\n",
              " 'feature_guardian_frore_rolling_A',\n",
              " 'feature_guardian_frore_rolling_B',\n",
              " 'feature_distressed_bloated_disquietude_A',\n",
              " 'feature_distressed_bloated_disquietude_B',\n",
              " 'feature_unsizable_ancestral_collocutor_A',\n",
              " 'feature_unsizable_ancestral_collocutor_B',\n",
              " 'feature_ecstatic_foundational_crinoidea_A',\n",
              " 'feature_ecstatic_foundational_crinoidea_B',\n",
              " 'feature_entopic_interpreted_subsidiary_A',\n",
              " 'feature_entopic_interpreted_subsidiary_B',\n",
              " 'feature_just_flavescent_draff_A',\n",
              " 'feature_just_flavescent_draff_B',\n",
              " 'feature_elaborate_intimate_bor_A',\n",
              " 'feature_elaborate_intimate_bor_B',\n",
              " 'feature_iffy_pretty_gumming_A',\n",
              " 'feature_iffy_pretty_gumming_B',\n",
              " 'feature_piping_geotactic_cusp_A',\n",
              " 'feature_piping_geotactic_cusp_B',\n",
              " 'feature_cyclopedic_maestoso_daguerreotypist_A',\n",
              " 'feature_cyclopedic_maestoso_daguerreotypist_B',\n",
              " 'feature_unreproved_cultish_glioma_A',\n",
              " 'feature_unreproved_cultish_glioma_B',\n",
              " 'feature_mancunian_stalky_charmeuse_A',\n",
              " 'feature_mancunian_stalky_charmeuse_B']"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "feature_names = train.filter(regex='^feature_').columns.tolist() #MG redefinition\n"
      ],
      "metadata": {
        "id": "1Bk5OUXl3Pvu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "pPYkhv3j3PbD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "PADDING_VALUE = -1 \n",
        "MAX_LEN = 6000\n",
        "FEATURE_DIM = len(feature_names)\n",
        "HIDDEN_DIM = 128\n",
        "OUTPUT_DIM = len(target_names)\n",
        "NUM_HEADS = 2\n",
        "NUM_LAYERS = 2"
      ],
      "metadata": {
        "id": "uoeXSuIxzF_u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# caching\n",
        "\n",
        "def pad_sequence(inputs, padding_value=-1, max_len=None):\n",
        "    if max_len is None:\n",
        "        max_len = max([input.shape[0] for input in inputs])\n",
        "    padded_inputs = []\n",
        "    masks = []\n",
        "    for input in inputs:\n",
        "        pad_len = max_len - input.shape[0]\n",
        "        padded_input = F.pad(input, (0, 0, 0, pad_len), value=padding_value)\n",
        "        mask = torch.ones((input.shape[0], 1), dtype=torch.float)\n",
        "        masks.append(\n",
        "            torch.cat((mask, torch.zeros((pad_len, 1), dtype=torch.float)), dim=0)\n",
        "        )\n",
        "        padded_inputs.append(padded_input)\n",
        "    return torch.stack(padded_inputs), torch.stack(masks)\n",
        "\n",
        "def convert_to_torch(era, data):\n",
        "\n",
        "    inputs = torch.from_numpy(\n",
        "                data[feature_names].values.astype(np.int8))\n",
        "    labels = torch.from_numpy(\n",
        "                data[target_names].values.astype(np.float32))\n",
        "    \n",
        "    padded_inputs, masks_inputs = pad_sequence(\n",
        "            [inputs], padding_value=PADDING_VALUE, max_len=MAX_LEN)\n",
        "    padded_labels, masks_labels = pad_sequence(\n",
        "            [labels], padding_value=PADDING_VALUE, max_len=MAX_LEN)\n",
        "\n",
        "    return {\n",
        "        era: (\n",
        "            padded_inputs,\n",
        "            padded_labels,\n",
        "            masks_inputs\n",
        "        )\n",
        "    }\n",
        "\n",
        "def get_era2data(df):\n",
        "    res = Parallel(n_jobs=-1, prefer=\"threads\")(\n",
        "        delayed(convert_to_torch)(era, data)\n",
        "        for era, data in tqdm(df.groupby(\"era_int\")))\n",
        "    era2data = {}\n",
        "    for r in tqdm(res):\n",
        "        era2data.update(r)\n",
        "    return era2data\n",
        "\n",
        "era2data_train = get_era2data(train)\n",
        "era2data_validation = get_era2data(validation)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0IfvuZc5zH3S",
        "outputId": "00f9e25e-cf2c-46b7-9728-ca752467cd8a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 574/574 [00:03<00:00, 168.15it/s]\n",
            "100%|██████████| 574/574 [00:00<00:00, 821402.42it/s]\n",
            "100%|██████████| 486/486 [00:09<00:00, 49.23it/s] \n",
            "100%|██████████| 486/486 [00:00<00:00, 779813.21it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Use better attention mechanism that doesn't give OOM\n",
        "\n",
        "\n",
        "THis might help\n",
        "- https://github.com/lucidrains/linear-attention-transformer"
      ],
      "metadata": {
        "id": "u1LmbrzVG6qp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "    \"\"\"\n",
        "    Didn't have much difference\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, d_model, max_len=5000):\n",
        "        super(PositionalEncoding, self).__init__()\n",
        "        self.dropout = nn.Dropout(p=0.1)\n",
        "\n",
        "        pe = torch.zeros(max_len, d_model)\n",
        "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
        "        div_term = torch.exp(\n",
        "            torch.arange(0, d_model, 2).float()\n",
        "            * (-torch.log(torch.tensor(10000.0)) / d_model)\n",
        "        )\n",
        "        pe[:, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 1::2] = torch.cos(position * div_term)\n",
        "        pe = pe.unsqueeze(0).transpose(0, 1)\n",
        "        self.register_buffer(\"pe\", pe)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.pe[: x.size(0), :]\n",
        "        return self.dropout(x)\n",
        "\n",
        "class LinearAttentionWithKQV(nn.Module):\n",
        "    \"\"\"\n",
        "    Light on compute. MAX_LEN x D_MODEL; \n",
        "    use this if OOM\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, d_model, dropout=0.1):\n",
        "        super(LinearAttentionWithKQV, self).__init__()\n",
        "        self.linear_k = nn.Linear(d_model, d_model)\n",
        "        self.linear_q = nn.Linear(d_model, d_model)\n",
        "        self.linear_v = nn.Linear(d_model, d_model)\n",
        "        self.dim = d_model\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, inputs, mask=None):\n",
        "        k = self.linear_k(inputs)\n",
        "        q = self.linear_q(inputs)\n",
        "        v = self.linear_v(inputs)\n",
        "        n = torch.sqrt(torch.tensor(self.dim, dtype=torch.float32))\n",
        "\n",
        "        scores = torch.bmm(k.transpose(1, 2), v) / n\n",
        "        scores = self.dropout(scores)\n",
        "\n",
        "        # print(scores.shape, mask.shape)\n",
        "        if mask is not None:\n",
        "            q = q.masked_fill(mask == 0, float(\"-inf\"))\n",
        "        \n",
        "        q = torch.softmax(q, dim=1)\n",
        "\n",
        "        attention_weights = torch.bmm(q, scores) / n\n",
        "\n",
        "        return attention_weights\n",
        "\n",
        "\n",
        "class FeedForwardLayer(nn.Module):\n",
        "    def __init__(self, d_model, d_ff=128, dropout=0.1):\n",
        "        super().__init__()\n",
        "        self.linear_1 = nn.Linear(d_model, d_ff)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.linear_2 = nn.Linear(d_ff, d_model)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.dropout(F.relu(self.linear_1(x)))\n",
        "        x = self.linear_2(x)\n",
        "        return x\n",
        "\n",
        "class SelfAttention(nn.Module):\n",
        "    \"\"\"\n",
        "    Heavy on compute. Vanilla attention MAXLEN x MAXLEN\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, d_model, dropout=0.1):\n",
        "        super(SelfAttention, self).__init__()\n",
        "        self.linear_k = nn.Linear(d_model, d_model)\n",
        "        self.linear_q = nn.Linear(d_model, d_model)\n",
        "        self.linear_v = nn.Linear(d_model, d_model)\n",
        "        self.dim = d_model\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, inputs, mask=None):\n",
        "        k = self.linear_k(inputs)\n",
        "        q = self.linear_q(inputs)\n",
        "        v = self.linear_v(inputs)\n",
        "        n = torch.sqrt(torch.tensor(self.dim, dtype=torch.float32))\n",
        "\n",
        "        scores = torch.bmm(q, k.transpose(1, 2)) / n\n",
        "        # print(scores.shape)\n",
        "        if mask is not None:\n",
        "            scores = scores .masked_fill(mask == 0, float(\"-inf\"))\n",
        "\n",
        "        attention_weights = torch.softmax(scores, dim=1)\n",
        "        attention_weights = self.dropout(attention_weights)\n",
        "\n",
        "        attention_weights = torch.bmm(attention_weights, v) / n\n",
        "        \n",
        "        return attention_weights\n",
        "\n",
        "class MultiHeadLinearAttention(nn.Module):\n",
        "    def __init__(self, d_model, num_heads):\n",
        "        super(MultiHeadLinearAttention, self).__init__()\n",
        "        self.d_model = d_model\n",
        "        self.num_heads = num_heads\n",
        "        #self.attention = LinearAttentionWithKQV(d_model, dropout=0.1)  # drop-in attention\n",
        "        self.attention = SelfAttention(d_model, dropout=0.15)\n",
        "        self.layers = nn.ModuleList(\n",
        "            [nn.Linear(d_model, d_model) for _ in range(num_heads)]\n",
        "        )\n",
        "        self.fc = nn.Linear(num_heads * d_model, d_model)\n",
        "\n",
        "    def forward(self, inputs, mask=None):\n",
        "        x = inputs\n",
        "        head_outputs = []\n",
        "        for layer in range(self.num_heads):\n",
        "            attention_weights = self.attention(x, mask)\n",
        "            head_output = x * attention_weights\n",
        "            head_output = self.layers[layer](head_output)\n",
        "            head_output = F.relu(head_output)\n",
        "            head_outputs.append(head_output)\n",
        "\n",
        "        concatenated = torch.cat(head_outputs, dim=-1)\n",
        "        output = self.fc(concatenated)\n",
        "\n",
        "        return output\n",
        "\n",
        "\n",
        "class TransformerEncoder(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        input_dim,\n",
        "        d_model,\n",
        "        output_dim,\n",
        "        num_heads,\n",
        "        num_layers,\n",
        "        dropout_prob=0.15,\n",
        "        max_len=5000,\n",
        "    ):\n",
        "        super(TransformerEncoder, self).__init__()\n",
        "        self.input_dim = input_dim\n",
        "        self.output_dim = output_dim\n",
        "        self.num_heads = num_heads\n",
        "        self.num_layers = num_layers\n",
        "        self.dropout_prob = dropout_prob\n",
        "        self.d_model = d_model\n",
        "\n",
        "        self.positional_encoding = PositionalEncoding(d_model, max_len)\n",
        "        self.attention = MultiHeadLinearAttention(d_model, num_heads)\n",
        "        self.fc = nn.Sequential(\n",
        "            nn.Linear(d_model, d_model),\n",
        "        )\n",
        "        self.layers = nn.ModuleList(\n",
        "            [\n",
        "                nn.Sequential(FeedForwardLayer(d_model=d_model))\n",
        "                for _ in range(num_layers)\n",
        "            ]\n",
        "        )\n",
        "        self.mapper = nn.Sequential(\n",
        "            nn.Linear(input_dim, d_model), nn.Linear(d_model, d_model)\n",
        "        )\n",
        "\n",
        "    def forward(self, inputs, mask=None):\n",
        "        x = self.mapper(inputs)\n",
        "        pe = self.positional_encoding(x)\n",
        "        x = x + pe # works without PE as well\n",
        "        for layer in range(self.num_layers):\n",
        "            attention_weights = self.attention(x, mask)\n",
        "            x = x + attention_weights\n",
        "            x = F.layer_norm(x, x.shape[1:])\n",
        "            x = F.dropout(x, p=self.dropout_prob)\n",
        "\n",
        "            op = self.layers[layer](x)\n",
        "            x = x + op\n",
        "            x = F.layer_norm(x, x.shape[1:])\n",
        "            x = F.dropout(x, p=self.dropout_prob)\n",
        "\n",
        "        outputs = self.fc(x)\n",
        "        return outputs\n",
        "\n",
        "\n",
        "class Transformer(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        input_dim,\n",
        "        d_model,\n",
        "        output_dim,\n",
        "        num_heads,\n",
        "        num_layers,\n",
        "        dropout_prob=0.15,\n",
        "        max_len=5000,\n",
        "    ):\n",
        "        super(Transformer, self).__init__()\n",
        "        self.input_dim = input_dim\n",
        "        self.output_dim = output_dim\n",
        "        self.num_heads = num_heads\n",
        "        self.num_layers = num_layers\n",
        "        self.dropout_prob = dropout_prob\n",
        "        self.d_model = d_model\n",
        "\n",
        "        self.encoder = TransformerEncoder(\n",
        "            input_dim=input_dim,\n",
        "            d_model=d_model,\n",
        "            output_dim=output_dim,\n",
        "            num_heads=num_heads,\n",
        "            num_layers=num_layers,\n",
        "            max_len=max_len,\n",
        "        )\n",
        "        self.fc = nn.Sequential(\n",
        "            nn.Linear(d_model, d_model // 2),\n",
        "            nn.SELU(),\n",
        "            nn.Linear(d_model // 2, self.output_dim),\n",
        "            nn.Sigmoid(),\n",
        "        )\n",
        "\n",
        "    def forward(self, inputs, mask=None):\n",
        "        emb = self.encoder(inputs, mask)\n",
        "        outputs = self.fc(emb)\n",
        "        return outputs\n",
        "\n",
        "\n",
        "def test_model():\n",
        "\n",
        "    inputs = [\n",
        "        torch.randint(0, 4, (5, FEATURE_DIM)).float(),\n",
        "        torch.randint(0, 4, (3, FEATURE_DIM)).float(),\n",
        "    ]\n",
        "    labels = [\n",
        "        torch.randint(0, 2, (5, OUTPUT_DIM)).float(),\n",
        "        torch.randint(0, 2, (3, OUTPUT_DIM)).float(),\n",
        "    ]\n",
        "\n",
        "    padded_inputs, masks_inputs = pad_sequence(inputs, padding_value=0, max_len=MAX_LEN)\n",
        "    padded_labels, masks_labels = pad_sequence(labels, padding_value=0, max_len=MAX_LEN)\n",
        "\n",
        "    transformer = Transformer(\n",
        "        input_dim=FEATURE_DIM,\n",
        "        d_model=HIDDEN_DIM,\n",
        "        output_dim=OUTPUT_DIM,\n",
        "        num_heads=NUM_HEADS,\n",
        "        num_layers=NUM_LAYERS,\n",
        "        max_len=MAX_LEN,\n",
        "    )\n",
        "\n",
        "    with torch.no_grad():\n",
        "        outputs = transformer(padded_inputs, masks_inputs)\n",
        "\n",
        "    assert torch.isnan(outputs).sum() == 0\n",
        "    assert outputs.shape[:2] == padded_inputs.shape[:2]\n",
        "    assert outputs.shape[-1] == len(target_names)\n",
        "\n",
        "    print(\"Input Shape:\", padded_inputs.shape)\n",
        "    print(\"Output Shape:\", outputs.shape)\n",
        "\n",
        "    del transformer\n",
        "    del inputs, labels\n",
        "    del padded_inputs, masks_inputs, padded_labels, masks_labels\n",
        "    del outputs\n",
        "\n",
        "    gc.collect()\n",
        "\n",
        "test_model()"
      ],
      "metadata": {
        "id": "vdB3mQO62yhp",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "28688b85-2b7d-4c02-ef23-36519417b7a4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input Shape: torch.Size([2, 6000, 64])\n",
            "Output Shape: torch.Size([2, 6000, 36])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# pearsonr in torch differentiable\n",
        "def pearsonr(x, y):\n",
        "    mx = x.mean()\n",
        "    my = y.mean()\n",
        "    xm, ym = x - mx, y - my\n",
        "    r_num = torch.sum(xm * ym)\n",
        "    r_den = torch.sqrt(torch.sum(xm ** 2) * torch.sum(ym ** 2))\n",
        "    r = r_num / r_den\n",
        "    return r"
      ],
      "metadata": {
        "id": "kNkeZDr12ylJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def calculate_loss(outputs, padded_labels, masks_inputs, padded_inputs=None, target_weight_softmax=None):\n",
        "\n",
        "    # MSE on all targets; additionally, on primary target\n",
        "    if target_weight_softmax is not None:\n",
        "        _mse = criterion(\n",
        "            outputs * masks_inputs * target_weight_softmax, \n",
        "            padded_labels * masks_inputs * target_weight_softmax\n",
        "        ) * 0.1\n",
        "\n",
        "    else:\n",
        "        _mse = criterion(outputs * masks_inputs, padded_labels * masks_inputs) * 0.1\n",
        "\n",
        "    _mse += criterion(outputs[:, 0] * masks_inputs, padded_labels[:, 0] * masks_inputs)\n",
        "\n",
        "    # Corr with only primary target; adjust as needed\n",
        "    corr = pearsonr(\n",
        "        outputs[0][:, 0][masks_inputs.view(-1).nonzero()].view(-1, 1),\n",
        "        padded_labels[0][:, 0][masks_inputs.view(-1).nonzero()].view(-1, 1),\n",
        "    )\n",
        "\n",
        "    loss = _mse - corr #+ some_complex_constraints\n",
        "    return loss, _mse, corr\n",
        "\n",
        "\n",
        "# Training loop\n",
        "def train_on_batch(transformer, optimizer, batch):\n",
        "\n",
        "    padded_inputs = batch[0].to(device=device)\n",
        "    padded_labels = batch[1].to(device=device)\n",
        "    masks_inputs = batch[2].to(device=device)\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    outputs = transformer(padded_inputs / 4.0, masks_inputs)\n",
        "\n",
        "    target_weight_softmax = None\n",
        "    #random_weights = torch.rand(padded_labels.shape[-1], device=device)\n",
        "    #target_weight_softmax = F.softmax(random_weights)\n",
        "\n",
        "    loss, _mse, _corr = calculate_loss(outputs, padded_labels, masks_inputs, target_weight_softmax=target_weight_softmax)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    return loss.item(), _mse.item(), _corr.item()\n",
        "\n",
        "\n",
        "def evaluate_on_batch(transformer, batch):\n",
        "\n",
        "    padded_inputs = batch[0].to(device=device)\n",
        "    padded_labels = batch[1].to(device=device)\n",
        "    masks_inputs = batch[2].to(device=device)\n",
        "\n",
        "    outputs = transformer(padded_inputs / 4.0, masks_inputs)\n",
        "    loss, _mse, _corr = calculate_loss(outputs, padded_labels, masks_inputs)\n",
        "    return loss.item(), _mse.item(), _corr.item()\n",
        "\n",
        "\n",
        "def train_model(transformer, optimizer, num_epochs, train_loader, val_loader):\n",
        "    for epoch in range(num_epochs):\n",
        "        total_loss = []\n",
        "        total_corr = []\n",
        "        print(f\"\\nEPOCH: {epoch+1}/{num_epochs}\")\n",
        "        for era_num in tqdm(train_loader):\n",
        "            batch = train_loader[era_num]\n",
        "            loss, _mse, _corr = train_on_batch(transformer, optimizer, batch)\n",
        "            total_loss.append(loss)\n",
        "            total_corr.append(_corr)\n",
        "        print(\n",
        "            f\"Train Loss: {np.mean(total_loss):.4f} | Train Corr: {np.mean(total_corr):.4f}\"\n",
        "        )\n",
        "\n",
        "        transformer.eval()\n",
        "        with torch.no_grad():\n",
        "            \n",
        "            total_loss = []\n",
        "            total_corr = []\n",
        "            for era_num in tqdm(val_loader):\n",
        "                batch = val_loader[era_num]\n",
        "                loss, _mse, _corr = evaluate_on_batch(transformer, batch)\n",
        "                total_loss.append(loss)\n",
        "                total_corr.append(_corr)\n",
        "            print(\n",
        "                f\"Val Loss: {np.mean(total_loss):.4f} | Val Corr: {np.mean(total_corr):.4f}\"\n",
        "            )\n",
        "\n",
        "        torch.cuda.empty_cache()\n",
        "        _ = gc.collect()\n",
        "    \n",
        "    torch.save(transformer.state_dict(), f\"transformer.pth\")\n",
        "\n",
        "    return transformer\n",
        "\n",
        "# Define the model, loss function, and optimizer\n",
        "\n",
        "gc.collect()\n",
        "transformer = Transformer(\n",
        "    input_dim=FEATURE_DIM,\n",
        "    d_model=HIDDEN_DIM,\n",
        "    output_dim=OUTPUT_DIM,\n",
        "    num_heads=NUM_HEADS,\n",
        "    num_layers=NUM_LAYERS,\n",
        ")\n",
        "\n",
        "# load model from checkpoint\n",
        "if \"transformer.pth\" in os.listdir():\n",
        "    transformer.load_state_dict(torch.load(\"transformer.pth\"))\n",
        "\n",
        "transformer.to(device=device)\n",
        "criterion = nn.MSELoss()\n",
        "optimizer = optim.Adam(transformer.parameters(), lr=1e-4)\n",
        "\n",
        "# Number of training iterations\n",
        "# Train for longer with low LR\n",
        "\n",
        "num_epochs = 15\n",
        "\n",
        "transformer = train_model(transformer, optimizer, num_epochs, era2data_train, era2data_validation)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 417
        },
        "id": "PblNRqlXXArI",
        "outputId": "44d07d90-51cb-41a6-97e3-31d0ee26875a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "EPOCH: 1/10\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 5/574 [01:15<2:22:31, 15.03s/it]\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-17-d5c03ea66df9>\u001b[0m in \u001b[0;36m<cell line: 116>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    114\u001b[0m \u001b[0mnum_epochs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m10\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 116\u001b[0;31m \u001b[0mtransformer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtransformer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_epochs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mera2data_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mera2data_validation\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-17-d5c03ea66df9>\u001b[0m in \u001b[0;36mtrain_model\u001b[0;34m(transformer, optimizer, num_epochs, train_loader, val_loader)\u001b[0m\n\u001b[1;32m     62\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mera_num\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtqdm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     63\u001b[0m             \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_loader\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mera_num\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 64\u001b[0;31m             \u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_mse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_corr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_on_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtransformer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     65\u001b[0m             \u001b[0mtotal_loss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     66\u001b[0m             \u001b[0mtotal_corr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_corr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-17-d5c03ea66df9>\u001b[0m in \u001b[0;36mtrain_on_batch\u001b[0;34m(transformer, optimizer, batch)\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m     \u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_mse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_corr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcalculate_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpadded_labels\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmasks_inputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_weight_softmax\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtarget_weight_softmax\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 41\u001b[0;31m     \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     42\u001b[0m     \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_mse\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_corr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    485\u001b[0m                 \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    486\u001b[0m             )\n\u001b[0;32m--> 487\u001b[0;31m         torch.autograd.backward(\n\u001b[0m\u001b[1;32m    488\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    489\u001b[0m         )\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    198\u001b[0m     \u001b[0;31m# some Python versions print out the first line of a multi-line function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    199\u001b[0m     \u001b[0;31m# calls in the traceback and some print out the last line\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 200\u001b[0;31m     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n\u001b[0m\u001b[1;32m    201\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    202\u001b[0m         allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# function to predict on era2data_eval\n",
        "def predict_on_era2data(era2data, transformer, device=\"cpu\"):\n",
        "\n",
        "    \"\"\"\n",
        "    Currently returns only primary target.\n",
        "    outputs[0][:, 0]: target_nomi_v4_20\n",
        "    \"\"\"\n",
        "    transformer.eval()\n",
        "    with torch.no_grad():\n",
        "        preds = []\n",
        "        for era_num in tqdm(era2data):\n",
        "            batch = era2data[era_num]\n",
        "\n",
        "            padded_inputs = batch[0].to(device=device)\n",
        "            padded_labels = batch[1].to(device=device)\n",
        "            masks_inputs = batch[2].to(device=device)\n",
        "\n",
        "            outputs = transformer(padded_inputs / 4.0, masks_inputs)\n",
        "            preds.append(\n",
        "                outputs[0][masks_inputs.view(-1).nonzero()]\n",
        "                .squeeze(1)\n",
        "                .detach()\n",
        "                .cpu()\n",
        "                .numpy()\n",
        "            )\n",
        "    preds = np.concatenate(preds)\n",
        "    return preds\n",
        "\n",
        "\n",
        "def unif(df):\n",
        "    \"\"\"from example scripts\"\"\"\n",
        "    x = (df.rank(method=\"first\") - 0.5) / len(df)\n",
        "    return pd.Series(x, index=df.index)\n",
        "\n",
        "def calculate_metrics(scores):\n",
        "    return {\n",
        "        \"mean\": scores.mean(),\n",
        "        \"std\": scores.std(),\n",
        "        \"min\": scores.min(),\n",
        "        \"max\": scores.max(),\n",
        "        \"sharpe\": scores.mean() / scores.std(),\n",
        "        \"max_dd\": (scores.cummax() - scores).max(),\n",
        "    }\n",
        "\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "transformer = transformer.to(device)\n",
        "# predict on era2data_eval\n",
        "preds_train = predict_on_era2data(era2data_train, transformer, device=device)\n",
        "preds_valid = predict_on_era2data(era2data_validation, transformer, device=device)\n",
        "\n",
        "train.loc[:, PREDICTION_NAME] = preds_train[:, 0]\n",
        "validation.loc[:, PREDICTION_NAME] = preds_valid[:, 0]\n",
        "\n",
        "\n",
        "scores_train = (\n",
        "    train[[\"era_int\", PREDICTION_NAME, TARGET_NAME]]\n",
        "    .groupby(\"era_int\")\n",
        "    .progress_apply(\n",
        "        lambda d: unif(d[PREDICTION_NAME]).corr(d[TARGET_NAME])\n",
        "    )\n",
        "    .sort_index()\n",
        "    .rename(\"corr\")\n",
        ")\n",
        "scores_valid = (\n",
        "    validation[[\"era_int\", PREDICTION_NAME, TARGET_NAME]]\n",
        "    .groupby(\"era_int\")\n",
        "    .progress_apply(\n",
        "        lambda d: unif(d[PREDICTION_NAME]).corr(d[TARGET_NAME])\n",
        "    )\n",
        "    .sort_index()\n",
        "    .rename(\"corr\")\n",
        ")\n"
      ],
      "metadata": {
        "id": "9JjM4fqm3Vz1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "metrics = calculate_metrics(scores_train)\n",
        "print(json.dumps(metrics, indent=2))\n",
        "scores_train.reset_index().plot(\n",
        "    kind=\"bar\",\n",
        "    figsize=(15, 5),\n",
        "    x=\"era_int\",\n",
        "    y=\"corr\",\n",
        "    xticks=scores_train.index.to_list()[::10],\n",
        ")\n",
        "\n",
        "plt.title(\"Training corr\")\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "KCLWWMVd3VxS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "metrics = calculate_metrics(scores_valid)\n",
        "print(json.dumps(metrics, indent=2))\n",
        "scores_valid.reset_index().plot(\n",
        "    kind=\"bar\",\n",
        "    figsize=(15, 5),\n",
        "    x=\"era_int\",\n",
        "    y=\"corr\",\n",
        ")\n",
        "plt.xticks(\n",
        "    ticks=range(0, len(scores_valid), 10), labels=scores_valid.index.to_list()[::10]\n",
        ")\n",
        "plt.title(\"validation corr\")\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "g7gb1Xq_2qWm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "validation[\"example_preds\"] = validation_example_preds"
      ],
      "metadata": {
        "id": "azMJcUmsE6gH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "diagnosis = validation[[\"era_int\", PREDICTION_NAME, TARGET_NAME, \"example_preds\"]].reset_index()\n",
        "combined = pd.merge(\n",
        "    diagnosis,\n",
        "    meta_model,\n",
        "    on=[\"id\"],\n",
        "    how=\"right\"\n",
        ").dropna(axis=0)\n",
        "\n",
        "pred2meta_model_corr = (\n",
        "    combined[[\"era_int\", \"numerai_meta_model\", PREDICTION_NAME]]\n",
        "    .groupby(\"era_int\")\n",
        "    .progress_apply(\n",
        "        lambda d: unif(d[PREDICTION_NAME]).corr(d[\"numerai_meta_model\"])\n",
        "    )\n",
        "    .sort_index()\n",
        "    .rename(\"pred2meta_model_corr\")\n",
        ")\n",
        "example_preds_corr = (\n",
        "    combined[[\"era_int\", \"example_preds\", PREDICTION_NAME]]\n",
        "    .groupby(\"era_int\")\n",
        "    .progress_apply(\n",
        "        lambda d: unif(d[PREDICTION_NAME]).corr(d[\"example_preds\"])\n",
        "    )\n",
        "    .sort_index()\n",
        "    .rename(\"pred2example_preds_corr\")\n",
        ")\n",
        "\n",
        "example2meta_model_corr = (\n",
        "    combined[[\"era_int\", \"example_preds\", \"numerai_meta_model\"]]\n",
        "    .groupby(\"era_int\")\n",
        "    .progress_apply(\n",
        "        lambda d: unif(d[\"numerai_meta_model\"]).corr(d[\"example_preds\"])\n",
        "    )\n",
        "    .sort_index()\n",
        "    .rename(\"example2meta_model_cor\")\n",
        ")"
      ],
      "metadata": {
        "id": "RZ5ht2vd9qCW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "correlations = pd.concat([\n",
        "    pred2meta_model_corr,\n",
        "    example_preds_corr,\n",
        "    example2meta_model_corr,\n",
        "], axis=1)\n",
        "\n",
        "print(correlations.mean(0))\n",
        "\n",
        "correlations.plot(figsize=(16, 8)).legend(fontsize=15, loc=\"upper left\")\n",
        "\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "HQfFrdy6zVzg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# torch.save(transformer.state_dict(), \"transformer.pth\")"
      ],
      "metadata": {
        "id": "LyPZ9CCDZcYu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if \"transformer.pth\" in os.listdir():\n",
        "    transformer = transformer = Transformer(\n",
        "        input_dim=FEATURE_DIM,\n",
        "        d_model=HIDDEN_DIM,\n",
        "        output_dim=OUTPUT_DIM,\n",
        "        num_heads=NUM_HEADS,\n",
        "        num_layers=NUM_LAYERS,\n",
        "    )\n",
        "    transformer.load_state_dict(torch.load(\"transformer.pth\"))\n",
        "    transformer = transformer.to(device)\n",
        "    transformer.eval()"
      ],
      "metadata": {
        "id": "t72o1IVqiQnx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def prepare_live(data):\n",
        "\n",
        "    inputs = torch.from_numpy(\n",
        "                data[feature_names].values.astype(np.int8))\n",
        "    \n",
        "    padded_inputs, masks_inputs = pad_sequence(\n",
        "            [inputs], padding_value=PADDING_VALUE, max_len=MAX_LEN)\n",
        "    \n",
        "    return padded_inputs, masks_inputs"
      ],
      "metadata": {
        "id": "oQulgceLZz7s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "input_live, mask_live = prepare_live(live)\n",
        "input_live, mask_live = input_live.to(device), mask_live.to(device)\n",
        "preds_live = transformer(input_live/4.0, mask_live)[0][mask_live.view(-1).nonzero()].squeeze(1).detach().cpu().numpy()\n",
        "live[\"prediction\"] = preds_live[:, 0]"
      ],
      "metadata": {
        "id": "yOFXP1PGhe30"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "napi = NumerAPI()\n",
        "current_round = napi.get_current_round()\n",
        "\n",
        "validation[\"prediction\"].to_csv(f\"validation_predictions_{current_round}.csv\")\n",
        "live[\"prediction\"].to_csv(f\"live_predictions_{current_round}.csv\")"
      ],
      "metadata": {
        "id": "8REYqIyxis4E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "dLvfAa7GkRte"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Era Clustering using Embeddings"
      ],
      "metadata": {
        "id": "aEhKBzOenuaz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def embeddings_on_era2data(era2data, transformer, device=\"cpu\"):\n",
        "\n",
        "    \"\"\"\n",
        "    Embeddings from encoder (6000, HIDDEN_DIM) for each era\n",
        "    \"\"\"\n",
        "    transformer.eval()\n",
        "    with torch.no_grad():\n",
        "        preds = []\n",
        "        for era_num in tqdm(era2data):\n",
        "            batch = era2data[era_num]\n",
        "\n",
        "            padded_inputs = batch[0].to(device=device)\n",
        "            padded_labels = batch[1].to(device=device)\n",
        "            masks_inputs = batch[2].to(device=device)\n",
        "\n",
        "            outputs = transformer(padded_inputs / 4.0, masks_inputs)\n",
        "            preds.append(\n",
        "                outputs[0][masks_inputs.view(-1).nonzero()]\n",
        "                .squeeze(1)\n",
        "                .detach()\n",
        "                .cpu()\n",
        "                .numpy()[0, :]\n",
        "            )\n",
        "\n",
        "    preds = np.array(preds)\n",
        "    return preds"
      ],
      "metadata": {
        "id": "biaOgFALis1l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "era_embeds_train = embeddings_on_era2data(era2data_train, transformer=transformer.encoder.to(device), device=device)\n",
        "era_embeds_valid = embeddings_on_era2data(era2data_validation, transformer=transformer.encoder.to(device), device=device)\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "GIPxKj8HkZNO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load data\n",
        "data_train = era_embeds_train.copy()\n",
        "data_valid = era_embeds_valid.copy()\n",
        "\n",
        "data_train = normalize(data_train)\n",
        "data_valid = normalize(data_valid)\n",
        "\n",
        "# Apply PCA to reduce dimensionality\n",
        "pca = PCA(n_components=5, svd_solver=\"full\")\n",
        "pca.fit_transform(data_train)\n",
        "data_pca_train = pca.transform(data_train)\n",
        "data_pca_valid = pca.transform(data_valid)\n",
        "\n",
        "# Apply k-means++ clustering\n",
        "kmeans = KMeans(n_clusters=3, init='k-means++', random_state=0, algorithm=\"lloyd\", n_init=10, max_iter=100)\n",
        "kmeans.fit(data_pca_train)\n",
        "clusters_train = kmeans.predict(data_pca_train)\n",
        "clusters_valid = kmeans.predict(data_pca_valid)\n",
        "\n",
        "fig, axs = plt.subplots(1, 2, figsize=(16, 6)) \n",
        "axs[0].scatter(data_pca_train[:, 0], data_pca_train[:, 1], c=clusters_train, linewidths=1)\n",
        "axs[0].set_title('Era clustered on Train set')\n",
        "axs[0].set_xlabel('Principal Component 1')\n",
        "axs[0].set_ylabel('Principal Component 2')\n",
        "\n",
        "axs[1].scatter(data_pca_valid[:, 0], data_pca_valid[:, 1], c=clusters_valid, linewidths=1)\n",
        "axs[1].set_title('Era clustered on Validation set')\n",
        "axs[1].set_xlabel('Principal Component 1')\n",
        "axs[1].set_ylabel('Principal Component 2')\n",
        "\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "1g-UoFgMl2rt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "9vAdcNPz06ve"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}