Upload training_code/model/dpo_data.py with huggingface_hub

training_code/model/dpo_data.py

@@ -0,0 +1,144 @@
+"""
+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),
+    }