fsrs-optimizer / memory_states.py
derek-thomas's picture
derek-thomas HF staff
Comment out print statements
303cbb8
raw
history blame
1.46 kB
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