|
| 1 | +# Copyright 2023–2026 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""DPO specific input pipeline utilities.""" |
| 16 | + |
| 17 | +import dataclasses |
| 18 | +import grain.python as grain |
| 19 | +import numpy as np |
| 20 | + |
| 21 | + |
| 22 | +@dataclasses.dataclass |
| 23 | +class DPODataFormatting(grain.MapTransform): |
| 24 | + """Prepares DPO data. |
| 25 | + Renames input columns, extracts common prefix if needed, generates masks, and performs |
| 26 | + DPO-aware padding (left-padded prompts, right-padded responses). |
| 27 | + """ |
| 28 | + |
| 29 | + pad_id: int |
| 30 | + max_target_length: int |
| 31 | + data_column_names: tuple[str, ...] |
| 32 | + max_prompt_length: int | None = None |
| 33 | + |
| 34 | + def map(self, element): |
| 35 | + "Apply the dataset transformations for DPO." |
| 36 | + # 1. Reformat/Extract Columns |
| 37 | + try: |
| 38 | + if len(self.data_column_names) == 3: |
| 39 | + input_ids = element[self.data_column_names[0]] |
| 40 | + chosen_ids = element[self.data_column_names[1]] |
| 41 | + rejected_ids = element[self.data_column_names[2]] |
| 42 | + elif len(self.data_column_names) == 2: |
| 43 | + # Support for datasets like Anthropic/hh-rlhf where prompt is a common prefix |
| 44 | + full_chosen = element[self.data_column_names[0]] |
| 45 | + full_rejected = element[self.data_column_names[1]] |
| 46 | + |
| 47 | + # Find common prefix length |
| 48 | + prefix_len = 0 |
| 49 | + for c, r in zip(full_chosen, full_rejected): |
| 50 | + if c != r: |
| 51 | + break |
| 52 | + prefix_len += 1 |
| 53 | + input_ids = full_chosen[:prefix_len] |
| 54 | + chosen_ids = full_chosen[prefix_len:] |
| 55 | + rejected_ids = full_rejected[prefix_len:] |
| 56 | + else: |
| 57 | + raise ValueError(f"DPODataFormatting expects 2 or 3 columns, got {len(self.data_column_names)}") |
| 58 | + except KeyError as e: |
| 59 | + raise KeyError( |
| 60 | + f"Column '{e.args[0]}' not found in the dataset. " |
| 61 | + f"Expected columns: {self.data_column_names}. " |
| 62 | + f"Available columns: {list(element.keys())}. " |
| 63 | + "Please verify that 'train_data_columns' and 'eval_data_columns' match your dataset." |
| 64 | + ) from e |
| 65 | + |
| 66 | + # 2. Padding and Masking |
| 67 | + max_prompt_length = self.max_prompt_length or (self.max_target_length // 2) |
| 68 | + max_response_length = self.max_target_length - max_prompt_length |
| 69 | + |
| 70 | + assert max_prompt_length > 0, ( |
| 71 | + "max_prompt_length must be positive. " "Check the configs for 'max_prompt_length' and 'max_target_length'." |
| 72 | + ) |
| 73 | + assert max_response_length > 0, ( |
| 74 | + "max_response_length must be positive. " "Check the configs for 'max_prompt_length' and 'max_target_length'." |
| 75 | + ) |
| 76 | + |
| 77 | + prompt_ids = self._pad(input_ids, max_prompt_length, left=True) |
| 78 | + chosen_ids = self._pad(chosen_ids, max_response_length, left=False) |
| 79 | + rejected_ids = self._pad(rejected_ids, max_response_length, left=False) |
| 80 | + |
| 81 | + # Remove old columns if they exist |
| 82 | + for key in self.data_column_names: |
| 83 | + if key in element: |
| 84 | + del element[key] |
| 85 | + |
| 86 | + element["prompt_ids"] = prompt_ids |
| 87 | + element["chosen_ids"] = chosen_ids |
| 88 | + element["rejected_ids"] = rejected_ids |
| 89 | + element["prompt_mask"] = (prompt_ids != self.pad_id).astype(np.int32) |
| 90 | + element["chosen_mask"] = (chosen_ids != self.pad_id).astype(np.int32) |
| 91 | + element["rejected_mask"] = (rejected_ids != self.pad_id).astype(np.int32) |
| 92 | + return element |
| 93 | + |
| 94 | + def _pad(self, x, length, left=False): |
| 95 | + """Pads or trims an array to a specific length. |
| 96 | +
|
| 97 | + When left=True (for prompts), trims from the left to keep the suffix (closest context). |
| 98 | + When left=False (for responses), trims from the right to keep the prefix. |
| 99 | + """ |
| 100 | + x = np.asarray(x) |
| 101 | + pad_amount = max(length - x.shape[0], 0) |
| 102 | + if left: |
| 103 | + pad_width = ((pad_amount, 0),) |
| 104 | + x_trimmed = x[-length:] |
| 105 | + else: |
| 106 | + pad_width = ((0, pad_amount),) |
| 107 | + x_trimmed = x[:length] |
| 108 | + return np.pad(x_trimmed, pad_width, constant_values=self.pad_id).astype(np.int32) |
0 commit comments