zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-07-02 06:41:13 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

better create_dataset

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2025-02-04 10:43:00 +01:00
parent bd1a42ec2f
commit 7b0141455e
1 changed files with 13 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
llms/mlx_lm/tuner/datasets.py
View File

@ -141,9 +141,19 @@ def create_dataset(
completion_feature = completion_feature or "completion"
sample = data[0]
if "messages" in sample:
return ChatDataset(data, tokenizer)
elif "prompt" in sample and "answer" in sample:
if args.training_mode == "normal":
if "messages" in sample:
return ChatDataset(data, tokenizer)
elif prompt_feature in sample and completion_feature in sample:
return CompletionsDataset(data, tokenizer, prompt_feature, completion_feature)
elif "text" in sample:
return Dataset(data, tokenizer)
else:
raise ValueError(
"Unsupported data format, check the supported formats here:\n"
"https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md#data."
)
else:
return GRPODataset(
data=data,
tokenizer=tokenizer,
@ -152,15 +162,6 @@ def create_dataset(
use_chat_template=args.use_chat_template,
use_prompt=args.use_prompt
)
elif prompt_feature in sample and completion_feature in sample:
return CompletionsDataset(data, tokenizer, prompt_feature, completion_feature)
elif "text" in sample:
return Dataset(data, tokenizer)
else:
raise ValueError(
"Unsupported data format, check the supported formats here:\n"
"https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md#data."
)
def load_local_dataset(