change name

huggingface · Jun 26, 2024 · cf6271f · cf6271f
1 parent 7d45ef1
commit cf6271f
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 1 deletion.
diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py
@@ -104,6 +104,7 @@
         "DataCollatorForTokenClassification",
         "DataCollatorForWholeWordMask",
         "DataCollatorWithPadding",
+        "DataCollatorWithFlattening",
         "DefaultDataCollator",
         "default_data_collator",
     ],
@@ -4692,6 +4693,7 @@
         DataCollatorForTokenClassification,
         DataCollatorForWholeWordMask,
         DataCollatorWithPadding,
+        DataCollatorWithFlattening,
         DefaultDataCollator,
         default_data_collator,
     )

diff --git a/src/transformers/data/__init__.py b/src/transformers/data/__init__.py
@@ -20,6 +20,7 @@
     DataCollatorForTokenClassification,
     DataCollatorForWholeWordMask,
     DataCollatorWithPadding,
+    DataCollatorWithFlattening,
     DefaultDataCollator,
     default_data_collator,
 )

diff --git a/src/transformers/data/data_collator.py b/src/transformers/data/data_collator.py
@@ -1613,7 +1613,7 @@ def numpy_mask_tokens(self, inputs: Any) -> Tuple[Any, Any, Any, Any]:
         return inputs.astype(np.int64), perm_mask, target_mapping, labels.astype(np.int64)
 
 @dataclass
-class DataCollatorBatchFlattening(DefaultDataCollator):
+class DataCollatorWithFlattening(DefaultDataCollator):
     return_tensors: str = "pt"
     def __init__(self, return_position_ids=True):
         self.return_position_ids=return_position_ids