| from typing import Optional |
| """ |
| Alignment datasets and formatting utilities. |
| |
| Used for: |
| 1. Supervised Fine-Tuning, SFT |
| 2. Direct Preference Optimization, DPO |
| |
| Input formats: |
| |
| SFT JSONL: |
| {"instruction": "...", "response": "..."} |
| |
| DPO JSONL: |
| {"prompt": "...", "chosen": "...", "rejected": "..."} |
| |
| Output: |
| PyTorch datasets returning tokenized tensors. |
| """ |
|
|
| import json |
| from pathlib import Path |
| from typing import Dict, List |
|
|
| import torch |
| from torch.utils.data import Dataset |
| from tokenizers import Tokenizer |
|
|
|
|
| def read_jsonl(path: str) -> list[dict]: |
| """ |
| Input: |
| path: JSONL file path. |
| |
| Output: |
| List of dictionaries. |
| """ |
| rows = [] |
|
|
| with open(path, "r", encoding="utf-8") as f: |
| for line in f: |
| line = line.strip() |
|
|
| if not line: |
| continue |
|
|
| rows.append(json.loads(line)) |
|
|
| return rows |
|
|
|
|
| def format_instruction_prompt(instruction: str) -> str: |
| """ |
| Input: |
| instruction: User instruction text. |
| |
| Output: |
| Prompt string used during SFT/alignment. Delegates to the single shared |
| formatter so SFT, DPO, and inference never drift apart. |
| """ |
| from src.common.prompts import format_prompt |
| return format_prompt(instruction) |
|
|
|
|
| def format_response(response: str) -> str: |
| """ |
| Input: |
| response: Assistant response. |
| |
| Output: |
| Clean response text. |
| """ |
| return response.strip() |
|
|
|
|
| def build_sft_example( |
| tokenizer: Tokenizer, |
| instruction: str, |
| response: str, |
| bos_id: Optional[int], |
| eos_id: Optional[int], |
| block_size: int, |
| ) -> Optional[dict]: |
| """ |
| Builds one SFT training example. |
| |
| The model sees: |
| prompt + response |
| |
| Loss is calculated only on response tokens. |
| |
| Input: |
| tokenizer: custom BPE tokenizer. |
| instruction: input instruction. |
| response: expected answer. |
| bos_id: beginning-of-sequence token id. |
| eos_id: end-of-sequence token id. |
| block_size: max sequence length. |
| |
| Output: |
| dict with input_ids and labels, or None if invalid. |
| """ |
| prompt = format_instruction_prompt(instruction) |
| response = format_response(response) |
|
|
| prompt_ids = tokenizer.encode(prompt).ids |
| response_ids = tokenizer.encode(response).ids |
|
|
| if bos_id is not None: |
| prompt_ids = [bos_id] + prompt_ids |
|
|
| if eos_id is not None: |
| response_ids = response_ids + [eos_id] |
|
|
| input_ids = prompt_ids + response_ids |
|
|
| |
| labels = [-100] * len(prompt_ids) + response_ids |
|
|
| if len(input_ids) > block_size: |
| input_ids = input_ids[:block_size] |
| labels = labels[:block_size] |
|
|
| |
| if all(x == -100 for x in labels): |
| return None |
|
|
| return { |
| "input_ids": input_ids, |
| "labels": labels, |
| } |
|
|
|
|
| def build_dpo_pair( |
| tokenizer: Tokenizer, |
| prompt: str, |
| chosen: str, |
| rejected: str, |
| bos_id: Optional[int], |
| eos_id: Optional[int], |
| block_size: int, |
| ) -> Optional[dict]: |
| """ |
| Builds one DPO preference pair. |
| |
| DPO compares: |
| log probability of chosen answer |
| vs. |
| log probability of rejected answer |
| |
| Loss is calculated only on answer tokens. |
| |
| Input: |
| tokenizer: custom BPE tokenizer. |
| prompt: user prompt. |
| chosen: preferred response. |
| rejected: non-preferred response. |
| bos_id/eos_id: optional special token ids. |
| block_size: maximum sequence length. |
| |
| Output: |
| dict containing chosen/rejected input_ids and labels. |
| """ |
| prompt_text = format_instruction_prompt(prompt) |
|
|
| prompt_ids = tokenizer.encode(prompt_text).ids |
|
|
| if bos_id is not None: |
| prompt_ids = [bos_id] + prompt_ids |
|
|
| chosen_ids = tokenizer.encode(format_response(chosen)).ids |
| rejected_ids = tokenizer.encode(format_response(rejected)).ids |
|
|
| if eos_id is not None: |
| chosen_ids = chosen_ids + [eos_id] |
| rejected_ids = rejected_ids + [eos_id] |
|
|
| chosen_input = prompt_ids + chosen_ids |
| rejected_input = prompt_ids + rejected_ids |
|
|
| chosen_labels = [-100] * len(prompt_ids) + chosen_ids |
| rejected_labels = [-100] * len(prompt_ids) + rejected_ids |
|
|
| if len(chosen_input) > block_size: |
| chosen_input = chosen_input[:block_size] |
| chosen_labels = chosen_labels[:block_size] |
|
|
| if len(rejected_input) > block_size: |
| rejected_input = rejected_input[:block_size] |
| rejected_labels = rejected_labels[:block_size] |
|
|
| if all(x == -100 for x in chosen_labels): |
| return None |
|
|
| if all(x == -100 for x in rejected_labels): |
| return None |
|
|
| return { |
| "chosen_input_ids": chosen_input, |
| "chosen_labels": chosen_labels, |
| "rejected_input_ids": rejected_input, |
| "rejected_labels": rejected_labels, |
| } |
|
|
|
|
| class SFTDataset(Dataset): |
| """ |
| Supervised fine-tuning dataset. |
| |
| Input: |
| JSONL file with instruction-response rows. |
| |
| Output per item: |
| input_ids |
| labels |
| """ |
|
|
| def __init__( |
| self, |
| path: str, |
| tokenizer: Tokenizer, |
| bos_id: Optional[int], |
| eos_id: Optional[int], |
| block_size: int, |
| ): |
| rows = read_jsonl(path) |
| examples = [] |
|
|
| for row in rows: |
| example = build_sft_example( |
| tokenizer=tokenizer, |
| instruction=row.get("prompt", row.get("instruction", "")), |
| response=row["response"], |
| bos_id=bos_id, |
| eos_id=eos_id, |
| block_size=block_size, |
| ) |
|
|
| if example is not None: |
| examples.append(example) |
|
|
| if not examples: |
| raise ValueError("No valid SFT examples found.") |
|
|
| self.examples = examples |
|
|
| def __len__(self): |
| return len(self.examples) |
|
|
| def __getitem__(self, idx): |
| return self.examples[idx] |
|
|
|
|
| class DPODataset(Dataset): |
| """ |
| DPO preference dataset. |
| |
| Input: |
| JSONL file with prompt/chosen/rejected rows. |
| |
| Output per item: |
| chosen_input_ids |
| chosen_labels |
| rejected_input_ids |
| rejected_labels |
| """ |
|
|
| def __init__( |
| self, |
| path: str, |
| tokenizer: Tokenizer, |
| bos_id: Optional[int], |
| eos_id: Optional[int], |
| block_size: int, |
| ): |
| rows = read_jsonl(path) |
| examples = [] |
|
|
| for row in rows: |
| example = build_dpo_pair( |
| tokenizer=tokenizer, |
| prompt=row["prompt"], |
| chosen=row["chosen"], |
| rejected=row["rejected"], |
| bos_id=bos_id, |
| eos_id=eos_id, |
| block_size=block_size, |
| ) |
|
|
| if example is not None: |
| examples.append(example) |
|
|
| if not examples: |
| raise ValueError("No valid DPO examples found.") |
|
|
| self.examples = examples |
|
|
| def __len__(self): |
| return len(self.examples) |
|
|
| def __getitem__(self, idx): |
| return self.examples[idx] |
|
|
|
|
| def pad_1d(values: List[int], max_len: int, pad_value: int) -> List[int]: |
| """ |
| Pads one integer list to max_len. |
| """ |
| return values + [pad_value] * (max_len - len(values)) |
|
|
|
|
| def collate_sft(batch: list[dict], pad_id: int) -> dict: |
| """ |
| Input: |
| batch of SFT examples. |
| |
| Output: |
| padded tensors. |
| """ |
| max_len = max(len(x["input_ids"]) for x in batch) |
|
|
| input_ids = [] |
| labels = [] |
|
|
| for item in batch: |
| input_ids.append(pad_1d(item["input_ids"], max_len, pad_id)) |
| labels.append(pad_1d(item["labels"], max_len, -100)) |
|
|
| return { |
| "input_ids": torch.tensor(input_ids, dtype=torch.long), |
| "labels": torch.tensor(labels, dtype=torch.long), |
| } |
|
|
|
|
| def collate_dpo(batch: list[dict], pad_id: int) -> dict: |
| """ |
| Input: |
| batch of DPO examples. |
| |
| Output: |
| padded chosen/rejected tensors. |
| """ |
| chosen_max = max(len(x["chosen_input_ids"]) for x in batch) |
| rejected_max = max(len(x["rejected_input_ids"]) for x in batch) |
|
|
| chosen_input_ids = [] |
| chosen_labels = [] |
| rejected_input_ids = [] |
| rejected_labels = [] |
|
|
| for item in batch: |
| chosen_input_ids.append( |
| pad_1d(item["chosen_input_ids"], chosen_max, pad_id) |
| ) |
| chosen_labels.append( |
| pad_1d(item["chosen_labels"], chosen_max, -100) |
| ) |
| rejected_input_ids.append( |
| pad_1d(item["rejected_input_ids"], rejected_max, pad_id) |
| ) |
| rejected_labels.append( |
| pad_1d(item["rejected_labels"], rejected_max, -100) |
| ) |
|
|
| return { |
| "chosen_input_ids": torch.tensor(chosen_input_ids, dtype=torch.long), |
| "chosen_labels": torch.tensor(chosen_labels, dtype=torch.long), |
| "rejected_input_ids": torch.tensor(rejected_input_ids, dtype=torch.long), |
| "rejected_labels": torch.tensor(rejected_labels, dtype=torch.long), |
| } |
|
|
