Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add datasets.Dataset to Trainer's train_dataset and eval_dataset type hints #30077

Merged
merged 3 commits into from
Apr 9, 2024
Merged

Add datasets.Dataset to Trainer's train_dataset and eval_dataset type hints #30077

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0ae8b2c
Add datasets.Dataset to Trainer's train_dataset and eval_dataset type…
ringohoffman Apr 6, 2024
1e9ab8d
Merge branch 'main' into add-datasets-dataset-to-trainer-dataset-type…
ringohoffman Apr 8, 2024
331012d
Add is_datasets_available check for importing datasets under TYPE_CHE…
ringohoffman Apr 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
10 changes: 6 additions & 4 deletions src/transformers/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,8 @@ def _get_fsdp_ckpt_kwargs():
if TYPE_CHECKING:
import optuna

if is_datasets_available():
import datasets

logger = logging.get_logger(__name__)

Expand Down Expand Up @@ -287,7 +289,7 @@ class Trainer:
The function to use to form a batch from a list of elements of `train_dataset` or `eval_dataset`. Will
default to [`default_data_collator`] if no `tokenizer` is provided, an instance of
[`DataCollatorWithPadding`] otherwise.
train_dataset (`torch.utils.data.Dataset` or `torch.utils.data.IterableDataset`, *optional*):
train_dataset (Union[`torch.utils.data.Dataset`, `torch.utils.data.IterableDataset`, `datasets.Dataset`], *optional*):
The dataset to use for training. If it is a [`~datasets.Dataset`], columns not accepted by the
`model.forward()` method are automatically removed.

Expand All @@ -296,7 +298,7 @@ class Trainer:
`torch.Generator` for the randomization that must be identical on all processes (and the Trainer will
manually set the seed of this `generator` at each epoch) or have a `set_epoch()` method that internally
sets the seed of the RNGs used.
eval_dataset (Union[`torch.utils.data.Dataset`, Dict[str, `torch.utils.data.Dataset`]), *optional*):
eval_dataset (Union[`torch.utils.data.Dataset`, Dict[str, `torch.utils.data.Dataset`, `datasets.Dataset`]), *optional*):
The dataset to use for evaluation. If it is a [`~datasets.Dataset`], columns not accepted by the
`model.forward()` method are automatically removed. If it is a dictionary, it will evaluate on each
dataset prepending the dictionary key to the metric name.
Expand Down Expand Up @@ -358,8 +360,8 @@ def __init__(
model: Union[PreTrainedModel, nn.Module] = None,
args: TrainingArguments = None,
data_collator: Optional[DataCollator] = None,
train_dataset: Optional[Union[Dataset, IterableDataset]] = None,
eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None,
ringohoffman marked this conversation as resolved.
Show resolved Hide resolved
eval_dataset: Optional[Union[Dataset, Dict[str, Dataset], "datasets.Dataset"]] = None,
tokenizer: Optional[PreTrainedTokenizerBase] = None,
model_init: Optional[Callable[[], PreTrainedModel]] = None,
compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
Expand Down