| from __future__ import annotations |
|
|
| import random |
| from dataclasses import dataclass |
| from typing import Iterable, Iterator, Sequence |
|
|
| from dovla_cil.data.datasets import CILDataset |
| from dovla_cil.data.schema import CILRecord |
| from dovla_cil.data.sharding import group_records |
|
|
|
|
| class BatchIndices(list[int]): |
| """Batch index list with same-state ranking metadata.""" |
|
|
| def __init__( |
| self, |
| indices: Iterable[int], |
| *, |
| group_ids: Iterable[str] = (), |
| pair_indices: Iterable[tuple[int, int]] = (), |
| ) -> None: |
| super().__init__(int(index) for index in indices) |
| self.group_ids = list(group_ids) |
| self.pair_indices = [tuple(pair) for pair in pair_indices] |
|
|
|
|
| class GroupAwareBatchSampler: |
| """Yield batches that preserve CIL same-state group structure. |
| |
| Modes: |
| - `full_group`: include complete groups, up to `batch_groups` groups per batch. |
| - `pairs`: sample positive/negative same-group pairs for ranking losses. |
| - `mixed`: sample `records_per_group` records per group while preferring reward-diverse sets. |
| """ |
|
|
| def __init__( |
| self, |
| dataset: CILDataset, |
| *, |
| mode: str = "full_group", |
| batch_groups: int = 1, |
| records_per_group: int | None = None, |
| pair_count_per_group: int = 1, |
| shuffle: bool = False, |
| seed: int = 17, |
| reward_margin: float = 1e-6, |
| group_ids: Sequence[str] | None = None, |
| ) -> None: |
| if mode not in {"full_group", "pairs", "mixed"}: |
| raise ValueError("mode must be 'full_group', 'pairs', or 'mixed'") |
| if batch_groups <= 0: |
| raise ValueError("batch_groups must be positive") |
| if records_per_group is not None and records_per_group <= 0: |
| raise ValueError("records_per_group must be positive when provided") |
| if pair_count_per_group <= 0: |
| raise ValueError("pair_count_per_group must be positive") |
| if reward_margin < 0: |
| raise ValueError("reward_margin must be non-negative") |
| self.dataset = dataset |
| self.mode = mode |
| self.batch_groups = int(batch_groups) |
| self.records_per_group = records_per_group |
| self.pair_count_per_group = int(pair_count_per_group) |
| self.shuffle = bool(shuffle) |
| self.seed = int(seed) |
| self.reward_margin = float(reward_margin) |
| unknown = [group_id for group_id in (group_ids or []) if group_id not in dataset.group_ids] |
| if unknown: |
| raise KeyError(f"Unknown CIL groups for sampler: {unknown[:3]}") |
| self.group_ids = list(group_ids) if group_ids is not None else list(dataset.group_ids) |
|
|
| def __iter__(self) -> Iterator[BatchIndices]: |
| rng = random.Random(self.seed) |
| group_ids = list(self.group_ids) |
| if self.shuffle: |
| rng.shuffle(group_ids) |
| for offset in range(0, len(group_ids), self.batch_groups): |
| selected_group_ids = group_ids[offset : offset + self.batch_groups] |
| if self.mode == "full_group": |
| yield self._full_group_batch(selected_group_ids) |
| elif self.mode == "pairs": |
| yield self._pair_batch(selected_group_ids, rng) |
| else: |
| yield self._mixed_batch(selected_group_ids, rng) |
|
|
| def __len__(self) -> int: |
| if not self.group_ids: |
| return 0 |
| return (len(self.group_ids) + self.batch_groups - 1) // self.batch_groups |
|
|
| def _full_group_batch(self, group_ids: Sequence[str]) -> BatchIndices: |
| indices: list[int] = [] |
| for group_id in group_ids: |
| indices.extend(self.dataset.group_indices(group_id)) |
| pair_indices = _local_pair_indices( |
| [self.dataset[index] for index in indices], reward_margin=self.reward_margin |
| ) |
| return BatchIndices(indices, group_ids=group_ids, pair_indices=pair_indices) |
|
|
| def _pair_batch(self, group_ids: Sequence[str], rng: random.Random) -> BatchIndices: |
| indices: list[int] = [] |
| pair_indices: list[tuple[int, int]] = [] |
| for group_id in group_ids: |
| group_indices = self.dataset.group_indices(group_id) |
| records = [self.dataset[index] for index in group_indices] |
| global_pairs = _sample_reward_ordered_pairs( |
| group_indices, |
| records, |
| pair_count=self.pair_count_per_group, |
| reward_margin=self.reward_margin, |
| rng=rng, |
| ) |
| local_by_global: dict[int, int] = {} |
| for better, worse in global_pairs: |
| if better not in local_by_global: |
| local_by_global[better] = len(indices) |
| indices.append(better) |
| if worse not in local_by_global: |
| local_by_global[worse] = len(indices) |
| indices.append(worse) |
| pair_indices.append((local_by_global[better], local_by_global[worse])) |
| return BatchIndices(indices, group_ids=group_ids, pair_indices=pair_indices) |
|
|
| def _mixed_batch(self, group_ids: Sequence[str], rng: random.Random) -> BatchIndices: |
| indices: list[int] = [] |
| for group_id in group_ids: |
| group_indices = self.dataset.group_indices(group_id) |
| records = [self.dataset[index] for index in group_indices] |
| selected = _sample_records_from_group( |
| group_indices, |
| records, |
| count=self.records_per_group or len(group_indices), |
| reward_margin=self.reward_margin, |
| rng=rng, |
| ) |
| indices.extend(selected) |
| pair_indices = _local_pair_indices( |
| [self.dataset[index] for index in indices], reward_margin=self.reward_margin |
| ) |
| return BatchIndices(indices, group_ids=group_ids, pair_indices=pair_indices) |
|
|
|
|
| @dataclass(frozen=True) |
| class GroupBatchSampler: |
| """Backward-compatible complete-group sampler over materialized records.""" |
|
|
| records: tuple[CILRecord, ...] |
| shuffle: bool = False |
| seed: int = 17 |
|
|
| @classmethod |
| def from_records( |
| cls, records: Iterable[CILRecord], *, shuffle: bool = False, seed: int = 17 |
| ) -> "GroupBatchSampler": |
| return cls(records=tuple(records), shuffle=shuffle, seed=seed) |
|
|
| def __iter__(self) -> Iterator[list[CILRecord]]: |
| groups = list(group_records(self.records).values()) |
| if self.shuffle: |
| rng = random.Random(self.seed) |
| rng.shuffle(groups) |
| yield from groups |
|
|
|
|
| def _sample_reward_ordered_pairs( |
| group_indices: list[int], |
| records: list[CILRecord], |
| *, |
| pair_count: int, |
| reward_margin: float, |
| rng: random.Random, |
| ) -> list[tuple[int, int]]: |
| candidates: list[tuple[int, int]] = [] |
| for left in range(len(records)): |
| for right in range(len(records)): |
| if left == right: |
| continue |
| reward_left = _ranking_reward(records[left]) |
| reward_right = _ranking_reward(records[right]) |
| if reward_left > reward_right + reward_margin: |
| candidates.append((group_indices[left], group_indices[right])) |
| rng.shuffle(candidates) |
| return candidates[:pair_count] |
|
|
|
|
| def _sample_records_from_group( |
| group_indices: list[int], |
| records: list[CILRecord], |
| *, |
| count: int, |
| reward_margin: float, |
| rng: random.Random, |
| ) -> list[int]: |
| if count >= len(group_indices): |
| return list(group_indices) |
| pairs = _sample_reward_ordered_pairs( |
| group_indices, |
| records, |
| pair_count=max(1, count // 2), |
| reward_margin=reward_margin, |
| rng=rng, |
| ) |
| selected: list[int] = [] |
| for better, worse in pairs: |
| for index in (better, worse): |
| if index not in selected: |
| selected.append(index) |
| if len(selected) >= count: |
| return selected |
| remaining = [index for index in group_indices if index not in selected] |
| rng.shuffle(remaining) |
| return selected + remaining[: max(0, count - len(selected))] |
|
|
|
|
| def _local_pair_indices(records: list[CILRecord], *, reward_margin: float) -> list[tuple[int, int]]: |
| pairs: list[tuple[int, int]] = [] |
| by_group = group_records(records) |
| offset_by_record_id = {record.record_id: index for index, record in enumerate(records)} |
| for group in by_group.values(): |
| for left in range(len(group)): |
| for right in range(len(group)): |
| if left == right: |
| continue |
| if _ranking_reward(group[left]) > _ranking_reward(group[right]) + reward_margin: |
| pairs.append( |
| ( |
| offset_by_record_id[group[left].record_id], |
| offset_by_record_id[group[right].record_id], |
| ) |
| ) |
| return pairs |
|
|
|
|
| def _ranking_reward(record: CILRecord) -> float: |
| return record.reward.score |
|
|
