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

fixes

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2025-02-03 19:43:49 +01:00
parent 05d921b788
commit 40bca770ae
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
llms/mlx_lm/tuner/grpo_trainer.py
View File

@ -271,9 +271,9 @@ def grpo_loss(
rewards = mx.zeros((len(all_completions),))
for reward_func in reward_funcs:
func_rewards = mx.array(reward_func(
prompts=prompt_text,
prompts=expanded_prompts,
completions=all_completion_texts,
answer=answer_text
answer=expanded_answers
))
rewards += func_rewards
@ -310,9 +310,9 @@ def grpo_loss(
reward_metrics = {}
for i, reward_func in enumerate(reward_funcs):
func_rewards = mx.array(reward_func(
prompts=prompt_text,
prompts=expanded_prompts,
completions=all_completion_texts,
answer=answer_text
answer=expanded_answers
))
reward_metrics[f'reward_func_{i}_mean'] = mx.mean(func_rewards)
reward_metrics[f'reward_func_{i}_std'] = mx.std(func_rewards)