test: pretokenized arrow dataset

Signed-off-by: Will Johnson <mwjohnson728@gmail.com>

Author: willmj

tests/artifacts/testdata/__init__.py

@@ -37,6 +37,9 @@
 TWITTER_COMPLAINTS_TOKENIZED_JSONL = os.path.join(
     DATA_DIR, "twitter_complaints_tokenized_with_maykeye_tinyllama_v0.jsonl"
 )
+TWITTER_COMPLAINTS_TOKENIZED_ARROW = os.path.join(
+    DATA_DIR, "twitter_complaints_tokenized_with_maykeye_tinyllama_v0.arrow"
+)
 EMPTY_DATA = os.path.join(DATA_DIR, "empty_data.json")
 MALFORMATTED_DATA = os.path.join(DATA_DIR, "malformatted_data.json")
 MODEL_NAME = "Maykeye/TinyLLama-v0"

tests/artifacts/testdata/twitter_complaints_tokenized_with_maykeye_tinyllama_v0.arrow

@@ -40,6 +40,7 @@
     TWITTER_COMPLAINTS_DATA_JSONL,
     TWITTER_COMPLAINTS_TOKENIZED_JSON,
     TWITTER_COMPLAINTS_TOKENIZED_JSONL,
+    TWITTER_COMPLAINTS_TOKENIZED_ARROW,
 )
 
 # Local
@@ -79,6 +80,20 @@
                 ]
             ),
         ),
+        (
+            TWITTER_COMPLAINTS_TOKENIZED_ARROW,
+            set(
+                [
+                    "Tweet text",
+                    "ID",
+                    "Label",
+                    "text_label",
+                    "output",
+                    "input_ids",
+                    "labels",
+                ]
+            ),
+        ),
         (
             TWITTER_COMPLAINTS_DATA_JSONL,
             set(["Tweet text", "ID", "Label", "text_label", "output"]),
@@ -519,6 +534,12 @@ def test_process_dataargs(data_args):
                 training_data_path=TWITTER_COMPLAINTS_TOKENIZED_JSONL,
             )
         ),
+        # JSONL pretokenized train datasets
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED_ARROW,
+            )
+        ),
     ],
 )
 def test_process_dataargs_pretokenized(data_args):