Init commit

.gitignore

@@ -0,0 +1 @@
+.idea/

app.py

@@ -0,0 +1,112 @@
+import gradio as gr
+import pytz
+
+from datetime import datetime
+
+from utilities import extract, create_time_series_features, train_model, process_personalized_collection, my_loss, \
+    cleanup
+from memory_states import get_my_memory_states
+from plot import make_plot
+
+
+def anki_optimizer(file, timezone, next_day_starts_at, revlog_start_date, requestRetention,
+                   progress=gr.Progress(track_tqdm=True)):
+    now = datetime.now()
+    prefix = now.strftime(f'%Y_%m_%d_%H_%M_%S')
+    proj_dir = extract(file, prefix)
+    type_sequence, df_out = create_time_series_features(revlog_start_date, timezone, next_day_starts_at, proj_dir)
+    w, dataset = train_model(proj_dir)
+    my_collection, rating_markdown = process_personalized_collection(requestRetention, w)
+    difficulty_distribution_padding, difficulty_distribution = get_my_memory_states(proj_dir, dataset, my_collection)
+    fig, suggested_retention_markdown = make_plot(proj_dir, type_sequence, w, difficulty_distribution_padding)
+    loss_markdown = my_loss(dataset, w)
+    difficulty_distribution = difficulty_distribution.to_string().replace("\n", "\n\n")
+    markdown_out = f"""
+{suggested_retention_markdown}
+
+# Loss Information
+{loss_markdown}
+
+# Difficulty Distribution
+{difficulty_distribution}
+
+# Ratings
+{rating_markdown}
+"""
+
+    w_markdown = f"""
+    # These are the weights for step 5
+    `var w = {w};`"""
+    files = ['prediction.tsv', 'revlog.csv', 'revlog_history.tsv', 'stability_for_analysis.tsv',
+             'expected_repetitions.csv']
+    files_out = [proj_dir / file for file in files]
+    cleanup(proj_dir, files)
+    return w_markdown, df_out, fig, markdown_out, files_out
+
+
+with gr.Blocks() as demo:
+    with gr.Tab("FSRS4Anki Optimizer"):
+        with gr.Box():
+            gr.Markdown("""
+            Based on the [tutorial](https://medium.com/@JarrettYe/how-to-use-the-next-generation-spaced-repetition-algorithm-fsrs-on-anki-5a591ca562e2) of [Jarrett Ye](https://github.com/L-M-Sherlock)
+            Check out the instructions on the next tab.
+            """)
+        with gr.Box():
+            with gr.Row():
+                file = gr.File(label='Review Logs')
+                timezone = gr.Dropdown(label="Choose your timezone", choices=pytz.all_timezones)
+            with gr.Row():
+                next_day_starts_at = gr.Number(value=4,
+                                               label="Replace it with your Anki's setting in Preferences -> Scheduling.",
+                                               precision=0)
+                with gr.Accordion(label="Advanced Settings", open=False):
+                    requestRetention = gr.Number(value=.9, label="Recommended to set between 0.8  0.9")
+            with gr.Row():
+                revlog_start_date = gr.Textbox(value="2006-10-05",
+                                               label="Replace it if you don't want the optimizer to use the review logs before a specific date.")
+        with gr.Row():
+            btn_plot = gr.Button('Optimize your Anki!')
+        with gr.Row():
+            w_output = gr.Markdown()
+    with gr.Tab("Instructions"):
+        with gr.Box():
+            gr.Markdown("""
+            # How to get personalized Anki parameters
+            If you have been using Anki for some time and have accumulated a lot of review logs, you can try this FSRS4Anki 
+            optimizer app to generate parameters for you.
+
+            This is based on the amazing work of [Jarrett Ye](https://github.com/L-M-Sherlock)
+            # Step 1 - Get the review logs to upload
+            1. Click the gear icon to the right of a deck’s name 
+            2. Export 
+            3. Check “Include scheduling information” and “Support older Anki versions”
+            ![](https://miro.medium.com/v2/resize:fit:1400/format:webp/1*W3Nnfarki2z7Ukyom4kMuw.png)
+            4. Export and upload that file to the app
+
+            # Step 2 - Get the `next_day_starts_at` parameter
+            1. Open preferences
+            2. Copy the next day starts at value and paste it in the app
+            ![](https://miro.medium.com/v2/resize:fit:1072/format:webp/1*qAUb6ry8UxFeCsjnKLXvsQ.png)
+
+            # Step 3 - Fill in the rest of the settings
+
+            # Step 4 Click run
+            
+            # Step 5 - Replace the default parameters in FSRS4Anki with the optimized parameters
+            ![](https://miro.medium.com/v2/resize:fit:1252/format:webp/1*NM4CR-n7nDk3nQN1Bi30EA.png)
+            """)
+    with gr.Tab("Analysis"):
+        with gr.Row():
+            markdown_output = gr.Markdown()
+            df_output = gr.DataFrame()
+        with gr.Row():
+            plot_output = gr.Plot()
+        with gr.Row():
+            files_output = gr.Files(label="Analysis Files")
+
+    btn_plot.click(anki_optimizer,
+                   inputs=[file, timezone, next_day_starts_at, revlog_start_date, requestRetention],
+                   outputs=[w_output, df_output, plot_output, markdown_output, files_output])
+demo.queue().launch(debug=True, show_error=True)
+
+# demo.queue().launch(debug=True)

