PR changes

Signed-off-by: Abhishek <[email protected]>
foundation-model-stack · Dec 11, 2024 · 4ba1c04 · 4ba1c04
1 parent e89002d
commit 4ba1c04
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 3 deletions.
diff --git a/.pylintrc b/.pylintrc
@@ -333,7 +333,7 @@ indent-string='    '
 max-line-length=100
 
 # Maximum number of lines in a module.
-max-module-lines=1200
+max-module-lines=1400
 
 # Allow the body of a class to be on the same line as the declaration if body
 # contains single statement.

diff --git a/tests/data/test_data_preprocessing_utils.py b/tests/data/test_data_preprocessing_utils.py
@@ -506,7 +506,11 @@ def test_process_dataconfig_file(data_config_path, data_path):
         ),
         (
             DATA_CONFIG_APPLY_CUSTOM_TEMPLATE_YAML,
-            [TWITTER_COMPLAINTS_DATA_JSONL, TWITTER_COMPLAINTS_DATA_JSONL],
+            [
+                TWITTER_COMPLAINTS_DATA_JSONL,
+                TWITTER_COMPLAINTS_DATA_JSONL,
+                TWITTER_COMPLAINTS_DATA_JSONL,
+            ],
         ),
         (
             DATA_CONFIG_APPLY_CUSTOM_TEMPLATE_YAML,
@@ -529,6 +533,7 @@ def test_process_dataconfig_file(data_config_path, data_path):
             [
                 TWITTER_COMPLAINTS_TOKENIZED_PARQUET,
                 TWITTER_COMPLAINTS_TOKENIZED_PARQUET,
+                TWITTER_COMPLAINTS_TOKENIZED_PARQUET,
             ],
         ),
         (
@@ -561,6 +566,7 @@ def test_process_dataconfig_file(data_config_path, data_path):
             [
                 TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_ARROW,
                 TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_ARROW,
+                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_ARROW,
             ],
         ),
     ],

diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -40,7 +40,10 @@
     MALFORMATTED_DATA,
     MODEL_NAME,
     TWITTER_COMPLAINTS_DATA_ARROW,
+    TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_ARROW,
+    TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSON,
     TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSONL,
+    TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_PARQUET,
     TWITTER_COMPLAINTS_DATA_JSON,
     TWITTER_COMPLAINTS_DATA_JSONL,
     TWITTER_COMPLAINTS_DATA_PARQUET,
@@ -772,13 +775,34 @@ def test_run_causallm_ft_pretokenized(dataset_path):
 @pytest.mark.parametrize(
     "datafiles, datasetconfigname",
     [
+        (
+            [
+                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSON,
+                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSON,
+            ],
+            DATA_CONFIG_MULTIPLE_DATASETS_SAMPLING_YAML,
+        ),
         (
             [
                 TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSONL,
                 TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSONL,
             ],
             DATA_CONFIG_MULTIPLE_DATASETS_SAMPLING_YAML,
-        )
+        ),
+        (
+            [
+                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_ARROW,
+                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_ARROW,
+            ],
+            DATA_CONFIG_MULTIPLE_DATASETS_SAMPLING_YAML,
+        ),
+        (
+            [
+                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_PARQUET,
+                TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_PARQUET,
+            ],
+            DATA_CONFIG_MULTIPLE_DATASETS_SAMPLING_YAML,
+        ),
     ],
 )
 def test_run_causallm_ft_and_inference_with_multiple_dataset(