diff --git a/sentence_transformers/training_args.py b/sentence_transformers/training_args.py index 0e91e52f1..da9a65dc6 100644 --- a/sentence_transformers/training_args.py +++ b/sentence_transformers/training_args.py @@ -2,6 +2,7 @@ import logging from dataclasses import dataclass, field +from typing import Optional, Union from transformers import TrainingArguments as TransformersTrainingArguments from transformers.training_args import ParallelMode @@ -170,7 +171,7 @@ class SentenceTransformerTrainingArguments(TransformersTrainingArguments): for valid options. Defaults to ``MultiDatasetBatchSamplers.PROPORTIONAL``. """ - prompts: str | None = field( + prompts: Optional[str] = field( # noqa: UP007 default=None, metadata={ "help": "The prompts to use for each column in the datasets. " @@ -178,10 +179,10 @@ class SentenceTransformerTrainingArguments(TransformersTrainingArguments): "to prompts, or 4) a mapping of dataset names to a mapping of column names to prompts." }, ) - batch_sampler: BatchSamplers | str = field( + batch_sampler: Union[BatchSamplers, str] = field( # noqa: UP007 default=BatchSamplers.BATCH_SAMPLER, metadata={"help": "The batch sampler to use."} ) - multi_dataset_batch_sampler: MultiDatasetBatchSamplers | str = field( + multi_dataset_batch_sampler: Union[MultiDatasetBatchSamplers, str] = field( # noqa: UP007 default=MultiDatasetBatchSamplers.PROPORTIONAL, metadata={"help": "The multi-dataset batch sampler to use."} )