chore: import upstream snapshot with attribution

swift/rl_core/resample.py

@@ -0,0 +1,72 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Shared resample logic for HF / Megatron / Megatron-Ray trainers.
+
+When ``truncation_strategy='delete'`` (or dynamic sampling), samples whose
+``template.encode`` fails (e.g. exceed ``max_length``, or multimodal processing
+errors) must be replaced with fresh ones drawn from a backup iterator until we
+have ``len(inputs)`` valid samples. The backends previously each carried a
+near-identical copy of this loop; they only differed in the iterator they pull
+from and whether the assistant response is stripped before encoding (prompt-only
+algorithms like GRPO strip; GKD off-policy distillation keeps it).
+"""
+from typing import Iterator, List
+
+from swift.template.base import Template
+from swift.utils import get_logger, remove_response
+
+logger = get_logger()
+
+
+def resample_encode_failed_inputs(
+    template: Template,
+    data_iterator: Iterator,
+    inputs: List[dict],
+    max_resample_rounds: int = 10,
+    strip_response: bool = True,
+) -> List[dict]:
+    """Replace samples whose encode fails with fresh ones from ``data_iterator``.
+
+    Caps the TOTAL encode attempts (fail-fast): a systematic failure
+    (e.g. ``max_length`` too small so every prompt is over-length) raises quickly
+    instead of churning through the iterator, and an empty batch breaks the loop
+    instead of spinning forever.
+
+    Args:
+        template: Template used to encode (and thereby validate) a sample.
+        data_iterator: Backup iterator yielding batches (lists) of fresh samples.
+        inputs: The current batch; its length is the required valid count.
+        max_resample_rounds: Resample budget; total attempts == required * (rounds + 1).
+        strip_response: Remove the assistant response (in place) before encoding.
+
+    Returns:
+        A list of valid samples with the same length as ``inputs``.
+
+    Raises:
+        RuntimeError: If not enough valid samples are collected after the budget.
+    """
+    required = len(inputs)
+    max_attempts = required * (max_resample_rounds + 1)
+    valid, pending = [], list(inputs)
+    attempts = n_dropped = 0
+
+    while len(valid) < required and attempts < max_attempts:
+        if not pending:
+            batch = list(next(data_iterator))
+            if not batch:  # guard: an empty batch would otherwise spin forever
+                break
+            pending.extend(batch)
+        item = pending.pop(0)
+        attempts += 1
+        try:
+            if strip_response:
+                remove_response(item['messages'])
+            template.encode(item)
+            valid.append(item)
+        except Exception as e:
+            n_dropped += 1
+            logger.info(f'Encoding failed for one sample; will resample. {e}')
+
+    if len(valid) < required:
+        raise RuntimeError(f'resample: only collected {len(valid)}/{required} valid samples after {attempts} '
+                           f'attempts ({n_dropped} failed). Increase `max_length` or adjust `truncation_strategy`.')
+    return valid[:required]