memory_states.py

@@ -0,0 +1,35 @@
+import numpy as np
+from functools import partial
+
+import pandas as pd
+
+
+def predict_memory_states(my_collection, group):
+    states = my_collection.states(*group.name)
+    group['stability'] = float(states[0])
+    group['difficulty'] = float(states[1])
+    group['count'] = len(group)
+    return pd.DataFrame({
+        'r_history': [group.name[1]],
+        't_history': [group.name[0]],
+        'stability': [round(float(states[0]), 2)],
+        'difficulty': [round(float(states[1]), 2)],
+        'count': [len(group)]
+        })
+
+
+def get_my_memory_states(proj_dir, dataset, my_collection):
+    prediction = dataset.groupby(by=['t_history', 'r_history']).progress_apply(
+        partial(predict_memory_states, my_collection))
+    prediction.reset_index(drop=True, inplace=True)
+    prediction.sort_values(by=['r_history'], inplace=True)
+    prediction.to_csv(proj_dir / "prediction.tsv", sep='\t', index=None)
+    print("prediction.tsv saved.")
+    prediction['difficulty'] = prediction['difficulty'].map(lambda x: int(round(x)))
+    difficulty_distribution = prediction.groupby(by=['difficulty'])['count'].sum() / prediction['count'].sum()
+    print(difficulty_distribution)
+    difficulty_distribution_padding = np.zeros(10)
+    for i in range(10):
+        if i + 1 in difficulty_distribution.index:
+            difficulty_distribution_padding[i] = difficulty_distribution.loc[i + 1]
+    return difficulty_distribution_padding, difficulty_distribution

model.py

