"""mlx-lm-lora 5.8.6 reward callback. Import loads no model."""

import importlib.util
from pathlib import Path

from mlx_lm_lora.trainer.grpo_reward_functions import register_reward_function

_spec = importlib.util.spec_from_file_location("price_lab", Path(__file__).with_name("lab.py"))
_lab = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_lab)


@register_reward_function("price_reward")
def price_reward(prompts, completions, answer, types=None):
    # Callback parameter is singular "answer" in this implementation.
    if len(completions) != len(answer):
        raise ValueError("Reward batch lengths differ")
    return [_lab.grade(text, int(target))["reward"]
            for text, target in zip(completions, answer)]