From aa1c8abdc67be20a1efcacb0e3102ac732a87967 Mon Sep 17 00:00:00 2001 From: madroid Date: Mon, 30 Sep 2024 22:36:21 +0800 Subject: [PATCH] LoRA: Support HuggingFace dataset via data parameter (#996) * LoRA: support huggingface dataset via `data` argument * LoRA: Extract the load_custom_hf_dataset function * LoRA: split small functions * fix spelling errors * handle load hf dataset error * fix pre-commit lint * update data argument help * nits and doc --------- Co-authored-by: Awni Hannun --- llms/mlx_lm/LORA.md | 8 ++- llms/mlx_lm/lora.py | 5 +- llms/mlx_lm/tuner/datasets.py | 131 +++++++++++++++++++++------------- 3 files changed, 93 insertions(+), 51 deletions(-) diff --git a/llms/mlx_lm/LORA.md b/llms/mlx_lm/LORA.md index 80c25b4..2d0dcf6 100644 --- a/llms/mlx_lm/LORA.md +++ b/llms/mlx_lm/LORA.md @@ -251,7 +251,13 @@ To use Hugging Face datasets, first install the `datasets` package: pip install datasets ``` -Specify the Hugging Face dataset arguments in a YAML config. For example: +If the Hugging Face dataset is already in a supported format, you can specify +it on the command line. For example, pass `--data mlx-community/wikisql` to +train on the pre-formatted WikiwSQL data. + +Otherwise, provide a mapping of keys in the dataset to the features MLX LM +expects. Use a YAML config to specify the Hugging Face dataset arguments. For +example: ``` hf_dataset: diff --git a/llms/mlx_lm/lora.py b/llms/mlx_lm/lora.py index 6923277..c96e75a 100644 --- a/llms/mlx_lm/lora.py +++ b/llms/mlx_lm/lora.py @@ -79,7 +79,10 @@ def build_parser(): parser.add_argument( "--data", type=str, - help="Directory with {train, valid, test}.jsonl files", + help=( + "Directory with {train, valid, test}.jsonl files or the name " + "of a Hugging Face dataset (e.g., 'mlx-community/wikisql')" + ), ) parser.add_argument( "--fine-tune-type", diff --git a/llms/mlx_lm/tuner/datasets.py b/llms/mlx_lm/tuner/datasets.py index 2b8abf4..20b32ef 100644 --- a/llms/mlx_lm/tuner/datasets.py +++ b/llms/mlx_lm/tuner/datasets.py @@ -76,17 +76,14 @@ class CompletionsDataset(Dataset): return text -def create_dataset(path: Path, tokenizer: PreTrainedTokenizer = None): - # Return empty dataset for non-existent paths - if not path.exists(): - return [] - with open(path, "r") as fid: - data = [json.loads(l) for l in fid] - if "messages" in data[0]: +def create_dataset(data, tokenizer: PreTrainedTokenizer = None): + sample = data[0] + + if "messages" in sample: return ChatDataset(data, tokenizer) - elif "prompt" in data[0] and "completion" in data[0]: + elif "prompt" in sample and "completion" in sample: return CompletionsDataset(data, tokenizer) - elif "text" in data[0]: + elif "text" in sample: return Dataset(data) else: raise ValueError( @@ -95,54 +92,90 @@ def create_dataset(path: Path, tokenizer: PreTrainedTokenizer = None): ) -def load_dataset(args, tokenizer: PreTrainedTokenizer): - if getattr(args, "hf_dataset", None) is not None: - import datasets +def load_local_dataset(data_path: Path, tokenizer: PreTrainedTokenizer): + def load_subset(path): + if not path.exists(): + return [] + with open(path, "r") as fid: + data = [json.loads(l) for l in fid] + return create_dataset(data, tokenizer) - hf_args = args.hf_dataset - dataset_name = hf_args["name"] - print(f"Loading Hugging Face dataset {dataset_name}.") - text_feature = hf_args.get("text_feature") - prompt_feature = hf_args.get("prompt_feature") - completion_feature = hf_args.get("completion_feature") + names = ("train", "valid", "test") + train, valid, test = [load_subset(data_path / f"{n}.jsonl") for n in names] + return train, valid, test - def create_hf_dataset(split: str = None): - ds = datasets.load_dataset( - dataset_name, - split=split, - **hf_args.get("config", {}), - ) - if prompt_feature and completion_feature: - return CompletionsDataset( - ds, tokenizer, prompt_feature, completion_feature - ) - elif text_feature: - return Dataset(train_ds, text_key=text_feature) - else: - raise ValueError( - "Specify either a prompt and completion feature or a text " - "feature for the Hugging Face dataset." - ) - if args.train: - train_split = hf_args.get("train_split", "train[:80%]") - valid_split = hf_args.get("valid_split", "train[-10%:]") - train = create_hf_dataset(split=train_split) - valid = create_hf_dataset(split=valid_split) - else: - train, valid = [], [] - if args.test: - test = create_hf_dataset(split=hf_args.get("test_split")) - else: - test = [] +def load_hf_dataset(data_id: str, tokenizer: PreTrainedTokenizer): + from datasets import exceptions, load_dataset + + try: + dataset = load_dataset(data_id) - else: names = ("train", "valid", "test") - data_path = Path(args.data) train, valid, test = [ - create_dataset(data_path / f"{n}.jsonl", tokenizer) for n in names + create_dataset(dataset[n], tokenizer) if n in dataset.keys() else [] + for n in names ] + + except exceptions.DatasetNotFoundError: + raise ValueError(f"Not found Hugging Face dataset: {data_id} .") + + return train, valid, test + + +def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer): + import datasets + + hf_args = args.hf_dataset + dataset_name = hf_args["name"] + print(f"Loading Hugging Face dataset {dataset_name}.") + text_feature = hf_args.get("text_feature") + prompt_feature = hf_args.get("prompt_feature") + completion_feature = hf_args.get("completion_feature") + + def create_hf_dataset(split: str = None): + ds = datasets.load_dataset( + dataset_name, + split=split, + **hf_args.get("config", {}), + ) + if prompt_feature and completion_feature: + return CompletionsDataset(ds, tokenizer, prompt_feature, completion_feature) + elif text_feature: + return Dataset(train_ds, text_key=text_feature) + else: + raise ValueError( + "Specify either a prompt and completion feature or a text " + "feature for the Hugging Face dataset." + ) + + if args.train: + train_split = hf_args.get("train_split", "train[:80%]") + valid_split = hf_args.get("valid_split", "train[-10%:]") + train = create_hf_dataset(split=train_split) + valid = create_hf_dataset(split=valid_split) + else: + train, valid = [], [] + if args.test: + test = create_hf_dataset(split=hf_args.get("test_split")) + else: + test = [] + + return train, valid, test + + +def load_dataset(args, tokenizer: PreTrainedTokenizer): + if getattr(args, "hf_dataset", None) is not None: + train, valid, test = load_custom_hf_dataset(args, tokenizer) + else: + data_path = Path(args.data) + if data_path.exists(): + train, valid, test = load_local_dataset(data_path, tokenizer) + else: + print(f"Loading Hugging Face dataset {args.data}.") + train, valid, test = load_hf_dataset(args.data, tokenizer) + if args.train and len(train) == 0: raise ValueError( "Training set not found or empty. Must provide training set for fine-tuning."