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

fix testing

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2025-02-04 11:02:00 +01:00
parent 9b489a6c0c
commit c2fcb6738b
2 changed files with 28 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
llms/mlx_lm/lora.py
View File

@ -312,11 +312,10 @@ def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set
else: else:
reference_model = model reference_model = model
test_loss, test_rewards = evaluate_dpo( test_loss, _, _, _ = evaluate_dpo(
model=model, model=model,
ref_model=reference_model, ref_model=reference_model,
dataset=test_set, dataset=test_set,
tokenizer=tokenizer,
batch_size=args.batch_size, batch_size=args.batch_size,
num_batches=args.test_batches, num_batches=args.test_batches,
max_seq_length=args.max_seq_length, max_seq_length=args.max_seq_length,
@ -324,7 +323,8 @@ def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set
delta=args.delta, delta=args.delta,
loss_type=args.dpo_loss_type, loss_type=args.dpo_loss_type,
) )
print(f"Test loss {test_loss:.3f}, Rewards: {test_rewards[0]:.3f}, {test_rewards[1]:.3f}") print(f"Test loss {test_loss:.3f}")
else: else:
test_loss = evaluate( test_loss = evaluate(
model=model, model=model,

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

@ -119,6 +119,7 @@ class CompletionsDataset:
def create_dataset( def create_dataset(
args,
data, data,
tokenizer: PreTrainedTokenizer, tokenizer: PreTrainedTokenizer,
prompt_feature: Optional[str] = None, prompt_feature: Optional[str] = None,
@ -127,24 +128,30 @@ def create_dataset(
prompt_feature = prompt_feature or "prompt" prompt_feature = prompt_feature or "prompt"
completion_feature = completion_feature or "completion" completion_feature = completion_feature or "completion"
sample = data[0] sample = data[0]
# Add DPO dataset support if args.training_mode == "normal":
if "chosen" in sample and "rejected" in sample: if "messages" in sample:
return DPODataset(data, tokenizer) return ChatDataset(data, tokenizer)
elif "messages" in sample: elif prompt_feature in sample and completion_feature in sample:
return ChatDataset(data, tokenizer) return CompletionsDataset(data, tokenizer, prompt_feature, completion_feature)
elif prompt_feature in sample and completion_feature in sample: elif "text" in sample:
return CompletionsDataset(data, tokenizer, prompt_feature, completion_feature) return Dataset(data, tokenizer)
elif "text" in sample: else:
return Dataset(data, tokenizer) 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."
)
elif args.training_mode == "dpo":
if "chosen" in sample and "rejected" in sample:
return DPODataset(data, tokenizer)
else: else:
raise ValueError( raise ValueError(
"Unsupported data format, check the supported formats here:\n" "Unsupported training mode, check the supported training modes and their formats here:\n"
"https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md#data." "https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md#training-modes."
) )
def load_local_dataset( def load_local_dataset(
args,
data_path: Path, data_path: Path,
tokenizer: PreTrainedTokenizer, tokenizer: PreTrainedTokenizer,
prompt_feature: Optional[str] = None, prompt_feature: Optional[str] = None,
@ -155,7 +162,7 @@ def load_local_dataset(
return [] return []
with open(path, "r") as fid: with open(path, "r") as fid:
data = [json.loads(l) for l in fid] data = [json.loads(l) for l in fid]
return create_dataset(data, tokenizer, prompt_feature, completion_feature) return create_dataset(args, data, tokenizer, prompt_feature, completion_feature)
names = ("train", "valid", "test") names = ("train", "valid", "test")
train, valid, test = [load_subset(data_path / f"{n}.jsonl") for n in names] train, valid, test = [load_subset(data_path / f"{n}.jsonl") for n in names]
@ -163,6 +170,7 @@ def load_local_dataset(
def load_hf_dataset( def load_hf_dataset(
args,
data_id: str, data_id: str,
tokenizer: PreTrainedTokenizer, tokenizer: PreTrainedTokenizer,
prompt_feature: Optional[str] = None, prompt_feature: Optional[str] = None,
@ -178,7 +186,7 @@ def load_hf_dataset(
train, valid, test = [ train, valid, test = [
( (
create_dataset( create_dataset(
dataset[n], tokenizer, prompt_feature, completion_feature args, dataset[n], tokenizer, prompt_feature, completion_feature
) )
if n in dataset.keys() if n in dataset.keys()
else [] else []
@ -243,12 +251,12 @@ def load_dataset(args, tokenizer: PreTrainedTokenizer):
completion_feature = getattr(args, "completion_feature", None) completion_feature = getattr(args, "completion_feature", None)
if data_path.exists(): if data_path.exists():
train, valid, test = load_local_dataset( train, valid, test = load_local_dataset(
data_path, tokenizer, prompt_feature, completion_feature args, data_path, tokenizer, prompt_feature, completion_feature
) )
else: else:
print(f"Loading Hugging Face dataset {args.data}.") print(f"Loading Hugging Face dataset {args.data}.")
train, valid, test = load_hf_dataset( train, valid, test = load_hf_dataset(
args.data, tokenizer, prompt_feature, completion_feature args, args.data, tokenizer, prompt_feature, completion_feature
) )
if args.train and len(train) == 0: if args.train and len(train) == 0: