Feat/allow pass dataset into pretrain (#34)

open-spaced-repetition · Sep 6, 2023 · e17f59d · e17f59d
1 parent 8676b77
commit e17f59d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "FSRS-Optimizer"
-version = "4.13.2"
+version = "4.13.3"
 readme = "README.md"
 dependencies = [
     "matplotlib>=3.7.0",

diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py
@@ -511,8 +511,11 @@ def define_model(self):
         https://github.com/open-spaced-repetition/fsrs4anki/wiki/The-Algorithm
         '''
 
-    def pretrain(self, verbose=True):
-        self.dataset = pd.read_csv("./revlog_history.tsv", sep='\t', index_col=None, dtype={'r_history': str ,'t_history': str} )
+    def pretrain(self, dataset=None, verbose=True):
+        if dataset is None:
+            self.dataset = pd.read_csv("./revlog_history.tsv", sep='\t', index_col=None, dtype={'r_history': str ,'t_history': str} )
+        else:
+            self.dataset = dataset
         self.dataset = self.dataset[(self.dataset['i'] > 1) & (self.dataset['delta_t'] > 0) & (self.dataset['t_history'].str.count(',0') == 0)]
         if self.dataset.empty:
             raise ValueError('Training data is inadequate.')