Upload run_summarization_flax.py

run_summarization_flax.py

@@ -431,23 +431,25 @@ def main():
         return
 
     # Get the column names for input/target.
-    dataset_columns = summarization_name_mapping.get(data_args.dataset_name, None)
-    if data_args.text_column is None:
-        text_column = dataset_columns[0] if dataset_columns is not None else column_names[0]
-    else:
-        text_column = data_args.text_column
-        if text_column not in column_names:
-            raise ValueError(
-                f"--text_column' value '{data_args.text_column}' needs to be one of: {', '.join(column_names)}"
-            )
-    if data_args.summary_column is None:
-        summary_column = dataset_columns[1] if dataset_columns is not None else column_names[1]
-    else:
-        summary_column = data_args.summary_column
-        if summary_column not in column_names:
-            raise ValueError(
-                f"--summary_column' value '{data_args.summary_column}' needs to be one of: {', '.join(column_names)}"
-            )
+    if not data_args.pretokenized:
+        
+        dataset_columns = summarization_name_mapping.get(data_args.dataset_name, None)
+        if data_args.text_column is None:
+            text_column = dataset_columns[0] if dataset_columns is not None else column_names[0]
+        else:
+            text_column = data_args.text_column
+            if text_column not in column_names:
+                raise ValueError(
+                    f"--text_column' value '{data_args.text_column}' needs to be one of: {', '.join(column_names)}"
+                )
+        if data_args.summary_column is None:
+            summary_column = dataset_columns[1] if dataset_columns is not None else column_names[1]
+        else:
+            summary_column = data_args.summary_column
+            if summary_column not in column_names:
+                raise ValueError(
+                    f"--summary_column' value '{data_args.summary_column}' needs to be one of: {', '.join(column_names)}"
+                )
 
     # Temporarily set max_target_length for training.
     max_target_length = data_args.max_target_length