dataset records separate

turbo-llm · Nov 23, 2024 · b38e9d9 · b38e9d9
1 parent 32b5c16
commit b38e9d9
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/turbo_alignment/dataset/chat/chat.py b/turbo_alignment/dataset/chat/chat.py
@@ -372,10 +372,11 @@ def get_slice(self, start: int, end: int) -> Self:
             random_cut=self._random_cut,
         )
 
-        new_instance.records = self.records[start:end]
+        dataset_records = [self[idx] for idx in range(len(self))]
 
+        new_instance.records = self.records[start:end]
         new_instance.original_records_map = {
-            record['id']: self.get_original_record_by_id(record['id']) for record in new_instance.records
+            record['id']: self.get_original_record_by_id(record['id']) for record in dataset_records
         }
 
         return new_instance