@@ -0,0 +1,93 @@
+import numpy as np
+import torch
+from torch import nn
+
+init_w = [1, 1, 5, -0.5, -0.5, 0.2, 1.4, -0.02, 0.8, 2, -0.2, 0.5, 1]
+
+
+class FSRS(nn.Module):
+    def __init__(self, w):
+        super(FSRS, self).__init__()
+        self.w = nn.Parameter(torch.FloatTensor(w))
+        self.zero = torch.FloatTensor([0.0])
+
+    def forward(self, x, s, d):
+        '''
+        :param x: [review interval, review response]
+        :param s: stability
+        :param d: difficulty
+        :return:
+        '''
+        if torch.equal(s, self.zero):
+            # first learn, init memory states
+            new_s = self.w[0] + self.w[1] * (x[1] - 1)
+            new_d = self.w[2] + self.w[3] * (x[1] - 3)
+            new_d = new_d.clamp(1, 10)
+        else:
+            r = torch.exp(np.log(0.9) * x[0] / s)
+            new_d = d + self.w[4] * (x[1] - 3)
+            new_d = self.mean_reversion(self.w[2], new_d)
+            new_d = new_d.clamp(1, 10)
+            # recall
+            if x[1] > 1:
+                new_s = s * (1 + torch.exp(self.w[6]) *
+                             (11 - new_d) *
+                             torch.pow(s, self.w[7]) *
+                             (torch.exp((1 - r) * self.w[8]) - 1))
+            # forget
+            else:
+                new_s = self.w[9] * torch.pow(new_d, self.w[10]) * torch.pow(
+                    s, self.w[11]) * torch.exp((1 - r) * self.w[12])
+        return new_s, new_d
+
+    def loss(self, s, t, r):
+        return - (r * np.log(0.9) * t / s + (1 - r) * torch.log(1 - torch.exp(np.log(0.9) * t / s)))
+
+    def mean_reversion(self, init, current):
+        return self.w[5] * init + (1-self.w[5]) * current
+
+
+class WeightClipper(object):
+    def __init__(self, frequency=1):
+        self.frequency = frequency
+
+    def __call__(self, module):
+        if hasattr(module, 'w'):
+            w = module.w.data
+            w[0] = w[0].clamp(0.1, 10)  # initStability
+            w[1] = w[1].clamp(0.1, 5)  # initStabilityRatingFactor
+            w[2] = w[2].clamp(1, 10)  # initDifficulty
+            w[3] = w[3].clamp(-5, -0.1)  # initDifficultyRatingFactor
+            w[4] = w[4].clamp(-5, -0.1)  # updateDifficultyRatingFactor
+            w[5] = w[5].clamp(0, 0.5)  # difficultyMeanReversionFactor
+            w[6] = w[6].clamp(0, 2)  # recallFactor
+            w[7] = w[7].clamp(-0.2, -0.01)  # recallStabilityDecay
+            w[8] = w[8].clamp(0.01, 1.5)  # recallRetrievabilityFactor
+            w[9] = w[9].clamp(0.5, 5)  # forgetFactor
+            w[10] = w[10].clamp(-2, -0.01)  # forgetDifficultyDecay
+            w[11] = w[11].clamp(0.01, 0.9)  # forgetStabilityDecay
+            w[12] = w[12].clamp(0.01, 2)  # forgetRetrievabilityFactor
+            module.w.data = w
+
+
+def lineToTensor(line):
+    ivl = line[0].split(',')
+    response = line[1].split(',')
+    tensor = torch.zeros(len(response), 2)
+    for li, response in enumerate(response):
+        tensor[li][0] = int(ivl[li])
+        tensor[li][1] = int(response)
+    return tensor
+
+
+class Collection:
+    def __init__(self, w):
+        self.model = FSRS(w)
+
+    def states(self, t_history, r_history):
+        with torch.no_grad():
+            line_tensor = lineToTensor(list(zip([t_history], [r_history]))[0])
+            output_t = [(self.model.zero, self.model.zero)]
+            for input_t in line_tensor:
+                output_t.append(self.model(input_t, *output_t[-1]))
+            return output_t[-1]

plot.py

