Repository: IBM/unitxt

Issue:
Bug in loading datasets with partial splits.
This fails:

from unitxt.loaders import LoadHF
```
# Initialize the LoadHF loader
loader = LoadHF(path="data/rag/loan_agreement",data_files = {
            "test": "all_entity_dataset.jsonl"
        })

# Load the dataset
dataset = loader.process()

# Print the loaded dataset
print(dataset.to_dataset())
```

There error
```
file "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/arrow_reader.py", line 225, in get_file_instructions
    file_instructions = make_file_instructions(
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/arrow_reader.py", line 134, in make_file_instructions
    absolute_instructions = instruction.to_absolute(name2len)
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/arrow_reader.py", line 663, in to_absolute
    return [_rel_to_abs_instr(rel_instr, name2len) for rel_instr in self._relative_instructions]
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/arrow_reader.py", line 663, in <listcomp>
    return [_rel_to_abs_instr(rel_instr, name2len) for rel_instr in self._relative_instructions]
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/arrow_reader.py", line 480, in _rel_to_abs_instr
    raise ValueError(f'Unknown split "{split}". Should be one of {list(name2len)}.')
ValueError: Unknown split "train". Should be one of ['test'].

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/yoavkatz/migrate/v1/unitxt/src/unitxt/stream.py", line 251, in to_dataset
    {
  File "/Users/yoavkatz/migrate/v1/unitxt/src/unitxt/stream.py", line 252, in <dictcomp>
    key: value.to_dataset(
  File "/Users/yoavkatz/migrate/v1/unitxt/src/unitxt/stream.py", line 40, in to_dataset
    return Dataset.from_generator(
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 1136, in from_generator
    ).read()
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/io/generator.py", line 49, in read
    self.builder.download_and_prepare(
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/builder.py", line 1029, in download_and_prepare
    self._download_and_prepare(
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/builder.py", line 1791, in _download_and_prepare
    super()._download_and_prepare(
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/builder.py", line 1124, in _download_and_prepare
    self._prepare_split(split_generator, **prepare_split_kwargs)
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/builder.py", line 1629, in _prepare_split
    for job_id, done, content in self._prepare_split_single(
  File "/Users/yoavkatz/miniforge3/envs/fme/lib/python3.10/site-packages/datasets/builder.py", line 1786, in _prepare_split_single
    raise DatasetGenerationError("An error occurred while generating the dataset") from e
```

Need to add `split="test"` to make it work. Started in recent upgrade.

```
from unitxt.loaders import LoadHF

# Initialize the LoadHF loader
loader = LoadHF(path="data/rag/loan_agreement",data_files = {
            "test": "all_entity_dataset.jsonl"
        },split="test")

# Load the dataset
dataset = loader.process()

# Print the loaded dataset
print(dataset.to_dataset())
```

Discussion:
@elronbandel: Fixed in https://github.com/IBM/unitxt/pull/1724

Solution patch:
diff --git a/src/unitxt/loaders.py b/src/unitxt/loaders.py
index af6f520c..1887f950 100644
--- a/src/unitxt/loaders.py
+++ b/src/unitxt/loaders.py
@@ -258,7 +258,7 @@ class LoadHF(LazyLoader):
         split:
             Optional specification of which split to load.
         data_files:
-            Optional specification of particular data files to load.
+            Optional specification of particular data files to load. When you provide a list of data_files to Hugging Face's load_dataset function without explicitly specifying the split argument, these files are automatically placed into the train split.
         revision:
             Optional. The revision of the dataset. Often the commit id. Use in case you want to set the dataset version.
         streaming (bool):
@@ -351,6 +351,10 @@ class LoadHF(LazyLoader):
     def get_splits(self):
         if self.splits is not None:
             return self.splits
+        if self.data_files is not None:
+            if isinstance(self.data_files, dict):
+                return list(self.data_files.keys())
+            return ["train"]
         try:
             return hf_get_dataset_splits(
                 path=self.path,


Test patch:
diff --git a/tests/library/test_loaders.py b/tests/library/test_loaders.py
index db6d8eab..48866c07 100644
--- a/tests/library/test_loaders.py
+++ b/tests/library/test_loaders.py
@@ -255,6 +255,128 @@ class TestLoaders(UnitxtTestCase):
         )  # that HF dataset only has the 'test' split
         self.assertEqual(instance["language"], "eng")
 
