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

Update trainer.py

Browse Source
This commit is contained in:
paNikitin 2025-02-24 09:12:31 +03:00
parent 231f5e870e
commit e2ace6fb0f
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
llms/mlx_lm/tuner/trainer.py
View File

@ -111,7 +111,9 @@ def cot_loss(
data_token_id = tokenizer.encode(args.data_token)[0] data_token_id = tokenizer.encode(args.data_token)[0]
reasoning_positions = mx.argmax(targets == reasoning_token_id, axis=1) reasoning_positions = mx.argmax(targets == reasoning_token_id, axis=1)
data_positions = mx.argmax(targets == data_token_id, axis=1) # find the LAST occurrence of data_token_id using slicing (in case generated dataset has multiple occurrences of [DATA])
data_positions = mx.argmax(targets[:, ::-1] == data_token_id, axis=1)
data_positions = targets.shape[1] - 1 - data_positions
seq_indices = mx.arange(targets.shape[1])[None, :] seq_indices = mx.arange(targets.shape[1])[None, :]