wehub-resource-sync
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
"""
|
|
Prompt iterator for RL (PPO / GRPO). Reads a JSONL of ``{"prompt", "gold"}`` (gold is the
|
|
verifiable numeric answer, or null) and yields lists of sampled rows. Rows are sharded
|
|
across ranks so each GPU optimizes on its own prompts.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Iterator
|
|
|
|
import numpy as np
|
|
|
|
|
|
def load_prompt_rows(path: str) -> list[dict]:
|
|
with open(path) as f:
|
|
return [json.loads(line) for line in f if line.strip()]
|
|
|
|
|
|
def get_prompt_iterator(
|
|
path: str,
|
|
prompts_per_iter: int,
|
|
*,
|
|
rank: int = 0,
|
|
world_size: int = 1,
|
|
seed: int = 0,
|
|
shuffle: bool = True,
|
|
) -> Iterator[list[dict]]:
|
|
"""Infinitely yield lists of ``prompts_per_iter`` rows (this rank's shard)."""
|
|
rows = load_prompt_rows(path)[rank::world_size]
|
|
rng = np.random.default_rng(seed + rank)
|
|
n = len(rows)
|
|
while True:
|
|
idx = rng.permutation(n) if shuffle else np.arange(n)
|
|
for s in range(0, n - prompts_per_iter + 1, prompts_per_iter):
|
|
yield [rows[i] for i in idx[s:s + prompts_per_iter]]
|