fsrs4anki_optimizer.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# FSRS4Anki v2.0.3 Optimizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lurCmW0Jqz3s"
      },
      "source": [
        "[![open in colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/open-spaced-repetition/fsrs4anki/blob/v2.0.3/fsrs4anki_optimizer.ipynb)\n",
        "\n",
        "↑ Click the above button to open the optimizer on Google Colab.\n",
        "\n",
        "> If you can't see the button and are located in the Chinese Mainland, please use a proxy or VPN."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wG7bBfGJFbMr"
      },
      "source": [
        "Upload your **Anki Deck Package (.apkg)** file or **Anki Collection Package (.colpkg)** file on the `Left sidebar -> Files`, drag and drop your file in the current directory (not the `sample_data` directory). \n",
        "\n",
        "No need to include media. Need to include scheduling information. \n",
        "\n",
        "> If you use the latest version of Anki, please check the box `Support older Anki versions (slower/larger files)` when you export.\n",
        "\n",
        "You can export it via `File -> Export...` or `Ctrl + E` in the main window of Anki.\n",
        "\n",
        "Then replace the `filename` with yours in the next code cell. And set the `timezone` and `next_day_starts_at` which can be found in your preferences of Anki.\n",
        "\n",
        "After that, just run all (`Runtime -> Run all` or `Ctrl + F9`) and wait for minutes. You can see the optimal parameters in section **3 Result**. Copy them, replace the parameters in `fsrs4anki_scheduler.js`, and paste them into the custom scheduling of your deck options (require Anki version >= 2.1.55).\n",
        "\n",
        "**NOTE**: The default output is generated from my review logs. If you find the output is the same as mine, maybe your notebook hasn't run there."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "iqP70_-3EUhi"
      },
      "outputs": [],
      "source": [
        "# Here are some settings that you need to replace before running this optimizer.\n",
        "\n",
        "filename = \"ALL__Learning.apkg\"\n",
        "# If you upload deck file, replace it with your deck filename. E.g., ALL__Learning.apkg\n",
        "# If you upload collection file, replace it with your colpgk filename. E.g., collection-2022-09-18@13-21-58.colpkg\n",
        "\n",
        "# Replace it with your timezone. I'm in China, so I use Asia/Shanghai.\n",
        "timezone = 'Asia/Shanghai'\n",
        "\n",
        "# Replace it with your Anki's setting in Prefernces -> Scheduling.\n",
        "next_day_starts_at = 4\n",
        "\n",
        "# Replace it if you don't want the optimizer to use the review logs before a specific date.\n",
        "revlog_start_date = \"2006-10-05\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bLFVNmG2qd06"
      },
      "source": [
        "## 1 Build dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EkzFeKawqgbs"
      },
      "source": [
        "### 1.1 Extract Anki collection & deck file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KD2js_wEr_Bs",
        "outputId": "42653d9e-316e-40bc-bd1d-f3a0e2b246c7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Extract successfully!\n"
          ]
        }
      ],
      "source": [
        "import zipfile\n",
        "# Extract the collection file or deck file to get the .anki21 database.\n",
        "with zipfile.ZipFile(f'./{filename}', 'r') as zip_ref:\n",
        "    zip_ref.extractall('./')\n",
        "    print(\"Extract successfully!\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dKpy4VfqGmaL"
      },
      "source": [
        "### 1.2 Create time-series feature\n",
        "\n",
        "The following code cell will extract the review logs from your Anki collection and preprocess them to a trainset which is saved in `revlog_history.tsv`.\n",
        "\n",
        " The time-series features are important in optimizing the model's parameters. For more detail, please see my paper: https://www.maimemo.com/paper/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "J2IIaY3PDaaG",
        "outputId": "607916c9-da95-48dd-fdab-6bd83fbbbb40"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "revlog.csv saved!\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 5166/5166 [00:18<00:00, 286.72it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Trainset saved!\n"
          ]
        }
      ],
      "source": [
        "import sqlite3\n",
        "import time\n",
        "import tqdm\n",
        "import pandas as pd\n",
        "import os\n",
        "from datetime import timedelta, datetime\n",
        "from tqdm import tqdm\n",
        "\n",
        "if os.path.isfile(\"collection.anki21b\"):\n",
        "    os.remove(\"collection.anki21b\")\n",
        "    raise Exception(\n",
        "        \"Please export the file with `support older Anki versions` if you use the latest version of Anki.\")\n",
        "elif os.path.isfile(\"collection.anki21\"):\n",
        "    con = sqlite3.connect(\"collection.anki21\")\n",
        "elif os.path.isfile(\"collection.anki2\"):\n",
        "    con = sqlite3.connect(\"collection.anki2\")\n",
        "else:\n",
        "    raise Exception(\"Collection not exist!\")\n",
        "cur = con.cursor()\n",
        "res = cur.execute(\"SELECT * FROM revlog\")\n",
        "revlog = res.fetchall()\n",
        "\n",
        "df = pd.DataFrame(revlog)\n",
        "df.columns = ['id', 'cid', 'usn', 'r', 'ivl',\n",
        "              'last_lvl', 'factor', 'time', 'type']\n",
        "df = df[(df['cid'] <= time.time() * 1000) &\n",
        "        (df['id'] <= time.time() * 1000) &\n",
        "        (df['id'] >= time.mktime(datetime.strptime(revlog_start_date, \"%Y-%m-%d\").timetuple()) * 1000)].copy()\n",
        "df['create_date'] = pd.to_datetime(df['cid'] // 1000, unit='s')\n",
        "df['create_date'] = df['create_date'].dt.tz_localize(\n",
        "    'UTC').dt.tz_convert(timezone)\n",
        "df['review_date'] = pd.to_datetime(df['id'] // 1000, unit='s')\n",
        "df['review_date'] = df['review_date'].dt.tz_localize(\n",
        "    'UTC').dt.tz_convert(timezone)\n",
        "df.drop(df[df['review_date'].dt.year < 2006].index, inplace=True)\n",
        "df.sort_values(by=['cid', 'id'], inplace=True, ignore_index=True)\n",
        "df.to_csv(\"revlog.csv\", index=False)\n",
        "print(\"revlog.csv saved!\")\n",
        "df = df[(df['type'] == 0) | (df['type'] == 1)].copy()\n",
        "df['real_days'] = df['review_date'] - timedelta(hours=next_day_starts_at)\n",
        "df['real_days'] = pd.DatetimeIndex(df['real_days'].dt.floor('D')).to_julian_date()\n",
        "df.drop_duplicates(['cid', 'real_days'], keep='first', inplace=True)\n",
        "df['delta_t'] = df.real_days.diff()\n",
        "df.dropna(inplace=True)\n",
        "df['delta_t'] = df['delta_t'].astype(dtype=int)\n",
        "df['i'] = 1\n",
        "df['r_history'] = \"\"\n",
        "df['t_history'] = \"\"\n",
        "col_idx = {key: i for i, key in enumerate(df.columns)}\n",
        "\n",
        "\n",
        "# code from https://github.com/L-M-Sherlock/anki_revlog_analysis/blob/main/revlog_analysis.py\n",
        "def get_feature(x):\n",
        "    for idx, log in enumerate(x.itertuples()):\n",
        "        if idx == 0:\n",
        "            x.iloc[idx, col_idx['delta_t']] = 0\n",
        "        if idx == x.shape[0] - 1:\n",
        "            break\n",
        "        x.iloc[idx + 1, col_idx['i']] = x.iloc[idx, col_idx['i']] + 1\n",
        "        x.iloc[idx + 1, col_idx['t_history']] = f\"{x.iloc[idx, col_idx['t_history']]},{x.iloc[idx, col_idx['delta_t']]}\"\n",
        "        x.iloc[idx + 1, col_idx['r_history']] = f\"{x.iloc[idx, col_idx['r_history']]},{x.iloc[idx, col_idx['r']]}\"\n",
        "    return x\n",
        "\n",
        "\n",
        "tqdm.pandas()\n",
        "df = df.groupby('cid', as_index=False).progress_apply(get_feature)\n",
        "df[\"t_history\"] = df[\"t_history\"].map(lambda x: x[1:] if len(x) > 1 else x)\n",
        "df[\"r_history\"] = df[\"r_history\"].map(lambda x: x[1:] if len(x) > 1 else x)\n",
        "df.to_csv('revlog_history.tsv', sep=\"\\t\", index=False)\n",
        "print(\"Trainset saved!\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k_SgzC-auWmu"
      },
      "source": [
        "## 2 Optimize parameter"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WrfBJjqCHEwJ"
      },
      "source": [
        "### 2.1 Define the model\n",
        "\n",
        "FSRS is a time-series model for predicting memory states."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "tdYp3GMLhTYm"
      },
      "outputs": [],
      "source": [
        "import math\n",
        "import sys\n",
        "import torch\n",
        "import numpy as np\n",
        "from torch import nn\n",
        "from sklearn.utils import shuffle\n",
        "\n",
        "initStability = 1\n",
        "initStabilityRatingFactor = 1\n",
        "initDifficulty = 1\n",
        "initDifficultyRatingFactor = -1\n",
        "updateDifficultyRatingFactor = -1\n",
        "difficultyMeanReversionFactor = 0.2\n",
        "recallFactor = 3\n",
        "recallDifficultyDecay = -0.8\n",
        "recallStabilityDecay = -0.2\n",
        "recallRetrievabilityFactor = 1.3\n",
        "forgetFactor = 2.2\n",
        "forgetDifficultyDecay = -0.3\n",
        "forgetStabilityDecay = 0.3\n",
        "forgetRetrievabilityFactor = 1.2\n",
        "\n",
        "class FSRS(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(FSRS, self).__init__()\n",
        "        self.f_s = nn.Parameter(torch.FloatTensor([initStability, initStabilityRatingFactor]))\n",
        "        # init stability\n",
        "        self.f_d = nn.Parameter(torch.FloatTensor([initDifficulty, initDifficultyRatingFactor, updateDifficultyRatingFactor, difficultyMeanReversionFactor]))\n",
        "        # init difficulty\n",
        "        self.s_w = nn.Parameter(torch.FloatTensor([recallFactor, recallDifficultyDecay, recallStabilityDecay, recallRetrievabilityFactor, forgetFactor, forgetDifficultyDecay, forgetStabilityDecay, forgetRetrievabilityFactor]))\n",
        "        self.zero = torch.FloatTensor([0.0])\n",
        "\n",
        "    def forward(self, x, s, d):\n",
        "        '''\n",
        "        :param x: [review interval, review response]\n",
        "        :param s: stability\n",
        "        :param d: difficulty\n",
        "        :return:\n",
        "        '''\n",
        "        if torch.equal(s, torch.FloatTensor([0.0])):\n",
        "            # first learn, init memory states\n",
        "            next_s = self.f_s[0] * (self.f_s[1] * (x[1] - 1) + 1)\n",
        "            next_d = self.f_d[0] * (self.f_d[1] * (x[1] - 4) + 1)\n",
        "        else:\n",
        "            r = torch.exp(np.log(0.9) * x[0] / s)\n",
        "            if x[1] > 1:\n",
        "                next_s = s * (1 + torch.exp(self.s_w[0]) * torch.pow(d, self.s_w[1]) *\n",
        "                              torch.pow(s, self.s_w[2]) *\n",
        "                              (torch.exp((1 - r) * self.s_w[3]) - 1))\n",
        "            else:\n",
        "                next_s = self.s_w[4] * torch.pow(d, self.s_w[5]) * torch.pow(s, self.s_w[6]) * torch.exp((1 - r) * self.s_w[7])\n",
        "            next_d = d + self.f_d[2] * (x[1] - 3)\n",
        "            next_d = self.mean_reversion(self.f_d[0] * (- self.f_d[1] + 1), next_d)\n",
        "        return next_s, self.constrain(next_d)\n",
        "\n",
        "    def loss(self, s, t, r):\n",
        "        return - (r * np.log(0.9) * t / s + (1 - r) * torch.log(1 - torch.exp(np.log(0.9) * t / s)))\n",
        "\n",
        "    def constrain(self, d):\n",
        "        return torch.relu(d - 1) + 1\n",
        "\n",
        "    def mean_reversion(self, init, current):\n",
        "        return self.f_d[3] * init + (1-self.f_d[3]) * current\n",
        "\n",
        "\n",
        "class WeightClipper(object):\n",
        "    def __init__(self, frequency=1):\n",
        "        self.frequency = frequency\n",
        "\n",
        "    def __call__(self, module):\n",
        "        if hasattr(module, 'f_s'):\n",
        "            w = module.f_s.data\n",
        "            w[0] = w[0].clamp(0.1, 10)  # initStability\n",
        "            w[1] = w[1].clamp(0.01, 10)  # initStabilityRatingFactor\n",
        "            module.f_s.data = w\n",
        "        if hasattr(module, 'f_d'):\n",
        "            w = module.f_d.data\n",
        "            w[0] = w[0].clamp(1, 10)  # initDifficulty\n",
        "            w[1] = w[1].clamp(-10, -0.01)  # initDifficultyRatingFactor\n",
        "            w[2] = w[2].clamp(-10, -0.01)  # updateDifficultyRatingFactor\n",
        "            w[3] = w[3].clamp(0, 1)  # difficultyMeanReversionFactor\n",
        "            module.f_d.data = w\n",
        "        if hasattr(module, 's_w'):\n",
        "            w = module.s_w.data\n",
        "            w[0] = w[0].clamp(0, 5)  # recallFactor\n",
        "            w[1] = w[1].clamp(-2, -0.01)  # recallDifficultyDecay\n",
        "            w[2] = w[2].clamp(-2, -0.01)  # recallStabilityDecay\n",
        "            w[3] = w[3].clamp(0.01, 2)  # recallRetrievabilityFactor\n",
        "            w[4] = w[4].clamp(0, 5)  # forgetFactor\n",
        "            w[5] = w[5].clamp(-2, -0.01)  # forgetDifficultyDecay\n",
        "            w[6] = w[6].clamp(0.01, 1)  # forgetStabilityDecay\n",
        "            w[7] = w[7].clamp(0.01, 2)  # forgetRetrievabilityFactor\n",
        "            module.s_w.data = w\n",
        "\n",
        "\n",
        "def lineToTensor(line):\n",
        "    ivl = line[0].split(',')\n",
        "    response = line[1].split(',')\n",
        "    tensor = torch.zeros(len(response), 2)\n",
        "    for li, response in enumerate(response):\n",
        "        tensor[li][0] = int(ivl[li])\n",
        "        tensor[li][1] = int(response)\n",
        "    return tensor\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8E1dYfgQLZAC"
      },
      "source": [
        "### 2.2 Train the model\n",
        "\n",
        "The `revlog_history.tsv` generated before will be used for training the FSRS model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Jht0gneShowU",
        "outputId": "aaa72b79-b454-483b-d746-df1a353b2c8f"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 41465/41465 [00:02<00:00, 20600.72it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Tensorized!\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:   0%|\u001b[31m          \u001b[0m| 49/41465 [00:00<01:33, 443.56it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 1\n",
            "f_s: [1.0005, 1.0]\n",
            "f_d: [1.0, -0.9995, -1.0, 0.2005]\n",
            "s_w: [3.0005, -0.7995, -0.1995, 1.3005, 2.2, -0.3, 0.3, 1.2]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  10%|\u001b[31m█         \u001b[0m| 4194/41465 [00:07<01:05, 569.24it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 4147\n",
            "f_s: [1.1787, 1.2177]\n",
            "f_d: [1.0032, -0.9741, -1.0962, 0.0792]\n",
            "s_w: [3.0362, -0.8431, -0.1384, 1.3329, 2.1557, -0.3478, 0.263, 1.1672]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  20%|\u001b[31m██        \u001b[0m| 8366/41465 [00:15<00:56, 584.60it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 8293\n",
            "f_s: [1.2675, 1.3764]\n",
            "f_d: [1.0092, -0.981, -1.116, 0.0485]\n",
            "s_w: [3.0417, -0.876, -0.1356, 1.3357, 2.1557, -0.3444, 0.2861, 1.1626]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  30%|\u001b[31m███       \u001b[0m| 12529/41465 [00:23<00:49, 580.15it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 12439\n",
            "f_s: [1.4104, 1.6081]\n",
            "f_d: [1.004, -0.9809, -1.144, 0.026]\n",
            "s_w: [3.0468, -0.9103, -0.1188, 1.3382, 2.1532, -0.3275, 0.2825, 1.1875]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  40%|\u001b[31m████      \u001b[0m| 16708/41465 [00:31<00:42, 580.14it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 16585\n",
            "f_s: [1.4912, 1.7228]\n",
            "f_d: [1.0246, -1.0256, -1.0955, 0.0635]\n",
            "s_w: [3.0565, -0.9029, -0.1192, 1.3439, 2.1617, -0.2983, 0.2644, 1.1984]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  50%|\u001b[31m█████     \u001b[0m| 20832/41465 [00:40<00:39, 519.82it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 20731\n",
            "f_s: [1.5297, 1.7955]\n",
            "f_d: [1.0009, -1.0033, -1.1073, 0.0404]\n",
            "s_w: [3.0891, -0.8993, -0.0775, 1.3721, 2.1375, -0.3252, 0.2676, 1.169]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  60%|\u001b[31m██████    \u001b[0m| 24988/41465 [00:47<00:29, 559.18it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 24877\n",
            "f_s: [1.5639, 1.8661]\n",
            "f_d: [1.0361, -1.0327, -1.1028, 0.0541]\n",
            "s_w: [3.0686, -0.9112, -0.106, 1.3512, 2.1429, -0.3039, 0.3036, 1.1938]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  70%|\u001b[31m███████   \u001b[0m| 29082/41465 [00:55<00:22, 542.99it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 29023\n",
            "f_s: [1.6234, 1.9308]\n",
            "f_d: [1.0678, -1.0522, -1.1199, 0.0519]\n",
            "s_w: [3.0447, -0.9389, -0.1301, 1.326, 2.1053, -0.3022, 0.28, 1.172]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  80%|\u001b[31m████████  \u001b[0m| 33219/41465 [01:03<00:15, 524.23it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 33169\n",
            "f_s: [1.6724, 1.9458]\n",
            "f_d: [1.048, -1.0358, -1.1234, 0.0705]\n",
            "s_w: [3.0564, -0.9348, -0.113, 1.3356, 2.1086, -0.2945, 0.2685, 1.1652]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train:  90%|\u001b[31m█████████ \u001b[0m| 37400/41465 [01:11<00:07, 528.12it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 37315\n",
            "f_s: [1.6507, 1.965]\n",
            "f_d: [1.0515, -1.0513, -1.1295, 0.0683]\n",
            "s_w: [3.0489, -0.9552, -0.1087, 1.3249, 2.1207, -0.2795, 0.2845, 1.1609]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: 100%|\u001b[31m██████████\u001b[0m| 41465/41465 [01:18<00:00, 526.77it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "iteration: 41461\n",
            "f_s: [1.6745, 2.0006]\n",
            "f_d: [1.0271, -1.0344, -1.1223, 0.0672]\n",
            "s_w: [3.0707, -0.9418, -0.0995, 1.3432, 2.1267, -0.2709, 0.3076, 1.1563]\n",
            "\n",
            "Training finished!\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "model = FSRS()\n",
        "clipper = WeightClipper()\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=5e-4)\n",
        "\n",
        "dataset = pd.read_csv(\"./revlog_history.tsv\", sep='\\t', index_col=None)\n",
        "dataset = dataset[(dataset['i'] > 1) & (dataset['delta_t'] > 0) & (dataset['t_history'].str.count('0') == 1)]\n",
        "dataset['tensor'] = dataset.progress_apply(lambda x: lineToTensor(\n",
        "    list(zip([x['t_history']], [x['r_history']]))[0]), axis=1)\n",
        "print(\"Tensorized!\")\n",
        "\n",
        "n_epoch = 1\n",
        "print_len = max(dataset.shape[0] // 10, 1)\n",
        "\n",
        "checkpoint = {\n",
        "    \"net\": model.state_dict(),\n",
        "    'optimizer': optimizer.state_dict(),\n",
        "    \"epoch\": -1\n",
        "}\n",
        "\n",
        "for k in range(n_epoch):\n",
        "    dataset = shuffle(dataset, random_state=2022 + k)\n",
        "    epoch_len = len(dataset)\n",
        "    for i, (_, row) in enumerate(tqdm(dataset.iterrows(), total=epoch_len, desc=\"train\", colour=\"red\")):\n",
        "        model.train()\n",
        "        optimizer.zero_grad()\n",
        "        output_t = [(model.zero, model.zero)]\n",
        "        for input_t in row['tensor']:\n",
        "            output_t.append(model(input_t, *output_t[-1]))\n",
        "        loss = model.loss(output_t[-1][0], row['delta_t'],\n",
        "                          {1: 0, 2: 1, 3: 1, 4: 1}[row['r']])\n",
        "        if np.isnan(loss.data.item()):\n",
        "            # Exception Case\n",
        "            print(row, output_t)\n",
        "            raise Exception('error case')\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        model.apply(clipper)\n",
        "\n",
        "        if (k * epoch_len + i) % print_len == 0:\n",
        "            print(f\"iteration: {k * epoch_len + i + 1}\")\n",
        "            for name, param in model.named_parameters():\n",
        "                print(f\"{name}: {list(map(lambda x: round(float(x), 4),param))}\")\n",
        "\n",
        "            checkpoint = {\n",
        "                \"net\": model.state_dict(),\n",
        "                \"optimizer\": optimizer.state_dict(),\n",
        "                \"iteration\": (k * epoch_len + i) // print_len\n",
        "            }\n",
        "\n",
        "torch.save(checkpoint, f'./model.pth')\n",
        "\n",
        "initStability, initStabilityRatingFactor = map(\n",
        "    lambda x: round(float(x), 4), dict(model.named_parameters())['f_s'].data)\n",
        "initDifficulty, initDifficultyRatingFactor, updateDifficultyRatingFactor, difficultyMeanReversionFactor = map(\n",
        "    lambda x: round(float(x), 4), dict(model.named_parameters())['f_d'].data)\n",
        "recallFactor, recallDifficultyDecay, recallStabilityDecay, recallRetrievabilityFactor, forgetFactor, forgetDifficultyDecay, forgetStabilityDecay, forgetRetrievabilityFactor = map(\n",
        "    lambda x: round(float(x), 4), dict(model.named_parameters())['s_w'].data)\n",
        "\n",
        "print(\"\\nTraining finished!\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZ4S2l7BWfzr"
      },
      "source": [
        "## 3 Result\n",
        "\n",
        "Copy the optimal parameters for FSRS for you in the output of next code cell after running."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NTnPSDA2QpUu",
        "outputId": "49f487b9-69a7-4e96-b35a-7e027f478fbd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "var f_s = [1.6746,2.0004];\n",
            "var f_d = [1.028,-1.0351,-1.1225,0.0667];\n",
            "var s_w = [3.0699,-0.9423,-0.1005,1.3424,2.1268,-0.2708,0.3077,1.1564];\n"
          ]
        }
      ],
      "source": [
        "print(f\"var f_s = [{initStability},{initStabilityRatingFactor}];\")\n",
        "print(f\"var f_d = [{initDifficulty},{initDifficultyRatingFactor},{updateDifficultyRatingFactor},{difficultyMeanReversionFactor}];\")\n",
        "print(f\"var s_w = [{recallFactor},{recallDifficultyDecay},{recallStabilityDecay},{recallRetrievabilityFactor},{forgetFactor},{forgetDifficultyDecay},{forgetStabilityDecay},{forgetRetrievabilityFactor}];\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I_zsoDyTaTrT"
      },
      "source": [
        "You can see the memory states and intervals generated by FSRS as if you press the good in each review at the due date scheduled by FSRS."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iws4rtP1WKBT",
        "outputId": "890d0287-1a17-4c59-fbbf-ee54d79cd383"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "1:again, 2:hard, 3:good, 4:easy\n",
            "\n",
            "first rating: 1\n",
            "rating history: 1,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3\n",
            "interval history: 0,2,3,5,9,15,26,43,71,116,188,302,480,756,1180,1825\n",
            "difficulty history: 0,4.2,4.1,3.9,3.8,3.7,3.6,3.5,3.4,3.3,3.2,3.2,3.1,3.0,3.0,2.9\n",
            "\n",
            "first rating: 2\n",
            "rating history: 2,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3\n",
            "interval history: 0,5,9,17,31,56,98,169,288,483,797,1298,2085,3306,5177,8011\n",
            "difficulty history: 0,3.2,3.1,3.0,3.0,2.9,2.8,2.8,2.7,2.7,2.7,2.6,2.6,2.6,2.5,2.5\n",
            "\n",
            "first rating: 3\n",
            "rating history: 3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3\n",
            "interval history: 0,8,18,39,81,161,310,579,1051,1857,3202,5399,8912,14428,22935,35842\n",
            "difficulty history: 0,2.1,2.1,2.1,2.1,2.1,2.1,2.1,2.1,2.1,2.1,2.1,2.1,2.1,2.1,2.1\n",
            "\n",
            "first rating: 4\n",
            "rating history: 4,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3\n",
            "interval history: 0,12,40,118,314,764,1722,3636,7258,13795,25118,44039,74668,122881,196898,308017\n",
            "difficulty history: 0,1.0,1.1,1.2,1.2,1.3,1.3,1.4,1.4,1.5,1.5,1.6,1.6,1.6,1.7,1.7\n",
            "\n"
          ]
        }
      ],
      "source": [
        "requestRetention = 0.9  # recommended setting: 0.8 ~ 0.9\n",
        "\n",
        "\n",
        "class Collection:\n",
        "    def __init__(self):\n",
        "        self.model = model\n",
        "\n",
        "    def states(self, t_history, r_history):\n",
        "        with torch.no_grad():\n",
        "            line_tensor = lineToTensor(list(zip([t_history], [r_history]))[0])\n",
        "            output_t = [(self.model.zero, self.model.zero)]\n",
        "            for input_t in line_tensor:\n",
        "                output_t.append(self.model(input_t, *output_t[-1]))\n",
        "            return output_t[-1]\n",
        "\n",
        "\n",
        "my_collection = Collection()\n",
        "print(\"1:again, 2:hard, 3:good, 4:easy\\n\")\n",
        "for first_rating in (1,2,3,4):\n",
        "    print(f'first rating: {first_rating}')\n",
        "    t_history = \"0\"\n",
        "    d_history = \"0\"\n",
        "    r_history = f\"{first_rating}\"  # the first rating of the new card\n",
        "    # print(\"stability, difficulty, lapses\")\n",
        "    for i in range(15):\n",
        "        states = my_collection.states(t_history, r_history)\n",
        "        # print('{0:9.2f} {1:11.2f} {2:7.0f}'.format(\n",
        "            # *list(map(lambda x: round(float(x), 4), states))))\n",
        "        next_t = round(float(np.log(requestRetention)/np.log(0.9) * states[0]))\n",
        "        difficulty = round(float(states[1]), 1)\n",
        "        t_history += f',{int(next_t)}'\n",
        "        d_history += f',{difficulty}'\n",
        "        r_history += f\",3\"\n",
        "    print(f\"rating history: {r_history}\")\n",
        "    print(f\"interval history: {t_history}\")\n",
        "    print(f\"difficulty history: {d_history}\")\n",
        "    print('')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "You can change the `test_rating_sequence` to see the scheduling intervals in different ratings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(tensor(8.3746), tensor(2.0922))\n",
            "(tensor(3.7407), tensor(4.1873))\n",
            "(tensor(2.4496), tensor(6.1425))\n",
            "(tensor(3.4695), tensor(5.8722))\n",
            "(tensor(5.0118), tensor(5.6199))\n",
            "(tensor(7.5984), tensor(5.3844))\n",
            "(tensor(11.7338), tensor(5.1647))\n",
            "(tensor(3.2731), tensor(7.0547))\n",
            "(tensor(4.5791), tensor(6.7235))\n",
            "(tensor(6.7862), tensor(6.4144))\n",
            "(tensor(9.8887), tensor(6.1259))\n",
            "(tensor(14.3442), tensor(5.8567))\n",
            "(tensor(20.6100), tensor(5.6054))\n",
            "(tensor(30.0604), tensor(5.3709))\n",
            "(tensor(43.5889), tensor(5.1521))\n",
            "(tensor(63.4707), tensor(4.9478))\n",
            "(tensor(5.5478), tensor(6.8523))\n",
            "rating history: 3,1,1,3,3,3,3,1,3,3,3,3,3,3,3,3,1\n",
            "interval history: 0,8,4,2,3,5,8,12,3,5,7,10,14,21,30,44,63,6\n",
            "difficulty history: 0,2.1,4.2,6.1,5.9,5.6,5.4,5.2,7.1,6.7,6.4,6.1,5.9,5.6,5.4,5.2,4.9,6.9\n"
          ]
        }
      ],
      "source": [
        "test_rating_sequence = \"3,1,1,3,3,3,3,1,3,3,3,3,3,3,3,3,1\"\n",
        "requestRetention = 0.9  # recommended setting: 0.8 ~ 0.9\n",
        "easyBonus = 1.3\n",
        "hardInterval = 1.2\n",
        "\n",
        "t_history = \"0\"\n",
        "d_history = \"0\"\n",
        "for i in range(len(test_rating_sequence.split(','))):\n",
        "    rating = test_rating_sequence[2*i]\n",
        "    last_t = int(t_history.split(',')[-1])\n",
        "    r_history = test_rating_sequence[:2*i+1]\n",
        "    states = my_collection.states(t_history, r_history)\n",
        "    print(states)\n",
        "    next_t = max(1,round(float(np.log(requestRetention)/np.log(0.9) * states[0])))\n",
        "    if rating == '4':\n",
        "        next_t = round(next_t * easyBonus)\n",
        "    elif rating == '2':\n",
        "        next_t = round(last_t * hardInterval)\n",
        "    t_history += f',{int(next_t)}'\n",
        "    difficulty = round(float(np.log(requestRetention)/np.log(0.9) * states[1]), 1)\n",
        "    d_history += f',{difficulty}'\n",
        "print(f\"rating history: {test_rating_sequence}\")\n",
        "print(f\"interval history: {t_history}\")\n",
        "print(f\"difficulty history: {d_history}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyMnk8/Ih2JAJZJ1PBkXQUBC",
      "collapsed_sections": [],
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3.8.13 ('fsrs4anki')",
      "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.8.13"
    },
    "vscode": {
      "interpreter": {
        "hash": "8dd9a290ffd10997e0b0d411ff1325a47862ea932e0fd309ade800e0e51d2b4b"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}