@@ -0,0 +1,92 @@
+from tqdm.auto import trange
+import gradio as gr
+import pandas as pd
+import numpy as np
+import plotly.express as px
+
+
+def make_plot(proj_dir, type_sequence, w, difficulty_distribution_padding, progress=gr.Progress(track_tqdm=True)):
+    base = 1.01
+    index_len = 800
+    index_offset = 150
+    d_range = 10
+    d_offset = 1
+    r_repetitions = 1
+    f_repetitions = 2.3
+    max_repetitions = 200000
+
+    type_block = dict()
+    type_count = dict()
+    last_t = type_sequence[0]
+    type_block[last_t] = 1
+    type_count[last_t] = 1
+    for t in type_sequence[1:]:
+        type_count[t] = type_count.setdefault(t, 0) + 1
+        if t != last_t:
+            type_block[t] = type_block.setdefault(t, 0) + 1
+        last_t = t
+    if 2 in type_count and 2 in type_block:
+        f_repetitions = round(type_count[2] / type_block[2] + 1, 1)
+
+    def stability2index(stability):
+        return int(round(np.log(stability) / np.log(base)) + index_offset)
+
+    def init_stability(d):
+        return max(((d - w[2]) / w[3] + 2) * w[1] + w[0], np.power(base, -index_offset))
+
+    def cal_next_recall_stability(s, r, d, response):
+        if response == 1:
+            return s * (1 + np.exp(w[6]) * (11 - d) * np.power(s, w[7]) * (np.exp((1 - r) * w[8]) - 1))
+        else:
+            return w[9] * np.power(d, w[10]) * np.power(s, w[11]) * np.exp((1 - r) * w[12])
+
+    stability_list = np.array([np.power(base, i - index_offset) for i in range(index_len)])
+    print(f"terminal stability: {stability_list.max(): .2f}")
+    df = pd.DataFrame(columns=["retention", "difficulty", "repetitions"])
+
+    for percentage in trange(96, 70, -2, desc='Repetition vs Retention plot'):
+        recall = percentage / 100
+        repetitions_list = np.zeros((d_range, index_len))
+        repetitions_list[:, :-1] = max_repetitions
+        for d in range(d_range, 0, -1):
+            s0 = init_stability(d)
+            s0_index = stability2index(s0)
+            diff = max_repetitions
+            while diff > 0.1:
+                s0_repetitions = repetitions_list[d - 1][s0_index]
+                for s_index in range(index_len - 2, -1, -1):
+                    stability = stability_list[s_index];
+                    interval = max(1, round(stability * np.log(recall) / np.log(0.9)))
+                    p_recall = np.power(0.9, interval / stability)
+                    recall_s = cal_next_recall_stability(stability, p_recall, d, 1)
+                    forget_d = min(d + d_offset, 10)
+                    forget_s = cal_next_recall_stability(stability, p_recall, forget_d, 0)
+                    recall_s_index = min(stability2index(recall_s), index_len - 1)
+                    forget_s_index = min(max(stability2index(forget_s), 0), index_len - 1)
+                    recall_repetitions = repetitions_list[d - 1][recall_s_index] + r_repetitions
+                    forget_repetitions = repetitions_list[forget_d - 1][forget_s_index] + f_repetitions
+                    exp_repetitions = p_recall * recall_repetitions + (1.0 - p_recall) * forget_repetitions
+                    if exp_repetitions < repetitions_list[d - 1][s_index]:
+                        repetitions_list[d - 1][s_index] = exp_repetitions
+                diff = s0_repetitions - repetitions_list[d - 1][s0_index]
+            df.loc[0 if pd.isnull(df.index.max()) else df.index.max() + 1] = [recall, d, s0_repetitions]
+
+    df.sort_values(by=["difficulty", "retention"], inplace=True)
+    df.to_csv(proj_dir/"expected_repetitions.csv", index=False)
+    print("expected_repetitions.csv saved.")
+
+    optimal_retention_list = np.zeros(10)
+    df2 = pd.DataFrame()
+    for d in range(1, d_range + 1):
+        retention = df[df["difficulty"] == d]["retention"]
+        repetitions = df[df["difficulty"] == d]["repetitions"]
+        optimal_retention = retention.iat[repetitions.argmin()]
+        optimal_retention_list[d - 1] = optimal_retention
+        df2 = df2.append(
+            pd.DataFrame({'retention': retention, 'expected repetitions': repetitions, 'd': d, 'r': optimal_retention}))
+
+    fig = px.line(df2, x="retention", y="expected repetitions", color='d', log_y=True)
+
+    print(f"\n-----suggested retention: {np.inner(difficulty_distribution_padding, optimal_retention_list):.2f}-----")
+    suggested_retention_markdown = f"""# Suggested Retention: `{np.inner(difficulty_distribution_padding, optimal_retention_list):.2f}`"""
+    return fig, suggested_retention_markdown

