Text-generation: Support text templates, Special tokens in corpus join, Batched processing

examples/open_llama/README.md

@@ -73,11 +73,21 @@ The relevant config file is [open_llama_sparsegpt_gpu.json](open_llama_sparsegpt
 
 Requirements file: [requirements-sparsegpt.txt](requirements-sparsegpt.txt)
 
-### Fine-tune Llama Model using QLoRA
+### Fine-tune Llama Model on a chatbot dataset using QLoRA
 This workflow fine-tunes LLaMA model using [QLoRA](https://arxiv.org/abs/2305.14314). The output model is still the input transformers model along with a quantization config and
 LoRA adapters that were fine-tuned on the training dataset.
 
-The relevant configfile is [llama_qlora.json](llama_qlora.json). It corresponds to the [guqnaco 7b example in the original qlora implementation](https://github.com/artidoro/qlora/blob/main/scripts/finetune_guanaco_7b.sh).
+The relevant config file is [llama_qlora.json](llama_qlora.json). It corresponds to the [guqnaco 7b example in the original qlora implementation](https://github.com/artidoro/qlora/blob/main/scripts/finetune_guanaco_7b.sh).
+
+Requirements file: [requirements-qlora.txt](requirements-qlora.txt)
+
+### Fine-tune Open Llama Model on a code generation dataset using QLoRA
+This workflow fine-tunes Open LLaMA model using [QLoRA] to generate code given a prompt.
+
+The relevant config file is [open_llama_qlora_tinycodes.json](open_llama_qlora_tinycodes.json). The code language is set to `Python` but can be changed to other languages by changing the `language` field in the config file.
+Supported languages are Python, TypeScript, JavaScript, Ruby, Julia, Rust, C++, Bash, Java, C#, and Go. Refer to the [dataset card](https://huggingface.co/datasets/nampdn-ai/tiny-codes) for more details on the dataset.
+
+Note: You must be logged in to HuggingFace using `huggingface-cli login` to download the dataset or update `token` field in the config file with your HuggingFace token.
 
 Requirements file: [requirements-qlora.txt](requirements-qlora.txt)
 

examples/open_llama/open_llama_qlora_tinycodes.json

@@ -0,0 +1,86 @@
+{
+    "input_model":{
+        "type": "PyTorchModel",
+        "config": {
+            "hf_config": {
+                "model_name": "openlm-research/open_llama_7b_v2",
+                "task": "text-generation"
+            }
+        }
+    },
+    "data_configs": {
+        "tiny-codes-train": {
+            "name": "tiny-codes-train",
+            "type": "HuggingfaceContainer",
+            "user_script": "qlora_user_script.py",
+            "components": {
+                "load_dataset": {
+                    "type": "load_tiny_code_dataset"
+                }
+            },
+            "params_config": {
+                "dataset_name": "nampdn-ai/tiny-codes",
+                "split": "train",
+                "component_kwargs": {
+                    "load_dataset": {
+                        "language": "Python",
+                        "token": true
+                    },
+                    "pre_process_data": {
+                        "dataset_type": "corpus",
+                        "corpus_strategy": "join",
+                        "text_template": "### Question: {prompt} \n### Answer: {response}",
+                        "source_max_len": 1024
+                    }
+                }
+            }
+        }
+    },
+    "passes": {
+        "qlora": {
+            "type": "QLoRA",
+            "config": {
+                "compute_dtype": "bfloat16",
+                "quant_type": "nf4",
+                "double_quant": true,
+                "lora_r": 64,
+                "lora_alpha": 16,
+                "lora_dropout": 0.1,
+                "train_data_config": "tiny-codes-train",
+                "eval_dataset_size": 1024,
+                "training_args": {
+                    "seed": 0,
+                    "data_seed": 42,
+                    "per_device_train_batch_size": 16,
+                    "per_device_eval_batch_size": 16,
+                    "gradient_accumulation_steps": 1,
+                    "gradient_checkpointing": true,
+                    "learning_rate": 0.0002,
+                    "max_steps": 1500,
+                    "logging_steps": 10,
+                    "evaluation_strategy": "steps",
+                    "eval_steps": 100,
+                    "save_steps": 500,
+                    "group_by_length": true,
+                    "adam_beta2": 0.999,
+                    "max_grad_norm": 0.3,
+                    "load_best_model_at_end": true
+                }
+            }
+        }
+    },
+    "engine": {
+        "log_severity_level": 0,
+        "search_strategy": false,
+        "evaluate_input_model": false,
+        "target": {
+            "type": "LocalSystem",
+            "config": {
+                "accelerators": ["gpu"]
+            }
+        },
+        "execution_providers": ["CPUExecutionProvider"],
+        "cache_dir": "cache",
+        "output_dir" : "models/qlora-tiny-codes"
+    }
+}

examples/open_llama/open_llama_sparsegpt_gpu.json

@@ -22,6 +22,7 @@
                         "dataset_type": "corpus",
                         "text_cols": ["text"],
                         "corpus_strategy": "join-random",
+                        "add_special_tokens": false,
                         "source_max_len": 2048,
                         "max_samples": 128,
                         "random_seed": 42
@@ -41,6 +42,7 @@
                         "dataset_type": "corpus",
                         "text_cols": ["text"],
                         "corpus_strategy": "join",
+                        "add_special_tokens": false,
                         "source_max_len": 2048
                     }
                 }

examples/open_llama/qlora_user_script.py

@@ -0,0 +1,17 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from typing import Union
+
+from datasets import load_dataset
+
+from olive.data.registry import Registry
+
+
+# TODO: remove custom dataset component once default dataset component supports filter, tokens and split
+@Registry.register_dataset()
+def load_tiny_code_dataset(dataset_name: str, split: str, language: str, token: Union[bool, str] = True):
+    dataset = load_dataset(dataset_name, split=split, token=token)
+    dataset = dataset.filter(lambda x: x["programming_language"] == language)
+    return dataset