Merge branch 'fix/make_cold_user_splitter_behavior_stable' into 'main'

Fix ColdUserRandomSplitter._core_split_polars randomness Closes #140 See merge request ai-lab-pmo/mltools/recsys/RePlay!247
sb-ai-lab · Feb 26, 2025 · 6fea0e1 · 6fea0e1
2 parents 940f8cd + 45a6ac7
commit 6fea0e1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/replay/splitters/cold_user_random_splitter.py b/replay/splitters/cold_user_random_splitter.py
@@ -94,7 +94,7 @@ def _core_split_polars(
     ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
         train_users = (
             interactions.select(self.query_column)
-            .unique()
+            .unique(maintain_order=True)
             .sample(fraction=(1 - threshold), seed=self.seed)
             .with_columns(pl.lit(False).alias("is_test"))
         )