Spaces:
Paused
Paused
File size: 4,827 Bytes
2f86d9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | # Copyright 2026 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Functional version of the `gm.data` transforms."""
from etils import enp
import flax
import jax
from kauldron.ktyping import Array, Bool, Int, PyTree # pylint: disable=g-multiple-import,g-importing-member
import numpy as np
# Do not @typechecked as `element` can be `list` too.
def pad(
element: PyTree[Array["sequence"]], # pyrefly: ignore[not-a-type, unknown-name]
max_length: int,
*,
truncate: bool = False,
fill_value: int = 0,
axis: int = -1,
) -> PyTree[Array["max_length"]]: # pyrefly: ignore[not-a-type, unknown-name]
"""Add zeros to the end of the sequence to reach the max length.
Supports padding multiple arrays at once.
Args:
element: The sequence to pad.
max_length: The max length of the sequence.
truncate: Whether to truncate the sequence to the max length. If `False`,
sequences longer than the `max_length` will raise an error.
fill_value: The value to fill the padded sequence with.
axis: The axis in which to pad the sequence (only -1 supported at the
moment).
Returns:
The padded sequence of length `max_length`.
"""
if axis != -1:
raise NotImplementedError("Only `axis=-1` is supported.")
return jax.tree.map(
lambda x: _pad(
x,
max_length=max_length,
fill_value=fill_value,
truncate=truncate,
),
element,
is_leaf=_is_list_array, # Also supports `[0, 1, ...]`
)
def _pad(
element: Array["sequence"], # pyrefly: ignore[unknown-name]
*,
max_length: int,
fill_value: int,
truncate: bool = False,
) -> Array["max_length"]: # pyrefly: ignore[unknown-name]
"""Inner padding implementation."""
# Use `xnp` so it supports both `np` and `jnp`.
xnp = enp.lazy.get_xnp(element, strict=False)
element = xnp.asarray(element)
# TODO(epot): Could add an `axis=` kwarg to support multi-dimensional arrays.
seq_length = element.shape[-1]
if not truncate and seq_length > max_length:
raise ValueError(
f"Cannot pad sequence of length {seq_length}. Is longer than the"
f" max length {max_length}. Set `truncate=True`."
)
sentence_tokens = element[..., :max_length]
pad_width = [(0, 0)] * (sentence_tokens.ndim - 1) + [
(0, max_length - sentence_tokens.shape[-1])
]
return xnp.pad(sentence_tokens, pad_width, constant_values=fill_value)
@flax.struct.dataclass
class Seq2SeqFields:
"""Fields for next token prediction."""
input: Int["*b l"]
target: Int["*b l"]
target_mask: Bool["*b l"]
# Note: There's no `batch` dimension here. It wouldn't make much sense as
# each example has a different length, so batching can only be applied
# after the output is padded.
def make_seq2seq_fields(
prompt: Int["prompt_len"], # pyrefly: ignore[unknown-name]
response: Int["response_len"], # pyrefly: ignore[unknown-name]
) -> Seq2SeqFields:
"""Create the model `input`, `target` and `loss_mask`.
From prompt and response token ids, generate the model `input`, `target` and
`loss_mask`.
Example:
```python
# Input:
prompt = [10, 11, 12, 13],
response = [20, 21, 1], # Here, response ends with EOS token.
# Ouptut:
out.input = [10, 11, 12, 13, 20, 21],
out.target = [11, 12, 13, 20, 21, 1],
out.target_mask = [ 0, 0, 0, 1, 1, 1],
```
Note:
- Input and target are the same sequence shifted by one token.
- The last token from the target is truncated from the input (as there's no
target for it)
Args:
prompt: The prompt tokens.
response: The response tokens.
Returns:
The input, target and mask, all of length `prompt_len + response_len - 1`.
"""
# Concatenate the prompt and response tokens.
sequence = np.concatenate([prompt, response])
# Create the loss mask.
target_mask = np.concatenate([
np.zeros((len(prompt) - 1,), dtype=np.bool_),
np.ones((len(response),), dtype=np.bool_),
])
return Seq2SeqFields(
input=sequence[:-1],
target=sequence[1:],
target_mask=target_mask,
)
def _is_list_array(x) -> bool:
"""Returns `True` if `x` is a list of ints, like `[0, 1, ...]`."""
return isinstance(x, list | tuple) and all(isinstance(x, int) for x in x)
|