fix: rewrite GRPO data pipeline for offline record-level access

- process_list_field returns List[List[int]] preserving per-response boundaries
- GRPODataset rewritten to record-level __getitem__ (no windowing/stride)
- grpo_collate_fn pads variable-length responses into [B, G, R] tensors
- JsonlStore detects nested List[List[int]] and stores List[Tensor] per record
- Store._normalize skips nested-list keys from cumsum bookkeeping
- Pipeline._flush handles nested lists without cross-record flattening
- Export grpo_collate_fn from astrai.dataset
- 6 new GRPO tests + 2 updated builder tests, 114 total pass
This commit is contained in:
2026-07-17 14:34:41 +08:00
parent c17aa0dc54
commit a1ea26d367
7 changed files with 502 additions and 75 deletions
+2
View File
@@ -1,6 +1,7 @@
from astrai.dataset.dataset import (
BaseDataset,
DatasetFactory,
grpo_collate_fn,
)
from astrai.dataset.sampler import ResumableDistributedSampler
from astrai.dataset.storage import (
@@ -21,6 +22,7 @@ from astrai.serialization import (
__all__ = [
"BaseDataset",
"DatasetFactory",
"grpo_collate_fn",
"Store",
"StoreFactory",
"H5Store",