Support for custom files for run_lora_clm.py (huggingface#1039)

Co-authored-by: regisss <15324346+regisss@users.noreply.github.com>

Author: vidyasiv

Makefile

@@ -72,6 +72,11 @@ fast_tests_table_transformers:
 	python -m pip install .[tests]
 	python -m pytest tests/test_table_transformer.py
 
+# Run non-performance regressions
+slow_tests_custom_file_input: test_installs
+	python -m pip install -r examples/language-modeling/requirements.txt
+	python -m pytest tests/test_custom_file_input.py
+
 # Run single-card non-regression tests
 slow_tests_1x: test_installs
 	python -m pytest tests/test_examples.py -v -s -k "single_card"

examples/language-modeling/README.md

@@ -693,6 +693,48 @@ DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED=1 LOWER_LIST=ops_bf16.txt python3 ..
 ```
 Default `peft_type` is `lora`, you could enable adalora or ia3 using `--peft_type adalora` or `--peft_type ia3`, or enable llama-adapter for llama model using `--peft_type llama-adapter`.
 
+#### Custom Files
+
+To run on your own training and validation files, use the following command:
+
+```bash
+python run_lora_clm.py \
+    --model_name_or_path bigcode/starcoder \
+    --train_file path_to_train_file \
+    --validation_file path_to_validation_file \
+    --per_device_train_batch_size 8 \
+    --per_device_eval_batch_size 8 \
+    --do_train \
+    --do_eval \
+    --output_dir /tmp/test-lora-clm \
+    --bf16 \
+    --use_habana \
+    --use_lazy_mode \
+    --use_hpu_graphs_for_inference \
+    --dataset_concatenation \
+    --throughput_warmup_steps 3
+```
+
+The format of the jsonlines files (with extensions .json or .jsonl) is expected to be
+
+```json
+{"text": "<text>"}
+{"text": "<text>"}
+{"text": "<text>"}
+{"text": "<text>"}
+```
+
+The format of the text files (with extensions .text or .txt) is expected to be
+
+```json
+"<text>"
+"<text>"
+"<text>"
+"<text>"
+```
+
+> Note: When using both custom files i.e `--train_file` and `--validation_file`, all files are expected to be of the same type i.e json or text.
+
 ### Prompt/Prefix/P-tuning
 
 To run prompt tuning finetuning, you can use `run_prompt_tuning_clm.py`.

examples/language-modeling/run_lora_clm.py

@@ -240,6 +240,9 @@ class DataArguments:
         default=False,
         metadata={"help": "Whether to keep in memory the loaded dataset. Defaults to False."},
     )
+    keep_linebreaks: bool = field(
+        default=True, metadata={"help": "Whether to keep line breaks when using TXT files or not."}
+    )
     dataset_seed: int = field(
         default=42,
         metadata={
@@ -510,6 +513,10 @@ def main():
         )
 
         if "validation" not in raw_datasets.keys() and training_args.do_eval:
+            if not data_args.validation_split_percentage:
+                raise ValueError(
+                    "Please set --validation_split_percentage as dataset does not contain `validation` key"
+                )
             raw_datasets["validation"] = load_dataset(
                 data_args.dataset_name,
                 data_args.dataset_config_name,
@@ -538,9 +545,11 @@ def main():
             if data_args.train_file is not None
             else data_args.validation_file.split(".")[-1]
         )
-        if extension == "txt":
+        if extension in ("txt", "text"):
             extension = "text"
             dataset_args["keep_linebreaks"] = data_args.keep_linebreaks
+        if extension in ("json", "jsonl"):
+            extension = "json"
         raw_datasets = load_dataset(
             extension,
             data_files=data_files,
@@ -551,6 +560,10 @@ def main():
 
         # If no validation data is there, validation_split_percentage will be used to divide the dataset.
         if "validation" not in raw_datasets.keys() and training_args.do_eval:
+            if not data_args.validation_split_percentage:
+                raise ValueError(
+                    "Please set --validation_split_percentage as dataset does not contain `validation` key"
+                )
             raw_datasets["validation"] = load_dataset(
                 extension,
                 data_files=data_files,
@@ -567,20 +580,32 @@ def main():
                 token=model_args.token,
                 **dataset_args,
             )
-
+    single_column_dataset = False
+    # For named dataset (timdettmers/openassistant-guanaco) or custom dataset with a single column "text"
     if (
-        data_args.dataset_name == "timdettmers/openassistant-guanaco"
-    ):  # from https://github.com/artidoro/qlora/blob/main/qlora.py#L621
+        training_args.do_train
+        and raw_datasets["train"].num_columns == 1
+        or training_args.do_eval
+        and raw_datasets["validation"].num_columns == 1
+    ):
+        single_column_dataset = True
         raw_datasets = raw_datasets.map(
             lambda x: {
                 "input": "",
                 "output": x["text"],
             }
         )
-        # Remove unused columns.
-        raw_datasets = raw_datasets.remove_columns(
-            [col for col in raw_datasets.column_names["train"] if col not in ["input", "output"]]
-        )
+        if training_args.do_train:
+            # Remove unused columns.
+            raw_datasets = raw_datasets.remove_columns(
+                [col for col in raw_datasets.column_names["train"] if col not in ["input", "output"]]
+            )
+
+        if training_args.do_eval:
+            # Remove unused columns.
+            raw_datasets = raw_datasets.remove_columns(
+                [col for col in raw_datasets.column_names["validation"] if col not in ["input", "output"]]
+            )
     else:
         # Preprocessing the datasets.
         for key in raw_datasets:
@@ -680,7 +705,7 @@ def tokenize(prompt, add_eos_token=True):
     def preprocess_function(examples):
         keys = list(examples.data.keys())
         if len(keys) != 2:
-            raise ValueError("Unsupported dataset format")
+            raise ValueError(f"Unsupported dataset format, number of keys {keys} !=2")
 
         st = [s + t for s, t in zip(examples[keys[0]], examples[keys[1]])]
 
@@ -717,17 +742,18 @@ def concatenate_data(dataset, max_seq_length):
                 concatenated_dataset[column] = reshaped_data
             return datasets.Dataset.from_dict(concatenated_dataset)
 
-        if data_args.dataset_name == "timdettmers/openassistant-guanaco":
+        if single_column_dataset:
             tokenized_datasets_ = tokenized_datasets["train"].remove_columns(["input", "output"])
             if training_args.do_eval:
-                tokenized_datasets_eval_ = tokenized_datasets["test"].remove_columns(["input", "output"])
+                tokenized_datasets_eval_ = tokenized_datasets["validation"].remove_columns(["input", "output"])
         else:
             tokenized_datasets_ = tokenized_datasets["train"].remove_columns(["prompt_sources", "prompt_targets"])
             if training_args.do_eval:
                 tokenized_datasets_eval_ = tokenized_datasets["validation"].remove_columns(
                     ["prompt_sources", "prompt_targets"]
                 )
-        tokenized_datasets["train"] = concatenate_data(tokenized_datasets_, data_args.max_seq_length)
+        if training_args.do_train:
+            tokenized_datasets["train"] = concatenate_data(tokenized_datasets_, data_args.max_seq_length)
         if training_args.do_eval:
             tokenized_datasets["validation"] = concatenate_data(tokenized_datasets_eval_, data_args.max_seq_length)
     if training_args.do_train:
@@ -849,7 +875,7 @@ def compute_metrics(eval_preds):
         trainer.log_metrics("train", metrics)
         trainer.save_metrics("train", metrics)
 
-        # Evaluation
+    # Evaluation
     if training_args.do_eval:
         logger.info("*** Evaluate ***")
         metrics = trainer.evaluate()