requirements.txt

@@ -0,0 +1,7 @@
+matplotlib==3.4.3
+numpy==1.23.3
+pandas==1.3.2
+scikit_learn==1.1.2
+torch==1.9.0
+tqdm==4.64.1
+plotly==5.13.0

utilities.py

@@ -0,0 +1,296 @@
+from functools import partial
+import datetime
+from zipfile import ZipFile
+
+import sqlite3
+import time
+
+import gradio as gr
+from tqdm.auto import tqdm
+import pandas as pd
+import numpy as np
+import os
+from datetime import timedelta, datetime
+from pathlib import Path
+
+import torch
+from sklearn.utils import shuffle
+
+from model import Collection, init_w, FSRS, WeightClipper, lineToTensor
+
+
+# Extract the collection file or deck file to get the .anki21 database.
+
+
+def extract(file, prefix):
+    proj_dir = Path(f'projects/{prefix}_{file.orig_name.replace(".", "_").replace("@", "_")}')
+    with ZipFile(file, 'r') as zip_ref:
+        zip_ref.extractall(proj_dir)
+        print(f"Extracted {file.orig_name} successfully!")
+    return proj_dir
+
+
+def create_time_series_features(revlog_start_date, timezone, next_day_starts_at, proj_dir,
+                                progress=gr.Progress(track_tqdm=True)):
+    if os.path.isfile(proj_dir / "collection.anki21b"):
+        os.remove(proj_dir / "collection.anki21b")
+        raise Exception(
+                "Please export the file with `support older Anki versions` if you use the latest version of Anki.")
+    elif os.path.isfile(proj_dir / "collection.anki21"):
+        con = sqlite3.connect(proj_dir / "collection.anki21")
+    elif os.path.isfile(proj_dir / "collection.anki2"):
+        con = sqlite3.connect(proj_dir / "collection.anki2")
+    else:
+        raise Exception("Collection not exist!")
+    cur = con.cursor()
+    res = cur.execute("SELECT * FROM revlog")
+    revlog = res.fetchall()
+
+    df = pd.DataFrame(revlog)
+    df.columns = ['id', 'cid', 'usn', 'r', 'ivl',
+                  'last_lvl', 'factor', 'time', 'type']
+    df = df[(df['cid'] <= time.time() * 1000) &
+            (df['id'] <= time.time() * 1000) &
+            (df['r'] > 0) &
+            (df['id'] >= time.mktime(datetime.strptime(revlog_start_date, "%Y-%m-%d").timetuple()) * 1000)].copy()
+    df['create_date'] = pd.to_datetime(df['cid'] // 1000, unit='s')
+    df['create_date'] = df['create_date'].dt.tz_localize(
+            'UTC').dt.tz_convert(timezone)
+    df['review_date'] = pd.to_datetime(df['id'] // 1000, unit='s')
+    df['review_date'] = df['review_date'].dt.tz_localize(
+            'UTC').dt.tz_convert(timezone)
+    df.drop(df[df['review_date'].dt.year < 2006].index, inplace=True)
+    df.sort_values(by=['cid', 'id'], inplace=True, ignore_index=True)
+    type_sequence = np.array(df['type'])
+    df.to_csv(proj_dir / "revlog.csv", index=False)
+    print("revlog.csv saved.")
+    df = df[(df['type'] == 0) | (df['type'] == 1)].copy()
+    df['real_days'] = df['review_date'] - timedelta(hours=next_day_starts_at)
+    df['real_days'] = pd.DatetimeIndex(df['real_days'].dt.floor('D')).to_julian_date()
+    df.drop_duplicates(['cid', 'real_days'], keep='first', inplace=True)
+    df['delta_t'] = df.real_days.diff()
+    df.dropna(inplace=True)
+    df['delta_t'] = df['delta_t'].astype(dtype=int)
+    df['i'] = 1
+    df['r_history'] = ""
+    df['t_history'] = ""
+    col_idx = {key: i for i, key in enumerate(df.columns)}
+
+    # code from https://github.com/L-M-Sherlock/anki_revlog_analysis/blob/main/revlog_analysis.py
+    def get_feature(x):
+        for idx, log in enumerate(x.itertuples()):
+            if idx == 0:
+                x.iloc[idx, col_idx['delta_t']] = 0
+            if idx == x.shape[0] - 1:
+                break
+            x.iloc[idx + 1, col_idx['i']] = x.iloc[idx, col_idx['i']] + 1
+            x.iloc[idx + 1, col_idx[
+                't_history']] = f"{x.iloc[idx, col_idx['t_history']]},{x.iloc[idx, col_idx['delta_t']]}"
+            x.iloc[idx + 1, col_idx['r_history']] = f"{x.iloc[idx, col_idx['r_history']]},{x.iloc[idx, col_idx['r']]}"
+        return x
+
+    tqdm.pandas(desc='Saving Trainset')
+    df = df.groupby('cid', as_index=False).progress_apply(get_feature)
+    df["t_history"] = df["t_history"].map(lambda x: x[1:] if len(x) > 1 else x)
+    df["r_history"] = df["r_history"].map(lambda x: x[1:] if len(x) > 1 else x)
+    df.to_csv(proj_dir / 'revlog_history.tsv', sep="\t", index=False)
+    print("Trainset saved.")
+
+    def cal_retention(group: pd.DataFrame) -> pd.DataFrame:
+        group['retention'] = round(group['r'].map(lambda x: {1: 0, 2: 1, 3: 1, 4: 1}[x]).mean(), 4)
+        group['total_cnt'] = group.shape[0]
+        return group
+
+    tqdm.pandas(desc='Calculating Retention')
+    df = df.groupby(by=['r_history', 'delta_t']).progress_apply(cal_retention)
+    print("Retention calculated.")
+    df = df.drop(columns=['id', 'cid', 'usn', 'ivl', 'last_lvl', 'factor', 'time', 'type', 'create_date', 'review_date',
+                          'real_days', 'r', 't_history'])
+    df.drop_duplicates(inplace=True)
+    df = df[(df['retention'] < 1) & (df['retention'] > 0)]
+
+    def cal_stability(group: pd.DataFrame) -> pd.DataFrame:
+        if group['i'].values[0] > 1:
+            r_ivl_cnt = sum(group['delta_t'] * group['retention'].map(np.log) * pow(group['total_cnt'], 2))
+            ivl_ivl_cnt = sum(group['delta_t'].map(lambda x: x ** 2) * pow(group['total_cnt'], 2))
+            group['stability'] = round(np.log(0.9) / (r_ivl_cnt / ivl_ivl_cnt), 1)
+        else:
+            group['stability'] = 0.0
+        group['group_cnt'] = sum(group['total_cnt'])
+        group['avg_retention'] = round(
+                sum(group['retention'] * pow(group['total_cnt'], 2)) / sum(pow(group['total_cnt'], 2)), 3)
+        group['avg_interval'] = round(
+                sum(group['delta_t'] * pow(group['total_cnt'], 2)) / sum(pow(group['total_cnt'], 2)), 1)
+        del group['total_cnt']
+        del group['retention']
+        del group['delta_t']
+        return group
+
+    tqdm.pandas(desc='Calculating Stability')
+    df = df.groupby(by=['r_history']).progress_apply(cal_stability)
+    print("Stability calculated.")
+    df.reset_index(drop=True, inplace=True)
+    df.drop_duplicates(inplace=True)
+    df.sort_values(by=['r_history'], inplace=True, ignore_index=True)
+
+    df_out = pd.DataFrame()
+    if df.shape[0] > 0:
+        for idx in tqdm(df.index):
+            item = df.loc[idx]
+            index = df[(df['i'] == item['i'] + 1) & (df['r_history'].str.startswith(item['r_history']))].index
+            df.loc[index, 'last_stability'] = item['stability']
+        df['factor'] = round(df['stability'] / df['last_stability'], 2)
+        df = df[(df['i'] >= 2) & (df['group_cnt'] >= 100)]
+        df['last_recall'] = df['r_history'].map(lambda x: x[-1])
+        df = df[df.groupby(['i', 'r_history'])['group_cnt'].transform(max) == df['group_cnt']]
+        df.to_csv(proj_dir / 'stability_for_analysis.tsv', sep='\t', index=None)
+        print("1:again, 2:hard, 3:good, 4:easy\n")
+        print(df[df['r_history'].str.contains(r'^[1-4][^124]*$', regex=True)][
+            ['r_history', 'avg_interval', 'avg_retention', 'stability', 'factor', 'group_cnt']].to_string(
+                index=False))
+        print("Analysis saved!")
+
+        df_out = df[df['r_history'].str.contains(r'^[1-4][^124]*$', regex=True)][
+            ['r_history', 'avg_interval', 'avg_retention', 'stability', 'factor', 'group_cnt']]
+    return type_sequence, df_out
+
+
+def train_model(proj_dir, progress=gr.Progress(track_tqdm=True)):
+    model = FSRS(init_w)
+
+    clipper = WeightClipper()
+    optimizer = torch.optim.Adam(model.parameters(), lr=5e-4)
+
+    dataset = pd.read_csv(proj_dir / "revlog_history.tsv", sep='\t', index_col=None,
+                          dtype={'r_history': str, 't_history': str})
+    dataset = dataset[(dataset['i'] > 1) & (dataset['delta_t'] > 0) & (dataset['t_history'].str.count(',0') == 0)]
+
+    tqdm.pandas(desc='Tensorizing Line')
+    dataset['tensor'] = dataset.progress_apply(lambda x: lineToTensor(list(zip([x['t_history']], [x['r_history']]))[0]),
+                                               axis=1)
+    print("Tensorized!")
+
+    pre_train_set = dataset[dataset['i'] == 2]
+    # pretrain
+    epoch_len = len(pre_train_set)
+    n_epoch = 1
+    pbar = tqdm(desc="Pre-training", colour="red", total=epoch_len * n_epoch)
+
+    for k in range(n_epoch):
+        for i, (_, row) in enumerate(shuffle(pre_train_set, random_state=2022 + k).iterrows()):
+            model.train()
+            optimizer.zero_grad()
+            output_t = [(model.zero, model.zero)]
+            for input_t in row['tensor']:
+                output_t.append(model(input_t, *output_t[-1]))
+            loss = model.loss(output_t[-1][0], row['delta_t'],
+                              {1: 0, 2: 1, 3: 1, 4: 1}[row['r']])
+            if np.isnan(loss.data.item()):
+                # Exception Case
+                print(row, output_t)
+                raise Exception('error case')
+            loss.backward()
+            optimizer.step()
+            model.apply(clipper)
+            pbar.update()
+    pbar.close()
+    for name, param in model.named_parameters():
+        print(f"{name}: {list(map(lambda x: round(float(x), 4), param))}")
+
+    train_set = dataset[dataset['i'] > 2]
+    epoch_len = len(train_set)
+    n_epoch = 1
+    print_len = max(epoch_len * n_epoch // 10, 1)
+    pbar = tqdm(desc="Training", total=epoch_len * n_epoch)
+
+    for k in range(n_epoch):
+        for i, (_, row) in enumerate(shuffle(train_set, random_state=2022 + k).iterrows()):
+            model.train()
+            optimizer.zero_grad()
+            output_t = [(model.zero, model.zero)]
+            for input_t in row['tensor']:
+                output_t.append(model(input_t, *output_t[-1]))
+            loss = model.loss(output_t[-1][0], row['delta_t'],
+                              {1: 0, 2: 1, 3: 1, 4: 1}[row['r']])
+            if np.isnan(loss.data.item()):
+                # Exception Case
+                print(row, output_t)
+                raise Exception('error case')
+            loss.backward()
+            for param in model.parameters():
+                param.grad[:2] = torch.zeros(2)
+            optimizer.step()
+            model.apply(clipper)
+            pbar.update()
+
+            if (k * epoch_len + i) % print_len == 0:
+                print(f"iteration: {k * epoch_len + i + 1}")
+                for name, param in model.named_parameters():
+                    print(f"{name}: {list(map(lambda x: round(float(x), 4), param))}")
+    pbar.close()
+
+    w = list(map(lambda x: round(float(x), 4), dict(model.named_parameters())['w'].data))
+
+    print("\nTraining finished!")
+    return w, dataset
+
+
+def process_personalized_collection(requestRetention, w):
+    my_collection = Collection(w)
+    rating_dict = {1: "again", 2: "hard", 3: "good", 4: "easy"}
+    rating_markdown = []
+    for first_rating in (1, 2, 3, 4):
+        rating_markdown.append(f'## First Rating: {first_rating} ({rating_dict[first_rating]})')
+        t_history = "0"
+        d_history = "0"
+        r_history = f"{first_rating}"  # the first rating of the new card
+        # print("stability, difficulty, lapses")
+        for i in range(10):
+            states = my_collection.states(t_history, r_history)
+            # print('{0:9.2f} {1:11.2f} {2:7.0f}'.format(
+            # *list(map(lambda x: round(float(x), 4), states))))
+            next_t = max(round(float(np.log(requestRetention) / np.log(0.9) * states[0])), 1)
+            difficulty = round(float(states[1]), 1)
+            t_history += f',{int(next_t)}'
+            d_history += f',{difficulty}'
+            r_history += f",3"
+        rating_markdown.append(f"*rating history*: {r_history}")
+        rating_markdown.append(f"*interval history*: {t_history}")
+        rating_markdown.append(f"*difficulty history*: {d_history}\n")
+    rating_markdown = '\n\n'.join(rating_markdown)
+    return my_collection, rating_markdown
+
+
+def log_loss(my_collection, row):
+    states = my_collection.states(row['t_history'], row['r_history'])
+    row['log_loss'] = float(my_collection.model.loss(states[0], row['delta_t'], {1: 0, 2: 1, 3: 1, 4: 1}[row['r']]))
+    return row
+
+
+def my_loss(dataset, w):
+    my_collection = Collection(init_w)
+    tqdm.pandas(desc='Calculating Loss before Training')
+    dataset = dataset.progress_apply(partial(log_loss, my_collection), axis=1)
+    print(f"Loss before training: {dataset['log_loss'].mean():.4f}")
+    my_collection = Collection(w)
+    tqdm.pandas(desc='Calculating Loss After Training')
+    dataset = dataset.progress_apply(partial(log_loss, my_collection), axis=1)
+    print(f"Loss after training: {dataset['log_loss'].mean():.4f}")
+    return f"""
+*Loss before training*: {dataset['log_loss'].mean():.4f}
+
+*Loss after training*: {dataset['log_loss'].mean():.4f}
+    """
+
+
+def cleanup(proj_dir: Path, files):
+    """
+    Delete all files in prefix that dont have filenames in files
+    :param proj_dir:
+    :param files:
+    :return:
+    """
+    for file in proj_dir.glob('*'):
+        if file.name not in files:
+            os.remove(file)
+