diff --git a/pyproject.toml b/pyproject.toml index 01630bf..3e90468 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "FSRS-Optimizer" -version = "5.2.1" +version = "5.2.2" readme = "README.md" dependencies = [ "matplotlib>=3.7.0", diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py index 2178198..2f61cd2 100644 --- a/src/fsrs_optimizer/fsrs_optimizer.py +++ b/src/fsrs_optimizer/fsrs_optimizer.py @@ -5,7 +5,7 @@ import numpy as np import os import math -from typing import List, Optional +from typing import List, Optional, Tuple from datetime import timedelta, datetime from collections import defaultdict import statsmodels.api as sm # type: ignore @@ -42,25 +42,25 @@ Relearning = 3 DEFAULT_PARAMETER = [ - 0.4072, - 1.1829, - 3.1262, - 15.4722, - 7.2102, - 0.5316, - 1.0651, - 0.0234, - 1.616, - 0.1544, - 1.0824, - 1.9813, - 0.0953, - 0.2975, - 2.2042, - 0.2407, - 2.9466, - 0.5034, - 0.6567, + 0.40255, + 1.18385, + 3.173, + 15.69105, + 7.1949, + 0.5345, + 1.4604, + 0.0046, + 1.54575, + 0.1192, + 1.01925, + 1.9395, + 0.11, + 0.29605, + 2.2698, + 0.2315, + 2.9898, + 0.51655, + 0.6621, ] S_MIN = 0.01 @@ -105,8 +105,12 @@ def init_d(self, rating: Tensor) -> Tensor: new_d = self.w[4] - torch.exp(self.w[5] * (rating - 1)) + 1 return new_d + def linear_damping(self, delta_d: Tensor, old_d: Tensor) -> Tensor: + return delta_d * (10 - old_d) / 9 + def next_d(self, state: Tensor, rating: Tensor) -> Tensor: - new_d = state[:, 1] - self.w[6] * (rating - 3) + delta_d = -self.w[6] * (rating - 3) + new_d = state[:, 1] + self.linear_damping(delta_d, state[:, 1]) new_d = self.mean_reversion(self.init_d(4), new_d) return new_d @@ -151,7 +155,9 @@ def step(self, X: Tensor, state: Tensor) -> Tensor: new_s = new_s.clamp(S_MIN, 36500) return torch.stack([new_s, new_d], dim=1) - def forward(self, inputs: Tensor, state: Optional[Tensor] = None) -> Tensor: + def forward( + self, inputs: Tensor, state: Optional[Tensor] = None + ) -> Tuple[Tensor, Tensor]: """ :param inputs: shape[seq_len, batch_size, 2] """ @@ -179,16 +185,16 @@ def __call__(self, module): w[2] = w[2].clamp(S_MIN, 100) w[3] = w[3].clamp(S_MIN, 100) w[4] = w[4].clamp(1, 10) - w[5] = w[5].clamp(0.01, 4) - w[6] = w[6].clamp(0.01, 4) - w[7] = w[7].clamp(0, 0.75) + w[5] = w[5].clamp(0.001, 4) + w[6] = w[6].clamp(0.001, 4) + w[7] = w[7].clamp(0.001, 0.75) w[8] = w[8].clamp(0, 4.5) w[9] = w[9].clamp(0, 0.8) - w[10] = w[10].clamp(0.01, 3.5) - w[11] = w[11].clamp(0.1, 5) - w[12] = w[12].clamp(0.01, 0.25) - w[13] = w[13].clamp(0.01, 0.9) - w[14] = w[14].clamp(0.01, 4) + w[10] = w[10].clamp(0.001, 3.5) + w[11] = w[11].clamp(0.001, 5) + w[12] = w[12].clamp(0.001, 0.25) + w[13] = w[13].clamp(0.001, 0.9) + w[14] = w[14].clamp(0, 4) w[15] = w[15].clamp(0, 1) w[16] = w[16].clamp(1, 6) w[17] = w[17].clamp(0, 2) @@ -2075,22 +2081,3 @@ def wrap_short_term_ratings(r_history, t_history): else: result.pop() return "".join(result) - - -if __name__ == "__main__": - model = FSRS(DEFAULT_PARAMETER) - stability = torch.tensor([5.0] * 4) - difficulty = torch.tensor([1.0, 2.0, 3.0, 4.0]) - retention = torch.tensor([0.9, 0.8, 0.7, 0.6]) - rating = torch.tensor([1, 2, 3, 4]) - state = torch.stack([stability, difficulty]).unsqueeze(0) - s_recall = model.stability_after_success(state, retention, rating) - print(s_recall) - s_forget = model.stability_after_failure(state, retention) - print(s_forget) - - retentions = torch.tensor([0.1, 0.2, 0.3, 0.4]) - labels = torch.tensor([0.0, 1.0, 0.0, 1.0]) - loss_fn = nn.BCELoss() - loss = loss_fn(retentions, labels) - print(loss) diff --git a/src/fsrs_optimizer/fsrs_simulator.py b/src/fsrs_optimizer/fsrs_simulator.py index aa96958..77df428 100644 --- a/src/fsrs_optimizer/fsrs_simulator.py +++ b/src/fsrs_optimizer/fsrs_simulator.py @@ -127,8 +127,12 @@ def init_d_with_short_term(rating): new_d = init_d(rating) - w[6] * rating_offset return np.clip(new_d, 1, 10) + def linear_damping(delta_d, old_d): + return delta_d * (10 - old_d) / 9 + def next_d(d, rating): - new_d = d - w[6] * (rating - 3) + delta_d = -w[6] * (rating - 3) + new_d = d + linear_damping(delta_d, d) new_d = mean_reversion(init_d(4), new_d) return np.clip(new_d, 1, 10) diff --git a/tests/model_test.py b/tests/model_test.py new file mode 100644 index 0000000..69e6972 --- /dev/null +++ b/tests/model_test.py @@ -0,0 +1,137 @@ +from src.fsrs_optimizer import * + + +class Test_Model: + def test_next_stability(self): + model = FSRS(DEFAULT_PARAMETER) + stability = torch.tensor([5.0] * 4) + difficulty = torch.tensor([1.0, 2.0, 3.0, 4.0]) + retention = torch.tensor([0.9, 0.8, 0.7, 0.6]) + rating = torch.tensor([1, 2, 3, 4]) + state = torch.stack([stability, difficulty]).unsqueeze(0) + s_recall = model.stability_after_success(state, retention, rating) + assert torch.allclose( + s_recall, torch.tensor([25.7761, 14.1219, 60.4044, 208.9760]), atol=1e-4 + ) + s_forget = model.stability_after_failure(state, retention) + assert torch.allclose( + s_forget, torch.tensor([1.7029, 1.9799, 2.3760, 2.8885]), atol=1e-4 + ) + s_short_term = model.stability_short_term(state, rating) + assert torch.allclose( + s_short_term, torch.tensor([2.5051, 4.1992, 7.0389, 11.7988]), atol=1e-4 + ) + + def test_next_difficulty(self): + model = FSRS(DEFAULT_PARAMETER) + stability = torch.tensor([5.0] * 4) + difficulty = torch.tensor([5.0] * 4) + rating = torch.tensor([1, 2, 3, 4]) + state = torch.stack([stability, difficulty]).unsqueeze(0) + d_recall = model.next_d(state, rating) + assert torch.allclose( + d_recall, + torch.tensor([6.6070, 5.7994, 4.9918, 4.1842]), + atol=1e-4, + ) + + def test_power_forgetting_curve(self): + delta_t = torch.tensor([0, 1, 2, 3, 4, 5]) + stability = torch.tensor([1, 2, 3, 4, 4, 2]) + retention = power_forgetting_curve(delta_t, stability) + assert torch.allclose( + retention, + torch.tensor([1.0, 0.946059, 0.9299294, 0.9221679, 0.90000004, 0.79394597]), + atol=1e-4, + ) + + def test_forward(self): + model = FSRS(DEFAULT_PARAMETER) + delta_ts = torch.tensor( + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [1.0, 1.0, 1.0, 1.0, 2.0, 2.0], + ] + ) + ratings = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0, 1.0, 2.0], + [1.0, 2.0, 3.0, 4.0, 1.0, 2.0], + ] + ) + inputs = torch.stack([delta_ts, ratings], dim=2) + _, state = model.forward(inputs) + stability = state[:, 0] + difficulty = state[:, 1] + assert torch.allclose( + stability, + torch.tensor([0.2619, 1.7073, 5.8691, 25.0123, 0.3403, 2.1482]), + atol=1e-4, + ) + assert torch.allclose( + difficulty, + torch.tensor([8.0827, 7.0405, 5.2729, 2.1301, 8.0827, 7.0405]), + atol=1e-4, + ) + + def test_loss_and_grad(self): + model = FSRS(DEFAULT_PARAMETER) + loss_fn = nn.BCELoss(reduction="none") + t_histories = torch.tensor( + [ + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + [0.0, 1.0, 1.0, 3.0], + [1.0, 3.0, 3.0, 5.0], + [3.0, 6.0, 6.0, 12.0], + ] + ) + r_histories = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0], + [3.0, 4.0, 2.0, 4.0], + [1.0, 4.0, 4.0, 3.0], + [4.0, 3.0, 3.0, 3.0], + [3.0, 1.0, 3.0, 3.0], + [2.0, 3.0, 3.0, 4.0], + ] + ) + delta_ts = torch.tensor([4.0, 11.0, 12.0, 23.0]) + labels = torch.tensor([1, 1, 1, 0], dtype=torch.float32, requires_grad=False) + inputs = torch.stack([t_histories, r_histories], dim=2) + seq_lens = inputs.shape[0] + real_batch_size = inputs.shape[1] + outputs, _ = model.forward(inputs) + stabilities = outputs[seq_lens - 1, torch.arange(real_batch_size), 0] + retentions = power_forgetting_curve(delta_ts, stabilities) + loss = loss_fn(retentions, labels).sum() + assert round(loss.item(), 4) == 4.4467 + loss.backward() + assert torch.allclose( + model.w.grad, + torch.tensor( + [ + -0.0583, + -0.0068, + -0.0026, + 0.0105, + -0.0513, + 1.3643, + 0.0837, + -0.9502, + 0.5345, + -2.8929, + 0.5142, + -0.0131, + 0.0419, + -0.1183, + -0.0009, + -0.1445, + 0.2024, + 0.2141, + 0.0323, + ] + ), + atol=1e-4, + ) diff --git a/tests/simulator_test.py b/tests/simulator_test.py index 3afb02c..3f4cf2a 100644 --- a/tests/simulator_test.py +++ b/tests/simulator_test.py @@ -10,7 +10,7 @@ def test_simulate(self): memorized_cnt_per_day, cost_per_day, ) = simulate(w=DEFAULT_PARAMETER, request_retention=0.9) - assert memorized_cnt_per_day[-1] == 5918.574208243532 + assert memorized_cnt_per_day[-1] == 5875.025236206539 def test_optimal_retention(self): default_params = { @@ -24,4 +24,4 @@ def test_optimal_retention(self): "loss_aversion": 2.5, } r = optimal_retention(**default_params) - assert r == 0.8346739534878145 + assert r == 0.8263932