From 13f0192617be084c36083ba711f3d125f4a2e24b Mon Sep 17 00:00:00 2001 From: Kerwin Date: Sun, 8 Feb 2026 20:18:57 +0800 Subject: [PATCH] feat: implement select_dataset() for dynamic dataset filtering --- slime/rollout/data_source.py | 112 ++++++++++++++++++++++++++++++++++- slime/utils/data.py | 4 ++ slime/utils/types.py | 1 + 3 files changed, 116 insertions(+), 1 deletion(-) diff --git a/slime/rollout/data_source.py b/slime/rollout/data_source.py index fa14c65f00..e9a28fe0ed 100644 --- a/slime/rollout/data_source.py +++ b/slime/rollout/data_source.py @@ -78,8 +78,15 @@ def __init__(self, args): ) if self.args.rollout_shuffle: self.dataset.shuffle(self.epoch_id) + + # Dynamic dataset support: immutable full dataset for adjust_dataset to read, + # and for select_dataset to re-filter from. + self.full_samples: list[Sample] = list(self.dataset.origin_samples) + self._last_difficulty_mapping: dict[int, int] | None = None else: self.dataset = None + self.full_samples = [] + self._last_difficulty_mapping = None def get_samples(self, num_samples): # TODO further improve code @@ -111,6 +118,102 @@ def get_samples(self, num_samples): samples.append(group) return samples + def select_dataset(self, difficulty_mapping: dict[int, int], + keep_min: int = 1, keep_max: int = 6) -> dict: + """ + Filter the training dataset based on difficulty scores. + + Retains only samples whose difficulty falls within [keep_min, keep_max]. + Samples not present in difficulty_mapping are kept by default (unevaluated). + + Args: + difficulty_mapping: Maps Sample.global_index to difficulty value + (integer 0-8, number of successful rollouts out of 8). + keep_min: Minimum difficulty to retain (inclusive). Default 1. + keep_max: Maximum difficulty to retain (inclusive). Default 6. + + Returns: + A report dict with statistics about the adjustment. + """ + if self.dataset is None: + logger.warning("select_dataset called but dataset is None. Skipping.") + return {"skipped": True} + + self._last_difficulty_mapping = difficulty_mapping + + total_full = len(self.full_samples) + total_before = len(self.dataset.origin_samples) + prev_ids = set(id(s) for s in self.dataset.origin_samples) + + removed_easy = 0 + removed_hard = 0 + retained = 0 + unevaluated = 0 + restored = 0 + + new_origin = [] + for sample in self.full_samples: + gidx = sample.global_index + if gidx not in difficulty_mapping: + new_origin.append(sample) + unevaluated += 1 + continue + + difficulty = difficulty_mapping[gidx] + if difficulty > keep_max: + removed_easy += 1 + elif difficulty < keep_min: + removed_hard += 1 + else: + new_origin.append(sample) + if id(sample) not in prev_ids: + restored += 1 + else: + retained += 1 + + # Safety: if ALL samples would be removed, keep dataset unchanged + if len(new_origin) == 0: + logger.warning( + "select_dataset: filtering would remove ALL samples. " + "Keeping dataset unchanged." + ) + return { + "total_full": total_full, + "total_before": total_before, + "total_after": total_before, + "removed_easy": 0, + "removed_hard": 0, + "retained": total_before, + "restored": 0, + "unevaluated": 0, + "warning": "all_removed_fallback", + } + + # Update BOTH origin_samples and samples to avoid the shuffle trap: + # Dataset.shuffle() uses len(self.samples) for permutation but indexes + # into origin_samples, so both must have the same filtered content. + self.dataset.origin_samples = new_origin + self.dataset.samples = list(new_origin) + self.dataset.epoch_id = -1 # force next shuffle to execute + + self.sample_offset = 0 + + if self.args.rollout_shuffle: + self.dataset.shuffle(self.epoch_id) + + report = { + "total_full": total_full, + "total_before": total_before, + "total_after": len(self.dataset.origin_samples), + "removed_easy": removed_easy, + "removed_hard": removed_hard, + "retained": retained, + "restored": restored, + "unevaluated": unevaluated, + } + logger.info(f"select_dataset: {report}") + return report + def add_samples(self, samples: list[list[Sample]]): raise RuntimeError(f"Cannot add samples to {self.__class__.__name__}. This is a read-only data source.") @@ -124,6 +227,7 @@ def save(self, rollout_id): "sample_group_index": self.sample_group_index, "sample_index": self.sample_index, "metadata": self.metadata, + "last_difficulty_mapping": self._last_difficulty_mapping, } path = os.path.join(self.args.save, f"rollout/global_dataset_state_dict_{rollout_id}.pt") os.makedirs(os.path.dirname(path), exist_ok=True) @@ -150,7 +254,13 @@ def load(self, rollout_id=None): self.sample_index = state_dict.get("sample_index", 0) self.metadata = state_dict.get("metadata", {}) - if self.args.rollout_global_dataset and self.args.rollout_shuffle: + # Restore dynamic dataset state by replaying the last filtering operation. + self._last_difficulty_mapping = state_dict.get("last_difficulty_mapping", None) + if self._last_difficulty_mapping is not None: + self.select_dataset(self._last_difficulty_mapping) + # select_dataset resets sample_offset to 0; restore the exact checkpoint offset. + self.sample_offset = state_dict.get("sample_offset", 0) + elif self.args.rollout_global_dataset and self.args.rollout_shuffle: self.dataset.shuffle(self.epoch_id) diff --git a/slime/utils/data.py b/slime/utils/data.py index 0b9d6a6db4..ab47c68453 100644 --- a/slime/utils/data.py +++ b/slime/utils/data.py @@ -234,6 +234,10 @@ def __init__( else: self.origin_samples = origin_samples + # Assign stable global_index after length filtering + for idx, sample in enumerate(self.origin_samples): + sample.global_index = idx + self.epoch_id = -1 self.seed = seed self.samples = self.origin_samples diff --git a/slime/utils/types.py b/slime/utils/types.py index dc3b43d23b..a8ff7f1322 100644 --- a/slime/utils/types.py +++ b/slime/utils/types.py @@ -9,6 +9,7 @@ class Sample: """The sample generated""" + global_index: int | None = None # stable index assigned at data loading, invariant across rollouts group_index: int | None = None index: int | None = None # prompt