Submodules#

oumi.core.collators.text_collator_with_padding module#

class oumi.core.collators.text_collator_with_padding.TextCollatorWithPadding(tokenizer: PreTrainedTokenizerBase, *, max_length: int | None, truncation: bool = False, label_ignore_index: int | None = None, max_variable_sized_dims: int = 1)[source]#

Bases: object

__call__(batch) → dict[str, Any][source]#

Pads to the longest length present in the batch.

Parameters:

batch – List of batch items.

Returns:

Processed batch.

Return type:

Dict[str, torch.Tensor]

oumi.core.collators.text_completions_collator_with_padding module#

class oumi.core.collators.text_completions_collator_with_padding.TextCompletionsCollatorWithPadding(tokenizer: PreTrainedTokenizerBase, instruction_prefix: str, response_prefix: str)[source]#

Bases: object

__call__(batch) → dict[str, Any][source]#

Pads to the longest length present in the batch.

Parameters:

batch – List of batch items.

Returns:

Processed batch.

Return type:

Dict[str, torch.Tensor]

oumi.core.collators.vision_language_collator_with_padding module#

class oumi.core.collators.vision_language_collator_with_padding.VisionLanguageCollatorWithPadding(tokenizer: PreTrainedTokenizerBase, *, max_length: int | None, truncation: bool = False, label_ignore_index: int | None = None, allow_multi_image_inputs: bool = True, main_image_feature: str = 'pixel_values')[source]#

Bases: object

__call__(batch) → dict[str, Any][source]#

Custom collator for multi-modal vision-language training.

Parameters:

batch – List of batch items.

Returns:

Processed batch.

Return type:

Dict[str, torch.Tensor]

collate_images(images) → Tensor[source]#

Collate images for multi-modal training.

Parameters:

images – List of images to collate.

Returns:

Batch of processed images.

Return type:

torch.Tensor

oumi.core.collators.vision_language_sft_collator module#

class oumi.core.collators.vision_language_sft_collator.VisionLanguageSftCollator(tokenizer: PreTrainedTokenizerBase, processor_name: str, *, processor_kwargs: dict[str, Any] | None = None, max_length: int | None = None, truncation: bool = False, truncation_side: str = 'right', label_ignore_index: int | None = None, allow_multi_image_inputs: bool = True, trust_remote_code: bool = False)[source]#

Bases: object

__call__(batch) → dict[str, Any][source]#

Custom collator for multi-modal vision-language training.

Parameters:

batch – List of batch items.

Returns:

Processed batch.

Return type:

Dict[str, torch.Tensor]