wehub-resource-sync
80 lines
3.0 KiB
Python
80 lines
3.0 KiB
Python
"""
|
|
Batch iterator over preference pairs for reward-model and DPO training.
|
|
|
|
Reads a JSONL file of ``{"prompt", "chosen", "rejected"}`` (produced by
|
|
``scripts/prepare_preference_data.py``). Each side is rendered through the chat template
|
|
so we get, for the chosen and rejected responses to the same prompt:
|
|
- token ids of ``prompt + response + EOT``
|
|
- a response mask (1 over the completion, used by DPO)
|
|
- the true sequence length (used by the reward model to read the last-token reward)
|
|
|
|
Right-padding is safe here because the model's attention is causal: the last real token
|
|
never attends to padding that comes after it, and the response mask zeros padded
|
|
positions in the loss.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Iterator
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from src.post_training.chat_template import EOT_ID, encode_chat
|
|
|
|
|
|
def _encode_side(prompt: str, response: str, max_len: int) -> tuple[list[int], list[int]]:
|
|
ids, mask = encode_chat([{"role": "user", "content": prompt},
|
|
{"role": "assistant", "content": response}])
|
|
return ids[:max_len], mask[:max_len]
|
|
|
|
|
|
def _collate(rows: list[dict], max_len: int, device: str) -> dict:
|
|
enc = [( _encode_side(r["prompt"], r["chosen"], max_len),
|
|
_encode_side(r["prompt"], r["rejected"], max_len)) for r in rows]
|
|
# Pad chosen and rejected to a single common length so they can share one forward.
|
|
L = max(max(len(c[0]), len(j[0])) for c, j in enc)
|
|
|
|
def pad(seq, fill):
|
|
return seq + [fill] * (L - len(seq))
|
|
|
|
ch_ids, ch_mask, ch_len, rj_ids, rj_mask, rj_len = [], [], [], [], [], []
|
|
for (cids, cmask), (jids, jmask) in enc:
|
|
ch_len.append(len(cids)); rj_len.append(len(jids))
|
|
ch_ids.append(pad(cids, EOT_ID)); ch_mask.append(pad(cmask, 0))
|
|
rj_ids.append(pad(jids, EOT_ID)); rj_mask.append(pad(jmask, 0))
|
|
|
|
t = lambda a, dt: torch.tensor(a, dtype=dt, device=device)
|
|
return {
|
|
"chosen_ids": t(ch_ids, torch.long), "chosen_mask": t(ch_mask, torch.long), "chosen_len": t(ch_len, torch.long),
|
|
"rejected_ids": t(rj_ids, torch.long), "rejected_mask": t(rj_mask, torch.long), "rejected_len": t(rj_len, torch.long),
|
|
}
|
|
|
|
|
|
def get_preference_iterator(
|
|
path: str,
|
|
batch_size: int,
|
|
max_len: int,
|
|
device: str = "cpu",
|
|
*,
|
|
rank: int = 0,
|
|
world_size: int = 1,
|
|
shuffle: bool = True,
|
|
infinite: bool = True,
|
|
) -> Iterator[dict]:
|
|
"""Yield collated preference batches (dict of tensors). Rows are sharded across ranks."""
|
|
with open(path) as f:
|
|
rows = [json.loads(line) for line in f if line.strip()]
|
|
rows = rows[rank::world_size]
|
|
rng = np.random.default_rng(7 + rank)
|
|
while True:
|
|
order = np.arange(len(rows))
|
|
if shuffle:
|
|
rng.shuffle(order)
|
|
for s in range(0, len(order) - batch_size + 1, batch_size):
|
|
batch = [rows[i] for i in order[s:s + batch_size]]
|
|
yield _collate(batch, max_len, device)
|
|
if not infinite:
|
|
return
|