# Copyright 2023 The EASYDEL Author @erfanzar (Erfan Zare Chavoshi).
#
# 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
#
# https://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.
import logging
import random
import re
import typing as tp
import warnings
from contextlib import contextmanager
from functools import partial
from eformer.pytree import auto_pytree
import chex
import jax
import numpy as np
from jax import numpy as jnp
from ml_collections import ConfigDict
from ml_collections.config_dict import placeholder
from easydel.infra.utils import ProcessingClassType
from easydel.utils import traversals
from easydel.utils.helpers import get_logger
logger = get_logger(__name__)
[docs]class JaxDistributedConfig(object):
"""
From EasyLM
Utility class for initializing JAX distributed.
"""
[docs] @staticmethod
def get_default_config(updates=None):
config = ConfigDict()
config.initialize_jax_distributed = False
config.coordinator_address = placeholder(str)
config.num_processes = placeholder(int)
config.process_id = placeholder(int)
config.local_device_ids = placeholder(str)
if updates is not None:
config.update(ConfigDict(updates).copy_and_resolve_references())
return config
[docs] @classmethod
def initialize(cls, config=None):
config = cls.get_default_config(config)
if config.initialize_jax_distributed:
if config.local_device_ids is not None:
local_device_ids = [int(x) for x in config.local_device_ids.split(",")]
else:
local_device_ids = None
jax.distributed.initialize(
coordinator_address=config.coordinator_address,
num_processes=config.num_processes,
process_id=config.process_id,
local_device_ids=local_device_ids,
)
# fmt:off
[docs]def create_prompt_creator(processing_class):
def to_role_and_content(field):
return {
"conversation": [
{"role": "user", "content": field["conversation"][0]["input"]},
{"role": "assistant", "content": field["conversation"][0]["output"]},
]
}
def _pc(sample):
return conversations_formatting_function(processing_class, messages_field="conversation")(to_role_and_content(sample))
return _pc
# fmt:on
[docs]def create_constant_length_dataset(
processing_class,
dataset,
dataset_text_field: tp.Optional[str] = None,
formatting_func: tp.Optional[tp.Callable] = None,
infinite: bool = False,
seq_length: int = 1024,
num_of_sequences: int = 1024,
chars_per_token: float = 3.6,
eos_token_id: int = 0,
shuffle: bool = True,
append_concat_token: bool = True,
add_special_tokens: bool = True,
) -> tp.Callable[[], tp.Iterator[tp.Dict[str, jnp.ndarray]]]:
"""
Creates a generator function that yields constant length chunks of tokens from a stream of text files.
Args:
processing_class: The processor used for processing the data.
dataset: Dataset with text files.
dataset_text_field: Name of the field in the dataset that contains the text.
formatting_func: Function that formats the text before tokenization.
infinite: If True the iterator is reset after dataset reaches end else stops.
seq_length: Length of token sequences to return.
num_of_sequences: Number of token sequences to keep in buffer.
chars_per_token: Number of characters per token used to estimate number of tokens in text buffer.
eos_token_id: Id of the end of sequence token if the passed processing_class does not have an EOS token.
shuffle: Shuffle the examples before they are returned.
append_concat_token: If true, appends eos_token_id at the end of each sample being packed.
add_special_tokens: If true, processing_class adds special tokens to each sample being packed.
Returns:
A generator function that yields dictionaries containing input_ids and attention_mask as jnp.arrays
"""
if processing_class.eos_token_id is None:
warnings.warn(
"The passed processing_class does not have an EOS token. We will use the passed eos_token_id instead which "
f"corresponds to {eos_token_id}. If this is not the correct EOS token, make sure to pass the correct eos_token_id.",
stacklevel=1,
)
concat_token_id = (
processing_class.eos_token_id if processing_class.eos_token_id else eos_token_id
)
max_buffer_size = seq_length * chars_per_token * num_of_sequences
# Input validation and formatting function setup
if dataset_text_field is not None and formatting_func is not None:
warnings.warn(
"Only one of `dataset_text_field` and `formatting_func` should be provided. "
"Ignoring `dataset_text_field` and using `formatting_func`.",
stacklevel=1,
)
if formatting_func is not None:
if formatting_func.__code__.co_argcount > 1:
warnings.warn(
"The passed formatting_func has more than one argument. Usually that function should have a single argument "
"`example` which corresponds to the dictionary returned by each element of the dataset. Make sure you know "
"what you are doing.",
stacklevel=1,
)
elif dataset_text_field is not None:
formatting_func = lambda x: x[dataset_text_field] # noqa
else:
raise ValueError(
"Either `dataset_text_field` or `formatting_func` should be provided."
)
def constant_length_generator() -> tp.Iterator[tp.Dict[str, jnp.ndarray]]:
iterator = iter(dataset)
more_examples = True
while more_examples:
buffer, buffer_len = [], 0
# Fill the buffer
while True:
if buffer_len >= max_buffer_size:
break
try:
prompt = formatting_func(next(iterator))
if isinstance(prompt, list):
prompt = "".join(p for p in prompt)
buffer.append(prompt)
buffer_len += len(buffer[-1])
except StopIteration:
if infinite:
iterator = iter(dataset)
warnings.warn(
"The dataset reached end and the iterator is reset to the start.",
stacklevel=1,
)
else:
more_examples = False
break
if shuffle:
random.shuffle(buffer)
# Tokenize all texts in the buffer
tokens = processing_class(
text=buffer,
add_special_tokens=add_special_tokens,
truncation=False,
)
tokenized_inputs = tokens["input_ids"]
attention_masks = tokens["attention_mask"]
# Concatenate all tokens and attention masks
all_token_ids = []
all_attention_masks = []
for tokenized_input, attention_mask in zip(tokenized_inputs, attention_masks):
if append_concat_token:
tokenized_input = tokenized_input + [concat_token_id]
attention_mask = attention_mask + [1]
all_token_ids.extend(tokenized_input)
all_attention_masks.extend(attention_mask)
# Create fixed-length examples
examples = []
examples_attention_masks = []
for i in range(0, len(all_token_ids), seq_length):
input_ids = all_token_ids[i : i + seq_length]
org_attention_masks = all_attention_masks[i : i + seq_length]
if len(input_ids) == seq_length:
examples.append(input_ids)
examples_attention_masks.append(org_attention_masks)
if shuffle:
# Shuffle examples while keeping pairs together
combined = list(zip(examples, examples_attention_masks))
random.shuffle(combined)
examples, examples_attention_masks = zip(*combined)
# Yield examples
for example, example_attention_mask in zip(examples, examples_attention_masks):
yield {
"input_ids": jnp.asarray(example, dtype="i4"),
"attention_mask": jnp.asarray(example_attention_mask, dtype="i4"),
}
return constant_length_generator
def _collate_batch(
examples,
processing_class,
pad_to_multiple_of: tp.Optional[int] = None,
):
if isinstance(examples[0], (list, tuple)):
examples = [jnp.array(e, dtype=jnp.int64) for e in examples]
length_of_first = len(examples[0])
are_tensors_same_length = all(len(x) == length_of_first for x in examples)
if are_tensors_same_length and (
pad_to_multiple_of is None or length_of_first % pad_to_multiple_of == 0
):
return jnp.stack(examples, axis=0)
if processing_class._pad_token is None:
raise ValueError(
"You are attempting to pad samples but the processing_class you are using"
f" ({processing_class.__class__.__name__}) does not have a pad token."
)
max_length = max(len(x) for x in examples)
if pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
result = jnp.full(
shape=(len(examples), max_length),
fill_value=processing_class.pad_token_id,
dtype=examples[0].dtype,
)
for i, example in enumerate(examples):
if processing_class.padding_side == "right":
result[i, : example.shape[0]] = example
else:
result[i, -example.shape[0] :] = example
return result
[docs]def tolist(x):
"""from HF
Args:
x:
Returns: X as tp.List
"""
if isinstance(x, list):
return x
elif hasattr(x, "numpy"):
x = x.numpy()
return x.tolist()
[docs]class DataCollatorForCompletionOnlyLM:
"""Data collator used for completion tasks. It ensures that all the tokens of the labels are set to an 'ignore_index'
when they do not come from the assistant. This ensures that the loss is only
calculated on the completion made by the assistant.
"""
def __init__(
self,
processing_class: tp.Union[str, "PreTrainedTokenizerBase"], # type:ignore #noqa
response_template: tp.Union[str, tp.List[int]],
instruction_template: tp.Optional[tp.Union[str, tp.List[int]]] = None,
*args,
mlm: bool = False,
ignore_index: int = -100,
**kwargs,
):
from transformers import AutoTokenizer
if isinstance(processing_class, str):
processing_class = AutoTokenizer.from_pretrained(processing_class)
self.processing_class = processing_class
self.instruction_template = instruction_template
if isinstance(instruction_template, str):
self.instruction_token_ids = self.processing_class.encode(
self.instruction_template, add_special_tokens=False
)
else:
self.instruction_token_ids = instruction_template
self.response_template = response_template
if isinstance(response_template, str):
self.response_token_ids = self.processing_class.encode(
self.response_template, add_special_tokens=False
)
else:
self.response_token_ids = response_template
if (
not mlm
and self.instruction_template
and self.processing_class.pad_token_id == self.processing_class.eos_token_id
):
warnings.warn(
"The pad_token_id and eos_token_id values of this processing_class are identical. "
"If you are planning for multi-turn training, "
"it can result in the model continuously generating questions and answers without eos token. "
"To avoid this, set the pad_token_id to a different value.",
stacklevel=1,
)
self.ignore_index = ignore_index
def _whole_word_mask(self, input_tokens: tp.List[str], max_predictions=512):
from transformers import BertTokenizer, BertTokenizerFast
if not isinstance(self.processing_class, (BertTokenizer, BertTokenizerFast)):
warnings.warn(
"DataCollatorForWholeWordMask is only suitable for BertTokenizer-like tokenizers. "
"Please refer to the documentation for more information.",
stacklevel=1,
)
cand_indexes = []
for i, token in enumerate(input_tokens):
if token == "[CLS]" or token == "[SEP]":
continue
if len(cand_indexes) >= 1 and token.startswith("##"):
cand_indexes[-1].append(i)
else:
cand_indexes.append([i])
random.shuffle(cand_indexes)
num_to_predict = min(max_predictions, max(1, int(round(len(input_tokens) * 0.15))))
masked_lms = []
covered_indexes = set()
for index_set in cand_indexes:
if len(masked_lms) >= num_to_predict:
break
if len(masked_lms) + len(index_set) > num_to_predict:
continue
is_any_index_covered = False
for index in index_set:
if index in covered_indexes:
is_any_index_covered = True
break
if is_any_index_covered:
continue
for index in index_set:
covered_indexes.add(index)
masked_lms.append(index)
if len(covered_indexes) != len(masked_lms):
raise ValueError(
"Length of covered_indexes is not equal to length of masked_lms."
)
mask_labels = [1 if i in covered_indexes else 0 for i in range(len(input_tokens))]
return mask_labels
[docs] def jax_mask_tokens(
self, inputs: tp.Any, special_tokens_mask: tp.Optional[tp.Any] = None
) -> tp.Tuple[tp.Any, tp.Any]:
"""Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original."""
labels = np.copy(inputs)
probability_matrix = np.full(labels.shape, 0.15)
if special_tokens_mask is None:
special_tokens_mask = [
self.processing_class.get_special_tokens_mask(
val, already_has_special_tokens=True
)
for val in labels.tolist()
]
special_tokens_mask = np.array(special_tokens_mask, dtype=bool)
else:
special_tokens_mask = special_tokens_mask.astype(bool)
probability_matrix[special_tokens_mask] = 0
masked_indices = np.random.binomial(
1, probability_matrix, size=probability_matrix.shape
).astype(bool)
labels[~masked_indices] = -100
indices_replaced = (
np.random.binomial(1, 0.8, size=labels.shape).astype(bool) & masked_indices
)
inputs[indices_replaced] = self.processing_class.mask_token_id
indices_random = (
np.random.binomial(1, 0.5, size=labels.shape).astype(bool)
& masked_indices
& ~indices_replaced
)
random_words = np.random.randint(
low=0,
high=len(self.processing_class),
size=np.count_nonzero(indices_random),
dtype=np.int64,
)
inputs[indices_random] = random_words
return inputs, labels
[docs] def jax_call(
self, examples: tp.List[tp.Union[tp.List[int], tp.Any, tp.Dict[str, tp.Any]]]
) -> tp.Dict[str, tp.Any]:
if isinstance(examples[0], tp.Mapping):
input_ids = [e["input_ids"] for e in examples]
else:
input_ids = examples
examples = [{"input_ids": e} for e in examples]
batch_input = _collate_batch(
input_ids,
self.processing_class,
)
mask_labels = []
for e in examples:
ref_tokens = []
for ida in tolist(e["input_ids"]):
token = self.processing_class._convert_id_to_token(ida)
ref_tokens.append(token)
# For Chinese tokens, we need extra inf to mark sub-word, e.g [喜,欢]-> [喜，##欢]
if "chinese_ref" in e:
ref_pos = tolist(e["chinese_ref"])
len_seq = len(e["input_ids"])
for i in range(len_seq):
if i in ref_pos:
ref_tokens[i] = "##" + ref_tokens[i]
mask_labels.append(self._whole_word_mask(ref_tokens))
batch_mask = _collate_batch(
mask_labels,
self.processing_class,
)
inputs, labels = self.jax_mask_tokens(batch_input, batch_mask)
return {"input_ids": inputs, "labels": labels}
def __call__(
self, examples: tp.List[tp.Union[tp.List[int], tp.Any, tp.Dict[str, tp.Any]]]
) -> tp.Dict[str, tp.Any]:
batch = self.jax_call(examples)
if self.instruction_template is None:
for i in range(len(examples)):
response_token_ids_start_idx = None
for idx in jnp.where(batch["labels"][i] == self.response_token_ids[0])[0]:
if (
self.response_token_ids
== batch["labels"][i][idx : idx + len(self.response_token_ids)].tolist()
):
response_token_ids_start_idx = idx
if response_token_ids_start_idx is None:
warnings.warn(
f"Could not find response key `{self.response_template}` in the "
f"following instance: {self.processing_class.decode(batch['input_ids'][i])} "
f"This instance will be ignored in loss calculation. "
f"Note, if this happens often, consider increasing the `max_seq_length`.",
stacklevel=1,
)
batch["labels"][i, :] = self.ignore_index
else:
response_token_ids_end_idx = response_token_ids_start_idx + len(
self.response_token_ids
)
batch["labels"][i, :response_token_ids_end_idx] = self.ignore_index
else:
for i in range(len(examples)):
response_token_ids_idxs = []
human_token_ids_idxs = []
for assistant_idx in jnp.where(
batch["labels"][i] == self.response_token_ids[0]
)[0]:
if (
self.response_token_ids
== batch["labels"][i][
assistant_idx : assistant_idx + len(self.response_token_ids)
].tolist()
):
response_token_ids_idxs.append(assistant_idx + len(self.response_token_ids))
if len(response_token_ids_idxs) == 0:
warnings.warn(
f"Could not find response key `{self.response_template}` in the "
f"following instance: {self.processing_class.decode(batch['input_ids'][i])} "
f"This instance will be ignored in loss calculation. "
f"Note, if this happens often, consider increasing the `max_seq_length`.",
stacklevel=1,
)
batch["labels"][i, :] = self.ignore_index
human_token_ids = self.instruction_token_ids
for human_idx in jnp.where(batch["labels"][i] == human_token_ids[0])[0]:
if (
human_token_ids
== batch["labels"][i][human_idx : human_idx + len(human_token_ids)].tolist()
):
human_token_ids_idxs.append(human_idx)
if len(human_token_ids_idxs) == 0:
warnings.warn(
f"Could not find instruction key `{self.instruction_template}` in the "
f"following instance: {self.processing_class.decode(batch['input_ids'][i])} "
f"This instance will be ignored in loss calculation. "
f"Note, if this happens often, consider increasing the `max_seq_length`.",
stacklevel=1,
)
batch["labels"][i, :] = self.ignore_index
if (
len(human_token_ids_idxs) > 0
and len(response_token_ids_idxs) > 0
and human_token_ids_idxs[0] > response_token_ids_idxs[0]
):
human_token_ids_idxs = [0] + human_token_ids_idxs
for idx, (start, end) in enumerate(
zip(human_token_ids_idxs, response_token_ids_idxs)
):
if idx != 0:
batch["labels"][i, start:end] = self.ignore_index
else:
batch["labels"][i, :end] = self.ignore_index
if len(response_token_ids_idxs) < len(human_token_ids_idxs):
batch["labels"][i, human_token_ids_idxs[-1] :] = self.ignore_index
return batch
[docs]@auto_pytree
class RewardDataCollatorWithPadding:
r"""
Reward DataCollator class that pads the inputs to the maximum length of the batch.
Args:
tokenizer (`ProcessingClassType`):
The tokenizer used for encoding the data.
padding (`Union[bool, str, `PaddingStrategy`]`, `optional`, defaults to `True`):
padding_strategy to pass to the tokenizer.
max_length (`int` or `None`, `optional`, defaults to `None`):
If set will pad the sequence to a maximum provided value.
"""
tokenizer: ProcessingClassType
padding: tp.Union[bool, str] = "max_length"
max_length: tp.Optional[int] = None
truncation_mode: str = "keep_end"
def __call__(self, features: list[dict[str, tp.Any]]) -> dict[str, tp.Any]:
features_chosen = []
features_rejected = []
margin = []
has_margin = "margin" in features[0]
for feature in features:
if (
"input_ids_chosen" not in feature
or "input_ids_rejected" not in feature
or "attention_mask_chosen" not in feature
or "attention_mask_rejected" not in feature
):
raise ValueError(
"The features should include `input_ids_chosen`, `attention_mask_chosen`, `input_ids_rejected` and `attention_mask_rejected`"
)
features_chosen.append(
{
"input_ids": feature["input_ids_chosen"],
"attention_mask": feature["attention_mask_chosen"],
}
)
features_rejected.append(
{
"input_ids": feature["input_ids_rejected"],
"attention_mask": feature["attention_mask_rejected"],
}
)
if has_margin:
margin.append(feature["margin"])
batch_chosen = self.tokenizer.pad(
features_chosen,
padding=self.padding,
max_length=self.max_length,
return_tensors="jax",
)
batch_rejected = self.tokenizer.pad(
features_rejected,
padding=self.padding,
max_length=self.max_length,
return_tensors="jax",
)
batch = {
"input_ids_chosen": batch_chosen["input_ids"],
"attention_mask_chosen": batch_chosen["attention_mask"],
"input_ids_rejected": batch_rejected["input_ids"],
"attention_mask_rejected": batch_rejected["attention_mask"],
}
if has_margin:
margin = jnp.array(margin, dtype="f4")
batch["margin"] = margin
return batch
[docs]@auto_pytree
class DataCollatorForPreference:
r"""DPO DataCollator class that pads the tokenized inputs to the maximum length of the batch."""
max_prompt_length: int
max_completion_length: int
pad_token_id: int = 0
label_pad_token_id: int = -100
is_encoder_decoder: tp.Optional[bool] = False
def __call__(self, features: tp.List[tp.Dict[str, tp.Any]]) -> tp.Dict[str, tp.Any]:
prompt_input_ids = [jnp.array(feature["prompt_input_ids"]) for feature in features]
prompt_attention_mask = [jnp.ones_like(input_ids) for input_ids in prompt_input_ids]
chosen_input_ids = [jnp.array(feature["chosen_input_ids"]) for feature in features]
chosen_attention_mask = [jnp.ones_like(input_ids) for input_ids in chosen_input_ids]
rejected_input_ids = [
jnp.array(feature["rejected_input_ids"]) for feature in features
]
rejected_attention_mask = [
jnp.ones_like(input_ids) for input_ids in rejected_input_ids
]
pixel_values = None
pixel_attention_mask = None
if "pixel_values" in features[0]:
pixel_values = [jnp.array(feature["pixel_values"]) for feature in features]
if "pixel_attention_mask" in features[0]:
pixel_attention_mask = [
jnp.array(feature["pixel_attention_mask"]) for feature in features
]
ref_chosen_logps = None
ref_rejected_logps = None
if "ref_chosen_logps" in features[0] and "ref_rejected_logps" in features[0]:
ref_chosen_logps = jnp.array(
[feature["ref_chosen_logps"] for feature in features]
)
ref_rejected_logps = jnp.array(
[feature["ref_rejected_logps"] for feature in features]
)
# Pad sequences
output = {
"prompt_input_ids": pad(
prompt_input_ids,
self.max_prompt_length,
padding_value=self.pad_token_id,
padding_side="left",
),
"prompt_attention_mask": pad(
prompt_attention_mask,
self.max_prompt_length,
padding_value=0,
padding_side="left",
),
"chosen_input_ids": pad(
chosen_input_ids,
self.max_completion_length,
padding_value=self.pad_token_id,
),
"chosen_attention_mask": pad(
chosen_attention_mask,
self.max_completion_length,
padding_value=0,
),
"rejected_input_ids": pad(
rejected_input_ids,
self.max_completion_length,
padding_value=self.pad_token_id,
),
"rejected_attention_mask": pad(
rejected_attention_mask,
self.max_completion_length,
padding_value=0,
),
}
# Add optional outputs
if pixel_values is not None:
output["pixel_values"] = pad(
pixel_values,
self.max_prompt_length,
padding_value=0.0,
)
if pixel_attention_mask is not None:
output["pixel_attention_mask"] = pad(
pixel_attention_mask,
self.max_prompt_length,
padding_value=0,
)
if "image_sizes" in features[0]:
output["image_sizes"] = jnp.array(
[feature["image_sizes"] for feature in features]
)
if ref_chosen_logps is not None and ref_rejected_logps is not None:
output["ref_chosen_logps"] = ref_chosen_logps
output["ref_rejected_logps"] = ref_rejected_logps
return output
[docs]@auto_pytree
class DPODataCollatorWithPadding:
r"""
DPO DataCollator class that pads the tokenized inputs to the maximum length of the batch.
"""
max_prompt_length: int
max_completion_length: int
pad_token_id: int = 0
label_pad_token_id: int = -100
is_encoder_decoder: tp.Optional[bool] = False
output_arrays_only: bool = True
prepadded: bool = True
def __call__(self, features: list[dict[str, tp.Any]]) -> dict[str, tp.Any]:
camax_length = self.max_completion_length + self.max_prompt_length
padded_batch = {}
for k in features[0].keys():
if k.endswith(("_input_ids", "_attention_mask", "_labels", "_pixel_values")):
match k.split("_")[0]:
case "rejected":
max_length = self.max_completion_length
case "chosen":
max_length = self.max_completion_length
case "prompt":
max_length = self.max_prompt_length
case _:
max_length = camax_length
if self.is_encoder_decoder:
to_pad = [jnp.array(ex[k], dtype="i4") for ex in features]
if (k.startswith("prompt")) and (k.endswith("input_ids")):
if self.pad_token_id is None:
raise ValueError(
"Padding is enabled, but the tokenizer is not configured with a padding token."
" Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)"
" before calling the trainer."
)
padding_value = self.pad_token_id
elif k.endswith("_attention_mask"):
padding_value = 0
elif k.startswith(("chosen", "rejected", "completion")) or ("decoder" in k):
padding_value = self.label_pad_token_id
else:
raise ValueError(f"Unexpected key in batch '{k}'")
padded_batch[k] = pad_sequence(
to_pad,
batch_first=False,
padding_value=padding_value,
max_len=None if self.prepadded else max_length,
)
else:
if k.endswith("_input_ids"):
if self.pad_token_id is None:
raise ValueError(
"Padding is enabled, but the tokenizer is not configured with a padding token."
" Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)"
" before calling the trainer."
)
padding_value = self.pad_token_id
elif k.endswith("_labels"):
padding_value = self.label_pad_token_id
elif k.endswith("_attention_mask"):
padding_value = 0
elif k.endswith("_pixel_values"):
padding_value = 0
else:
raise ValueError(f"Unexpected key in batch '{k}'")
if k in ["prompt_input_ids", "prompt_attention_mask"]:
padding_side = "left"
else:
padding_side = "right"
if k.endswith("_pixel_values"):
dtype = jnp.float32
else:
dtype = jnp.int32
to_pad = [jnp.array(ex[k], dtype=dtype) for ex in features]
padded_batch[k] = pad(
to_pad,
None if self.prepadded else max_length,
padding_value=padding_value,
padding_side=padding_side,
)
elif k.endswith("_logps"):
padded_batch[k] = jnp.array([ex[k] for ex in features])
else:
padded_batch[k] = [ex[k] for ex in features]
if self.output_arrays_only:
val = padded_batch.get(k)
if hasattr(val, "dtype"):
if val.dtype not in [
jnp.float64,
jnp.float32,
jnp.float16,
jnp.int32,
jnp.int16,
jnp.int8,
]:
padded_batch.pop(k)
else:
padded_batch.pop(k)
return padded_batch
[docs]def shift_and_pad(mask, *tensors):
for i in range(mask.shape[0]):
first_one_idx = np.nonzero(mask[i])[0][0].item()
mask[i] = np.roll(mask[i], shift=-first_one_idx)
for tensor in tensors:
tensor[i] = np.roll(tensor[i], shift=-first_one_idx)
if not tensors:
return mask
else:
return mask, *tensors
[docs]def pad(
tensors: list[jnp.ndarray],
max_lenght: tp.Optional[int],
padding_value: int = 0,
padding_side: str = "right",
) -> jnp.ndarray:
"""
Pads a list of tensors to the same shape along the first dimension.
"""
output_shape = tensors[0].shape[:-1]
current_max = tensors[0].shape[-1]
if max_lenght is None:
max_lenght = current_max
x_lenght = max(current_max, max_lenght)
output_shape += (x_lenght,)
output = jnp.full(
(len(tensors), *output_shape),
padding_value,
dtype=tensors[0].dtype,
)
for i, t in enumerate(tensors):
if padding_side == "left":
seq_slice = slice(output_shape[0] - t.shape[0], output_shape[0])
elif padding_side == "right":
seq_slice = slice(0, t.shape[0])
else:
raise ValueError("padding_side must be 'left' or 'right'")
slices = (i,) + (seq_slice,) + tuple(slice(0, s) for s in t.shape[1:])
output = output.at[slices].set(t)
if padding_side == "left":
output = output[..., -max_lenght:]
elif padding_side == "right":
output = output[..., :max_lenght]
else:
raise ValueError("padding_side must be 'left' or 'right'")
return output
[docs]def pad_to_length(
tensor: chex.Array,
length: int,
pad_value: tp.Union[int, float],
axis: int = -1,
) -> chex.Array:
if tensor.shape[axis] >= length:
if tensor.ndim == 2:
tensor = tensor[:, :length]
return tensor
else:
pad_size = list(tensor.shape)
pad_size[axis] = length - tensor.shape[axis]
return jax.numpy.concatenate(
[
tensor,
pad_value * jax.numpy.ones(pad_size, dtype=tensor.dtype),
],
axis=axis,
)
[docs]def pad_sequence(
sequences,
batch_first=False,
padding_value=0,
max_len: int | None = None,
):
max_len = max(seq.shape[-1] for seq in sequences) if max_len is None else max_len
padding_value = jnp.array(padding_value).reshape(1)
if batch_first:
padded_seqs = [
(
jnp.concatenate(
[
seq.reshape(1, -1),
jnp.ones((1, max_len - seq.shape[-1])) * padding_value,
],
axis=1,
)
if seq.shape[-1] < max_len
else seq.reshape(1, -1)
)
for seq in sequences
]
else:
padded_seqs = [
(
jnp.concatenate(
[
jnp.ones((1, max_len - seq.shape[-1])) * padding_value,
seq.reshape(1, -1),
],
axis=1,
)
if seq.shape[-1] < max_len
else seq.reshape(1, -1)
)
for seq in sequences
]
return jnp.array(padded_seqs)
[docs]@contextmanager
def leave_alone_context_manager():
# Perform setup actions (none in this case)
yield
[docs]def add_bos_token_if_needed(
bos_token_id: tp.Optional[int],
prompt_len_input_ids: int,
prompt_tokens: tp.Dict[str, tp.List[int]],
chosen_prompt_len_input_ids: int,
chosen_tokens: tp.Dict[str, tp.List[int]],
rejected_prompt_len_input_ids: int,
rejected_tokens: tp.Dict[str, tp.List[int]],
):
if bos_token_id is not None:
if (
prompt_len_input_ids == 0 or bos_token_id != prompt_tokens["prompt_input_ids"][0]
):
prompt_tokens["prompt_input_ids"] = [bos_token_id] + prompt_tokens[
"prompt_input_ids"
]
prompt_tokens["prompt_attention_mask"] = [1] + prompt_tokens[
"prompt_attention_mask"
]
if (
chosen_prompt_len_input_ids == 0
or bos_token_id != chosen_tokens["prompt_input_ids"][0]
):
chosen_tokens["prompt_input_ids"] = [bos_token_id] + chosen_tokens[
"prompt_input_ids"
]
chosen_tokens["prompt_attention_mask"] = [1] + chosen_tokens[
"prompt_attention_mask"
]
if (
rejected_prompt_len_input_ids == 0
or bos_token_id != rejected_tokens["prompt_input_ids"][0]
):
rejected_tokens["prompt_input_ids"] = [bos_token_id] + rejected_tokens[
"prompt_input_ids"
]
rejected_tokens["prompt_attention_mask"] = [1] + rejected_tokens[
"prompt_attention_mask"
]
return prompt_tokens, chosen_tokens, rejected_tokens
[docs]def add_eos_token_if_needed(
eos_token_id: int,
chosen_tokens: tp.Dict[str, tp.List[int]],
rejected_tokens: tp.Dict[str, tp.List[int]],
):
if (
len(chosen_tokens["input_ids"]) == 0
or eos_token_id != chosen_tokens["input_ids"][-1]
):
chosen_tokens["input_ids"].append(eos_token_id)
chosen_tokens["attention_mask"].append(1)
if (
len(rejected_tokens["input_ids"]) == 0
or eos_token_id != rejected_tokens["input_ids"][-1]
):
rejected_tokens["input_ids"].append(eos_token_id)
rejected_tokens["attention_mask"].append(1)
return chosen_tokens, rejected_tokens
[docs]def first_true_indices(bools, dtype=jnp.int32):
"""
Takes an N-dimensional bool array and returns an (N-1)-dimensional array of integers giving
the position of the first True in each "row".
"""
row_len = bools.shape[-1]
zero_or_index = row_len * (~bools).astype(dtype) + jnp.arange(row_len, dtype=dtype)
return jnp.min(zero_or_index, axis=-1)
[docs]def truncate_right(input_ids, stop_token_id, pad_token_id):
"""
Truncates the input array from the right side after the first occurrence of the stop token.
"""
trunc_idxs = first_true_indices(input_ids == stop_token_id).reshape((-1, 1))
idxs = jnp.arange(input_ids.shape[1]).reshape((1, -1))
output_ids = jnp.where(idxs > trunc_idxs, pad_token_id, input_ids)
mask = jnp.where(idxs > trunc_idxs, 0, 1)
return output_ids, mask
[docs]@partial(jax.jit, static_argnums=(1,))
def compute_weight_stats(params, repattern: str):
"""Compute statistics for model weights in a JIT-compatible way.
Args:
params: Model parameters
repattern: parameters to analyze
Returns:
Dictionary of weight statistics
"""
stats = {}
for path, weight in traversals.flatten_dict(params).items():
weight = weight.value
pattern_search = ".".join([str(p) for p in path])
path = "/".join([str(p) for p in path])
if bool(re.match(repattern, pattern_search)):
stats[f"{path}/values"] = weight.flatten()
stats[f"{path}/mean"] = jnp.mean(weight)
stats[f"{path}/std"] = jnp.std(weight)
stats[f"{path}/min"] = jnp.min(weight)
stats[f"{path}/max"] = jnp.max(weight)
return stats