Skip to content

Commit

Permalink
Use typing.Optional and Union for Python 3.9
Browse files Browse the repository at this point in the history
  • Loading branch information
tomaarsen committed Jan 17, 2025
1 parent 907c5b8 commit 8f3cdc4
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions sentence_transformers/training_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -170,18 +171,18 @@ 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. "
"Either 1) a single string prompt, 2) a mapping of column names to prompts, 3) a mapping of dataset names "
"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."}
)

Expand Down

0 comments on commit 8f3cdc4

Please sign in to comment.