+    def test_load_from_hf_with_data_files_dict(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            dummy_dir = os.path.join(temp_dir, "dummy")
+            os.makedirs(dummy_dir, exist_ok=True)
+
+            dummy_test_file_path = os.path.join(dummy_dir, "dummy_test_file.jsonl")
+            dummy_random_file_path = os.path.join(dummy_dir, "dummy_random_file.jsonl")
+
+            sample_data = [
+                {"id": 1, "text": "Sample text 1", "label": "A"},
+            ]
+
+            with open(dummy_test_file_path, "w") as f:
+                for item in sample_data:
+                    f.write(json.dumps(item) + "\n")
+
+            with open(dummy_random_file_path, "w") as f:
+                for item in sample_data:
+                    f.write(json.dumps(item) + "\n")
+
+            loader = LoadHF(
+                path=dummy_dir,
+                data_files={
+                    "test": "dummy_test_file.jsonl",
+                    "random":  "dummy_random_file.jsonl",
+                }
+            )
+
+            dataset = loader.process()
+
+            self.assertIn("test", dataset)
+            self.assertIn("random", dataset)
+
+            test = list(dataset["test"])
+
+            self.assertEqual(len(test), 1)
+
+            random = list(dataset["random"])
+            self.assertEqual(len(random), 1)
+
+            self.assertIn("text", random[0])
+            self.assertIn("label", random[0])
+            self.assertIn("text", test[0])
+            self.assertIn("label", test[0])
+
+    def test_load_from_hf_with_data_files_list(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            dummy_dir = os.path.join(temp_dir, "dummy")
+            os.makedirs(dummy_dir, exist_ok=True)
+
+            dummy_test_file_path = os.path.join(dummy_dir, "dummy_test_file.jsonl")
+            dummy_random_file_path = os.path.join(dummy_dir, "dummy_random_file.jsonl")
+
+            sample_data = [
+                {"id": 1, "text": "Sample text 1", "label": "A"},
+            ]
+
+            with open(dummy_test_file_path, "w") as f:
+                for item in sample_data:
+                    f.write(json.dumps(item) + "\n")
+
+            with open(dummy_random_file_path, "w") as f:
+                for item in sample_data:
+                    f.write(json.dumps(item) + "\n")
+
+            loader = LoadHF(
+                path=dummy_dir,
+                data_files=[
+                    "dummy_test_file.jsonl",
+                    "dummy_random_file.jsonl",
+                ]
+            )
+
+            dataset = loader.process()
+
+            self.assertIn("train", dataset)
+
+            train = list(dataset["train"])
+
+            self.assertEqual(len(train), 2)
+
+
+            self.assertIn("text", train[0])
+            self.assertIn("label", train[0])
+            self.assertIn("text", train[1])
+            self.assertIn("label", train[1 ])
+
+    def test_load_from_hf_with_data_files_str(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            dummy_dir = os.path.join(temp_dir, "dummy")
+            os.makedirs(dummy_dir, exist_ok=True)
+
+            dummy_test_file_path = os.path.join(dummy_dir, "dummy_test_file.jsonl")
+
+            sample_data = [
+                {"id": 1, "text": "Sample text 1", "label": "A"},
+            ]
+
+            with open(dummy_test_file_path, "w") as f:
+                for item in sample_data:
+                    f.write(json.dumps(item) + "\n")
+
+
+            loader = LoadHF(
+                path=dummy_dir,
+                data_files="dummy_test_file.jsonl",
+            )
+
+            # Load the dataset
+            dataset = loader.process()
+
+            self.assertIn("train", dataset)
+
+            train = list(dataset["train"])
+
+            self.assertEqual(len(train), 1)
+
+
+            self.assertIn("text", train[0])
+            self.assertIn("label", train[0])
+
+
     def test_multiple_source_loader(self):
         # Using a context for the temporary directory
         with tempfile.TemporaryDirectory() as tmp_dir: