Fix trainer_seq2seq.py's __init__ type annotations (#34021)

* Fix `trainer_seq2seq.py`'s `__init__` type annotations * Update src/transformers/trainer_seq2seq.py Co-authored-by: Lysandre Debut <[email protected]> * Fix issue pointed out by `muellerzr` --------- Co-authored-by: Lysandre Debut <[email protected]>
huggingface · Oct 8, 2024 · 573942d · 573942d
1 parent 04b4e44
commit 573942d
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/src/transformers/trainer_seq2seq.py b/src/transformers/trainer_seq2seq.py
@@ -24,11 +24,16 @@
 from .generation.configuration_utils import GenerationConfig
 from .integrations.deepspeed import is_deepspeed_zero3_enabled
 from .trainer import Trainer
-from .utils import logging
+from .utils import is_datasets_available, logging
 from .utils.deprecation import deprecate_kwarg
 
 
+if is_datasets_available():
+    import datasets
+
 if TYPE_CHECKING:
+    from torch.utils.data import IterableDataset
+
     from .data.data_collator import DataCollator
     from .feature_extraction_utils import FeatureExtractionMixin
     from .image_processing_utils import BaseImageProcessor
@@ -50,7 +55,7 @@ def __init__(
         model: Union["PreTrainedModel", nn.Module] = None,
         args: "TrainingArguments" = None,
         data_collator: Optional["DataCollator"] = None,
-        train_dataset: Optional[Dataset] = None,
+        train_dataset: Optional[Union[Dataset, "IterableDataset", "datasets.Dataset"]] = None,
         eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
         processing_class: Optional[
             Union["PreTrainedTokenizerBase", "BaseImageProcessor", "FeatureExtractionMixin", "ProcessorMixin"]