8b2b99be03* tool mask support * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handle tool masks with older zoo builds * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep tool mask implementation in zoo --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
"""Compatibility checks for env/tool mask support with older unsloth_zoo."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import os
|
|
import textwrap
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
|
|
RL_SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl.py")
|
|
RL_REPLACEMENTS_SOURCE_PATH = os.path.join(
|
|
REPO_ROOT, "unsloth", "models", "rl_replacements.py"
|
|
)
|
|
|
|
|
|
def _read(path: str) -> str:
|
|
with open(path, "r") as fh:
|
|
return fh.read()
|
|
|
|
|
|
def _load_local_align_completion_tool_mask():
|
|
src = _read(RL_SOURCE_PATH)
|
|
tree = ast.parse(src)
|
|
for node in tree.body:
|
|
if isinstance(node, ast.If):
|
|
for item in node.body:
|
|
if (
|
|
isinstance(item, ast.FunctionDef)
|
|
and item.name == "align_completion_tool_mask"
|
|
):
|
|
function_src = ast.get_source_segment(src, item)
|
|
break
|
|
else:
|
|
continue
|
|
break
|
|
else:
|
|
raise AssertionError("local align_completion_tool_mask fallback is missing")
|
|
|
|
calls = []
|
|
|
|
def align_logprobs_with_mask(logprob_tensor, completion_mask, pad_value = None):
|
|
calls.append((logprob_tensor, completion_mask, pad_value))
|
|
return torch.tensor(
|
|
[[1, 0, 1], [0, 1, 1]],
|
|
device = completion_mask.device,
|
|
dtype = logprob_tensor.dtype,
|
|
)
|
|
|
|
namespace = {
|
|
"torch": torch,
|
|
"align_logprobs_with_mask": align_logprobs_with_mask,
|
|
}
|
|
exec(textwrap.dedent(function_src), namespace)
|
|
return namespace["align_completion_tool_mask"], calls
|
|
|
|
|
|
def test_rl_uses_optional_zoo_tool_mask_helper():
|
|
src = _read(RL_SOURCE_PATH)
|
|
assert 'RL_REPLACEMENTS.get("align_completion_tool_mask")' in src
|
|
assert 'RL_REPLACEMENTS["align_completion_tool_mask"]' not in src
|
|
|
|
|
|
def test_local_tool_mask_fallback_is_only_old_zoo_compat_shim():
|
|
align_completion_tool_mask, calls = _load_local_align_completion_tool_mask()
|
|
completion_mask = torch.tensor(
|
|
[[1, 1, 0], [1, 1, 1]],
|
|
dtype = torch.float32,
|
|
)
|
|
|
|
assert align_completion_tool_mask(None, completion_mask) is completion_mask
|
|
assert calls == []
|
|
|
|
same_shape_tool_mask = torch.tensor([[1, 0, 1], [0, 1, 1]], dtype = torch.bool)
|
|
with pytest.raises(RuntimeError, match = "Please upgrade unsloth_zoo"):
|
|
align_completion_tool_mask(same_shape_tool_mask, completion_mask)
|
|
|
|
|
|
def test_grpo_accumulated_loss_omits_none_tool_mask_for_old_zoo():
|
|
src = _read(RL_REPLACEMENTS_SOURCE_PATH)
|
|
assert "_grpo_accumulated_loss_kwargs = {}" in src
|
|
assert (
|
|
'if tool_mask is not None:\n _grpo_accumulated_loss_kwargs["tool_mask"] = tool_mask'
|
|
in src
|
|
)
|
|
assert src.count("**_grpo_accumulated_loss_kwargs") == 2
|
|
|
|
accelerated_loss_start = src.find('if hasattr(self.args, "loss_type"):')
|
|
assert accelerated_loss_start != -1
|
|
accelerated_loss_body = src[
|
|
accelerated_loss_start : src.find(
|
|
'if "train" in self._metrics:', accelerated_loss_start
|
|
)
|
|
]
|
|
assert "tool_mask = tool_mask" not in accelerated_loss_body
|
|
|
|
|
|
def test_rollout_output_patch_requires_real_tool_mask_symbol():
|
|
src = _read(RL_REPLACEMENTS_SOURCE_PATH)
|
|
assert 're.search(r"\\btool_mask\\b", function)' in src
|
|
assert 'output["tool_mask"]' in src
|