reset S0_dataset_group if dataset is not None

open-spaced-repetition · Jan 2, 2025 · f464c06 · f464c06
1 parent 96cea5f
commit f464c06
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py
@@ -1003,12 +1003,6 @@ def pretrain(self, dataset=None, verbose=True):
             )
         else:
             self.dataset = dataset
-        self.dataset = self.dataset[
-            (self.dataset["i"] > 1) & (self.dataset["delta_t"] > 0)
-        ]
-        if self.dataset.empty:
-            raise ValueError("Training data is inadequate.")
-        if self.S0_dataset_group is None:
             self.dataset["first_rating"] = self.dataset["r_history"].map(
                 lambda x: x[0] if len(x) > 0 else ""
             )
@@ -1018,6 +1012,11 @@ def pretrain(self, dataset=None, verbose=True):
                 .agg({"y": ["mean", "count"]})
                 .reset_index()
             )
+        self.dataset = self.dataset[
+            (self.dataset["i"] > 1) & (self.dataset["delta_t"] > 0)
+        ]
+        if self.dataset.empty:
+            raise ValueError("Training data is inadequate.")
         rating_stability = {}
         rating_count = {}
         average_recall = self.dataset["y"].mean()