145 lines
5.2 KiB
Python
145 lines
5.2 KiB
Python
"""
|
|
DPO data pipeline: loads UltraFeedback preference pairs.
|
|
|
|
Each example has a prompt + chosen response + rejected response.
|
|
We tokenize both (prompt+chosen) and (prompt+rejected), apply the same
|
|
chat template, and return them as pairs for DPO training.
|
|
"""
|
|
|
|
import torch
|
|
from torch.utils.data import Dataset, DataLoader
|
|
from datasets import load_dataset
|
|
|
|
|
|
CHAT_TEMPLATE = {
|
|
"user_start": "<|user|>\n",
|
|
"assistant_start": "<|assistant|>\n",
|
|
"turn_end": "\n<|end|>\n",
|
|
}
|
|
|
|
|
|
def format_preference_pair(prompt, chosen_msgs, rejected_msgs):
|
|
"""Build chat-templated strings for chosen and rejected."""
|
|
def build(messages):
|
|
text = CHAT_TEMPLATE["user_start"] + prompt.strip() + CHAT_TEMPLATE["turn_end"]
|
|
for msg in messages:
|
|
role = msg.get("role", "assistant")
|
|
content = msg.get("content", "").strip()
|
|
if role == "assistant":
|
|
text += CHAT_TEMPLATE["assistant_start"] + content + CHAT_TEMPLATE["turn_end"]
|
|
elif role == "user":
|
|
text += CHAT_TEMPLATE["user_start"] + content + CHAT_TEMPLATE["turn_end"]
|
|
return text
|
|
|
|
return build(chosen_msgs), build(rejected_msgs)
|
|
|
|
|
|
class DPODataset(Dataset):
|
|
"""
|
|
Loads UltraFeedback preference pairs and tokenizes them.
|
|
Returns (prompt_ids, chosen_ids, rejected_ids) with proper shifting.
|
|
"""
|
|
|
|
def __init__(self, tokenizer, max_seq_len=2048, split="train",
|
|
cache_dir=None, max_samples=None):
|
|
self.tokenizer = tokenizer
|
|
self.max_seq_len = max_seq_len
|
|
|
|
special_tokens = ["<|user|>", "<|assistant|>", "<|end|>"]
|
|
vocab = tokenizer.get_vocab()
|
|
new_tokens = [t for t in special_tokens if t not in vocab]
|
|
if new_tokens:
|
|
tokenizer.add_tokens(new_tokens, special_tokens=True)
|
|
|
|
self.assistant_token_id = tokenizer.encode("<|assistant|>", add_special_tokens=False)[0]
|
|
self.end_token_id = tokenizer.encode("<|end|>", add_special_tokens=False)[0]
|
|
self.user_token_id = tokenizer.encode("<|user|>", add_special_tokens=False)[0]
|
|
|
|
print(f"[DPO Data] Loading UltraFeedback preferences ({split})...")
|
|
ds = load_dataset(
|
|
"argilla/ultrafeedback-binarized-preferences-cleaned",
|
|
split=split,
|
|
cache_dir=cache_dir,
|
|
)
|
|
if max_samples:
|
|
ds = ds.select(range(min(max_samples, len(ds))))
|
|
print(f"[DPO Data] {len(ds)} preference pairs loaded")
|
|
|
|
self.examples = []
|
|
skipped = 0
|
|
for i, row in enumerate(ds):
|
|
prompt = row.get("prompt", "")
|
|
chosen = row.get("chosen", [])
|
|
rejected = row.get("rejected", [])
|
|
|
|
if not prompt or not chosen or not rejected:
|
|
skipped += 1
|
|
continue
|
|
|
|
chosen_text, rejected_text = format_preference_pair(prompt, chosen, rejected)
|
|
|
|
chosen_ids = tokenizer.encode(chosen_text, add_special_tokens=False)
|
|
rejected_ids = tokenizer.encode(rejected_text, add_special_tokens=False)
|
|
|
|
# Truncate if needed
|
|
if len(chosen_ids) > max_seq_len + 1:
|
|
chosen_ids = chosen_ids[:max_seq_len + 1]
|
|
if len(rejected_ids) > max_seq_len + 1:
|
|
rejected_ids = rejected_ids[:max_seq_len + 1]
|
|
|
|
if len(chosen_ids) < 10 or len(rejected_ids) < 10:
|
|
skipped += 1
|
|
continue
|
|
|
|
# Find where the prompt ends (first <|assistant|> token)
|
|
prompt_end = 0
|
|
for j, tid in enumerate(chosen_ids):
|
|
if tid == self.assistant_token_id:
|
|
prompt_end = j + 2 # skip <|assistant|> and \n
|
|
break
|
|
|
|
self.examples.append({
|
|
"chosen_ids": chosen_ids,
|
|
"rejected_ids": rejected_ids,
|
|
"prompt_len": prompt_end,
|
|
})
|
|
|
|
if (i + 1) % 20000 == 0:
|
|
print(f" Processed {i+1} pairs...")
|
|
|
|
print(f"[DPO Data] {len(self.examples)} pairs ready, {skipped} skipped")
|
|
|
|
def __len__(self):
|
|
return len(self.examples)
|
|
|
|
def __getitem__(self, idx):
|
|
ex = self.examples[idx]
|
|
return {
|
|
"chosen_ids": torch.tensor(ex["chosen_ids"], dtype=torch.long),
|
|
"rejected_ids": torch.tensor(ex["rejected_ids"], dtype=torch.long),
|
|
"prompt_len": ex["prompt_len"],
|
|
}
|
|
|
|
|
|
def dpo_collate_fn(batch, pad_id=0):
|
|
"""Pad chosen and rejected sequences separately."""
|
|
max_chosen = max(b["chosen_ids"].size(0) for b in batch)
|
|
max_rejected = max(b["rejected_ids"].size(0) for b in batch)
|
|
|
|
chosen_padded = []
|
|
rejected_padded = []
|
|
prompt_lens = []
|
|
|
|
for b in batch:
|
|
c_pad = max_chosen - b["chosen_ids"].size(0)
|
|
r_pad = max_rejected - b["rejected_ids"].size(0)
|
|
chosen_padded.append(torch.cat([b["chosen_ids"], torch.full((c_pad,), pad_id, dtype=torch.long)]))
|
|
rejected_padded.append(torch.cat([b["rejected_ids"], torch.full((r_pad,), pad_id, dtype=torch.long)]))
|
|
prompt_lens.append(b["prompt_len"])
|
|
|
|
return {
|
|
"chosen_ids": torch.stack(chosen_padded),
|
|
"rejected_ids": torch.stack(rejected_padded),
|
|
"prompt_lens": torch.tensor(prompt_lens, dtype=torch.long),
|
|
}
|