PR Changes

Signed-off-by: Abhishek <[email protected]>

Author: Abhishek-TAMU

.pylintrc

@@ -333,7 +333,7 @@ indent-string='    '
 max-line-length=100
 
 # Maximum number of lines in a module.
-max-module-lines=1400
+max-module-lines=1200
 
 # Allow the body of a class to be on the same line as the declaration if body
 # contains single statement.

tests/data/test_data_preprocessing_utils.py

@@ -23,7 +23,9 @@
 import datasets
 import pytest
 import yaml
-
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
 # First Party
 from tests.artifacts.predefined_data_configs import (
     DATA_CONFIG_APPLY_CUSTOM_TEMPLATE_YAML,
@@ -498,7 +500,7 @@ def test_process_dataconfig_file(data_config_path, data_path):
 
 
 @pytest.mark.parametrize(
-    "data_config_path, list_data_path",
+    "data_config_path, data_path_list",
     [
         (
             DATA_CONFIG_APPLY_CUSTOM_TEMPLATE_YAML,
@@ -571,11 +573,11 @@ def test_process_dataconfig_file(data_config_path, data_path):
         ),
     ],
 )
-def test_process_dataconfig_multiple_files(data_config_path, list_data_path):
+def test_process_dataconfig_multiple_files(data_config_path, data_path_list):
     """Ensure that datasets with multiple files are formatted and validated correctly based on the arguments passed in config file."""
     with open(data_config_path, "r") as f:
         yaml_content = yaml.safe_load(f)
-    yaml_content["datasets"][0]["data_paths"] = list_data_path
+    yaml_content["datasets"][0]["data_paths"] = data_path_list
     datasets_name = yaml_content["datasets"][0]["name"]
 
     # Modify input_field_name and output_field_name according to dataset
@@ -635,7 +637,7 @@ def test_process_dataconfig_multiple_files(data_config_path, list_data_path):
         ),
     ],
 )
-def test_process_dataconfig_multiple_datasets_datafiles(datafiles, datasetconfigname):
+def test_process_dataconfig_multiple_datasets_datafiles_sampling(datafiles, datasetconfigname):
     """Ensure that multiple datasets with multiple files are formatted and validated correctly."""
     with open(datasetconfigname, "r") as f:
         yaml_content = yaml.safe_load(f)
@@ -651,14 +653,26 @@ def test_process_dataconfig_multiple_datasets_datafiles(datafiles, datasetconfig
         data_args = configs.DataArguments(data_config_path=temp_yaml_file_path)
 
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    (train_set, _, _) = _process_dataconfig_file(data_args, tokenizer)
+    TRAIN_ARGS = configs.TrainingArguments(
+        packing=False,
+        max_seq_length=1024,
+        output_dir="tmp",
+    )
+    (train_set, eval_set, _, _, _, _) = process_dataargs(
+        data_args=data_args, tokenizer=tokenizer, train_args=TRAIN_ARGS
+    )
+
     assert isinstance(train_set, Dataset)
-    column_names = set(["input_ids", "attention_mask", "labels"])
-    assert set(train_set.column_names) == column_names
+    if eval_set:
+        assert isinstance(eval_set, Dataset)
+
+    assert set(["input_ids", "attention_mask", "labels"]).issubset(set(train_set.column_names))
+    if eval_set:
+        assert set(["input_ids", "attention_mask", "labels"]).issubset(set(eval_set.column_names))
 
 
 @pytest.mark.parametrize(
-    "data_config_path, list_data_path",
+    "data_config_path, data_path_list",
     [
         (
             DATA_CONFIG_APPLY_CUSTOM_TEMPLATE_YAML,
@@ -682,12 +696,12 @@ def test_process_dataconfig_multiple_datasets_datafiles(datafiles, datasetconfig
     ],
 )
 def test_process_dataconfig_multiple_files_varied_data_formats(
-    data_config_path, list_data_path
+    data_config_path, data_path_list
 ):
     """Ensure that datasets with multiple files with different formats raise assertion error when passed in config file."""
     with open(data_config_path, "r") as f:
         yaml_content = yaml.safe_load(f)
-    yaml_content["datasets"][0]["data_paths"] = list_data_path
+    yaml_content["datasets"][0]["data_paths"] = data_path_list
     datasets_name = yaml_content["datasets"][0]["name"]
 
     # Modify input_field_name and output_field_name according to dataset
@@ -719,7 +733,7 @@ def test_process_dataconfig_multiple_files_varied_data_formats(
 
 
 @pytest.mark.parametrize(
-    "data_config_path, list_data_path",
+    "data_config_path, data_path_list",
     [
         (
             DATA_CONFIG_APPLY_CUSTOM_TEMPLATE_YAML,
@@ -739,12 +753,12 @@ def test_process_dataconfig_multiple_files_varied_data_formats(
     ],
 )
 def test_process_dataconfig_multiple_files_varied_types(
-    data_config_path, list_data_path
+    data_config_path, data_path_list
 ):
     """Ensure that datasets with multiple files with different formats raise assertion error when passed in config file."""
     with open(data_config_path, "r") as f:
         yaml_content = yaml.safe_load(f)
-    yaml_content["datasets"][0]["data_paths"] = list_data_path
+    yaml_content["datasets"][0]["data_paths"] = data_path_list
     datasets_name = yaml_content["datasets"][0]["name"]
 
     # Modify input_field_name and output_field_name according to dataset
@@ -1048,51 +1062,3 @@ def test_process_dataset_configs_with_sampling_error(
         (_, _, _, _, _, _) = process_dataargs(
             data_args=data_args, tokenizer=tokenizer, train_args=TRAIN_ARGS
         )
-
-
-@pytest.mark.parametrize(
-    "datafiles, datasetconfigname",
-    [
-        (
-            [
-                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_ARROW,
-                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSONL,
-                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_PARQUET,
-            ],
-            DATA_CONFIG_MULTIPLE_DATASETS_SAMPLING_YAML,
-        ),
-    ],
-)
-def test_process_dataset_configs_with_sampling(datafiles, datasetconfigname):
-
-    data_args = configs.DataArguments()
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    TRAIN_ARGS = configs.TrainingArguments(
-        packing=False,
-        max_seq_length=1024,
-        output_dir="tmp",  # Not needed but positional
-    )
-
-    with tempfile.NamedTemporaryFile(
-        "w", delete=False, suffix=".yaml"
-    ) as temp_yaml_file:
-        with open(datasetconfigname, "r") as f:
-            data = yaml.safe_load(f)
-            datasets = data["datasets"]
-            for i in range(len(datasets)):
-                d = datasets[i]
-                d["data_paths"][0] = datafiles[i]
-            yaml.dump(data, temp_yaml_file)
-        data_args.data_config_path = temp_yaml_file.name
-
-    (train_set, eval_set, _, _, _, _) = process_dataargs(
-        data_args=data_args, tokenizer=tokenizer, train_args=TRAIN_ARGS
-    )
-
-    assert isinstance(train_set, Dataset)
-    if eval_set:
-        assert isinstance(eval_set, Dataset)
-
-    assert set(["input_ids", "labels"]).issubset(set(train_set.column_names))
-    if eval_set:
-        assert set(["input_ids", "labels"]).issubset(set(eval_set.column_names))