don't sort by length of reviews at first

Author: L-M-Sherlock

src/batch_shuffle.rs

@@ -105,6 +105,7 @@ mod tests {
         backend::{ndarray::NdArrayDevice, NdArray},
         tensor::Shape,
     };
+    use itertools::Itertools;
 
     use super::*;
     use crate::{
@@ -114,7 +115,10 @@ mod tests {
 
     #[test]
     fn test_simple_dataloader() {
-        let train_set = anki21_sample_file_converted_to_fsrs();
+        let train_set = anki21_sample_file_converted_to_fsrs()
+            .into_iter()
+            .sorted_by_cached_key(|item| item.reviews.len())
+            .collect();
         let (_pre_train_set, train_set) = prepare_training_data(train_set);
         let dataset = FSRSDataset::from(simple_weighted_fsrs_items(train_set));
         let batch_size = 512;

src/convertor_tests.rs

@@ -94,7 +94,7 @@ fn convert_to_fsrs_items(
     mut entries: Vec<RevlogEntry>,
     next_day_starts_at: i64,
     timezone: Tz,
-) -> Option<Vec<FSRSItem>> {
+) -> Option<Vec<(i64, FSRSItem)>> {
     // entries = filter_out_cram(entries);
     // entries = filter_out_manual(entries);
     entries = remove_revlog_before_last_first_learn(entries);
@@ -110,7 +110,7 @@ fn convert_to_fsrs_items(
             .iter()
             .enumerate()
             .skip(1)
-            .map(|(idx, _)| {
+            .map(|(idx, entry)| {
                 let reviews = entries
                     .iter()
                     .take(idx + 1)
@@ -119,9 +119,9 @@ fn convert_to_fsrs_items(
                         delta_t: r.last_interval.max(0) as u32,
                     })
                     .collect();
-                FSRSItem { reviews }
+                (entry.id, FSRSItem { reviews })
             })
-            .filter(|item| item.current().delta_t > 0)
+            .filter(|(_, item)| item.current().delta_t > 0)
             .collect(),
     )
 }
@@ -137,8 +137,8 @@ pub(crate) fn anki_to_fsrs(revlogs: Vec<RevlogEntry>) -> Vec<FSRSItem> {
         })
         .flatten()
         .collect_vec();
-    revlogs.sort_by_cached_key(|r| r.reviews.len());
-    revlogs
+    revlogs.sort_by_cached_key(|(id, _)| *id);
+    revlogs.into_iter().map(|(_, item)| item).collect()
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -260,6 +260,7 @@ fn conversion_works() {
         .into_iter()
         .filter_map(|entries| convert_to_fsrs_items(entries, 4, Tz::Asia__Shanghai))
         .flatten()
+        .map(|(_, item)| item)
         .collect_vec();
     assert_eq!(
         fsrs_items,
@@ -445,7 +446,8 @@ fn delta_t_is_correct() -> Result<()> {
             ],
             NEXT_DAY_AT,
             Tz::Asia__Shanghai
-        ),
+        )
+        .map(|items| items.into_iter().map(|(_, item)| item).collect_vec()),
         Some(vec![FSRSItem {
             reviews: vec![
                 FSRSReview {
@@ -470,7 +472,8 @@ fn delta_t_is_correct() -> Result<()> {
             ],
             NEXT_DAY_AT,
             Tz::Asia__Shanghai
-        ),
+        )
+        .map(|items| items.into_iter().map(|(_, item)| item).collect_vec()),
         Some(vec![
             FSRSItem {
                 reviews: vec![

src/dataset.rs

@@ -265,6 +265,12 @@ pub fn prepare_training_data(items: Vec<FSRSItem>) -> (Vec<FSRSItem>, Vec<FSRSIt
     (pretrainset.clone(), [pretrainset, trainset].concat())
 }
 
+pub(crate) fn sort_items_by_review_length(items: Vec<WeightedFSRSItem>) -> Vec<WeightedFSRSItem> {
+    let mut items = items;
+    items.sort_by_cached_key(|item| item.item.reviews.len());
+    items
+}
+
 pub(crate) fn simple_weighted_fsrs_items(items: Vec<FSRSItem>) -> Vec<WeightedFSRSItem> {
     items
         .into_iter()
@@ -294,21 +300,21 @@ mod tests {
     fn from_anki() {
         use burn::data::dataloader::Dataset;
 
-        let dataset = FSRSDataset::from(simple_weighted_fsrs_items(
+        let dataset = FSRSDataset::from(sort_items_by_review_length(simple_weighted_fsrs_items(
             anki21_sample_file_converted_to_fsrs(),
-        ));
+        )));
         assert_eq!(
             dataset.get(704).unwrap().item,
             FSRSItem {
                 reviews: vec![
                     FSRSReview {
-                        rating: 3,
-                        delta_t: 0,
+                        rating: 4,
+                        delta_t: 0
                     },
                     FSRSReview {
                         rating: 3,
-                        delta_t: 1,
-                    },
+                        delta_t: 3
+                    }
                 ],
             }
         );

src/training.rs

@@ -1,6 +1,9 @@
 use crate::batch_shuffle::{BatchTensorDataset, ShuffleDataLoader};
 use crate::cosine_annealing::CosineAnnealingLR;
-use crate::dataset::{prepare_training_data, recency_weighted_fsrs_items, FSRSDataset, FSRSItem};
+use crate::dataset::{
+    prepare_training_data, recency_weighted_fsrs_items, sort_items_by_review_length, FSRSDataset,
+    FSRSItem,
+};
 use crate::error::Result;
 use crate::model::{Model, ModelConfig};
 use crate::parameter_clipper::parameter_clipper;
@@ -238,7 +241,6 @@ impl<B: Backend> FSRS<B> {
             AdamConfig::new().with_epsilon(1e-8),
         );
         train_set.retain(|item| item.reviews.len() <= config.max_seq_len);
-        train_set.sort_by_cached_key(|item| item.reviews.len());
 
         if let Some(progress) = &progress {
             let progress_state = ProgressState {
@@ -308,7 +310,6 @@ impl<B: Backend> FSRS<B> {
             AdamConfig::new().with_epsilon(1e-8),
         );
         train_set.retain(|item| item.reviews.len() <= config.max_seq_len);
-        train_set.sort_by_cached_key(|item| item.reviews.len());
         let model =
             train::<Autodiff<B>>(train_set.clone(), train_set, &config, self.device(), None);
         let parameters: Vec<f32> = model.unwrap().w.val().to_data().convert().value;
@@ -328,14 +329,18 @@ fn train<B: AutodiffBackend>(
     // Training data
     let iterations = (train_set.len() / config.batch_size + 1) * config.num_epochs;
     let batch_dataset = BatchTensorDataset::<B>::new(
-        FSRSDataset::from(recency_weighted_fsrs_items(train_set)),
+        FSRSDataset::from(sort_items_by_review_length(recency_weighted_fsrs_items(
+            train_set,
+        ))),
         config.batch_size,
         device.clone(),
     );
     let dataloader_train = ShuffleDataLoader::new(batch_dataset, config.seed);
 
     let batch_dataset = BatchTensorDataset::<B::InnerBackend>::new(
-        FSRSDataset::from(recency_weighted_fsrs_items(test_set.clone())),
+        FSRSDataset::from(sort_items_by_review_length(recency_weighted_fsrs_items(
+            test_set.clone(),
+        ))),
         config.batch_size,
         device,
     );