mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-10 00:08:58 +00:00
3ce187da02
Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.
247 lines
7.8 KiB
Python
247 lines
7.8 KiB
Python
"""ORPO should use a processor's tokenizer for text-only row tokenization."""
|
|
|
|
import ast
|
|
import os
|
|
import re
|
|
|
|
|
|
REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
RL_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py")
|
|
|
|
|
|
def _load_orpo_rewriter(name = "orpo_trainer_text_tokenizer"):
|
|
src = open(RL_PATH).read()
|
|
tree = ast.parse(src)
|
|
ns = {"re": re}
|
|
# Materialise sibling module-level assignments (e.g. _PAD_FALLBACK) so
|
|
# any rewriter that references them at exec-time can resolve them.
|
|
for node in tree.body:
|
|
if isinstance(node, ast.Assign):
|
|
for target in node.targets:
|
|
if isinstance(target, ast.Name) and target.id.startswith("_"):
|
|
exec(ast.get_source_segment(src, node), ns)
|
|
for node in tree.body:
|
|
if isinstance(node, ast.FunctionDef) and node.name == name:
|
|
exec(ast.get_source_segment(src, node), ns)
|
|
return ns[name]
|
|
raise AssertionError(f"{name} not found")
|
|
|
|
|
|
class _Tokenizer:
|
|
bos_token_id = 1
|
|
eos_token_id = 2
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def __call__(
|
|
self,
|
|
text,
|
|
add_special_tokens = False,
|
|
**kwargs,
|
|
):
|
|
self.calls.append((text, add_special_tokens, kwargs))
|
|
ids = [ord(c) % 31 + 3 for c in text]
|
|
return {"input_ids": ids, "attention_mask": [1] * len(ids)}
|
|
|
|
|
|
class _Processor:
|
|
def __init__(self):
|
|
self.tokenizer = _Tokenizer()
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
raise AssertionError("text-only ORPO tokenization should not call processor")
|
|
|
|
|
|
class _Trainer:
|
|
def __init__(self):
|
|
self.processing_class = _Processor()
|
|
self.is_encoder_decoder = False
|
|
self.max_length = 2048
|
|
self.max_prompt_length = 1024
|
|
self.max_completion_length = 1024
|
|
self.truncation_mode = "keep_end"
|
|
self.label_pad_token_id = -100
|
|
self.padding_value = 0
|
|
|
|
|
|
def _exec_rewritten(
|
|
function_name,
|
|
source,
|
|
extra_ns = None,
|
|
):
|
|
rewriter = _load_orpo_rewriter()
|
|
rewritten = rewriter(function_name, source)
|
|
ns = {} if extra_ns is None else dict(extra_ns)
|
|
exec(rewritten, ns)
|
|
return ns[function_name]
|
|
|
|
|
|
def test_orpo_tokenize_row_returns_original_when_tokenizer_anchor_missing():
|
|
rewriter = _load_orpo_rewriter()
|
|
source = """
|
|
def tokenize_row(self, feature, model=None):
|
|
output = {}
|
|
output["prompt_input_ids"] = self.processing_class(feature["prompt"], add_special_tokens=False)["input_ids"]
|
|
return output
|
|
"""
|
|
|
|
rewritten = rewriter("tokenize_row", source)
|
|
|
|
assert rewritten == source
|
|
assert "tokenizer(" not in rewritten
|
|
|
|
|
|
def test_orpo_build_tokenized_answer_uses_processor_tokenizer():
|
|
source = """
|
|
def build_tokenized_answer(self, prompt, answer):
|
|
full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False)
|
|
prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"]
|
|
return full_tokenized["input_ids"][len(prompt_input_ids):]
|
|
"""
|
|
fn = _exec_rewritten("build_tokenized_answer", source)
|
|
trainer = _Trainer()
|
|
|
|
assert fn(trainer, "a", "b")
|
|
assert [call[0] for call in trainer.processing_class.tokenizer.calls] == ["ab", "a"]
|
|
|
|
|
|
def test_orpo_tokenize_row_uses_processor_tokenizer():
|
|
source = """
|
|
def tokenize_row(self, feature, model=None):
|
|
batch = {}
|
|
prompt = feature["prompt"]
|
|
chosen = feature["chosen"]
|
|
rejected = feature["rejected"]
|
|
if not self.is_encoder_decoder:
|
|
prompt_tokens = self.processing_class(prompt, add_special_tokens=False)
|
|
prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()}
|
|
chosen_tokens = self.build_tokenized_answer(prompt, chosen)
|
|
rejected_tokens = self.build_tokenized_answer(prompt, rejected)
|
|
prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"])
|
|
chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"])
|
|
rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"])
|
|
prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed(
|
|
self.processing_class.bos_token_id,
|
|
prompt_len_input_ids,
|
|
prompt_tokens,
|
|
chosen_prompt_len_input_ids,
|
|
chosen_tokens,
|
|
rejected_prompt_len_input_ids,
|
|
rejected_tokens,
|
|
)
|
|
chosen_tokens, rejected_tokens = add_eos_token_if_needed(
|
|
self.processing_class.eos_token_id, chosen_tokens, rejected_tokens
|
|
)
|
|
batch["prompt_input_ids"] = prompt_tokens["prompt_input_ids"]
|
|
batch["chosen_input_ids"] = chosen_tokens["input_ids"]
|
|
batch["rejected_input_ids"] = rejected_tokens["input_ids"]
|
|
return batch
|
|
"""
|
|
|
|
def add_bos_token_if_needed(*args):
|
|
return args[2], args[4], args[6]
|
|
|
|
def add_eos_token_if_needed(eos_token_id, chosen_tokens, rejected_tokens):
|
|
chosen_tokens["input_ids"] = chosen_tokens["input_ids"] + [eos_token_id]
|
|
rejected_tokens["input_ids"] = rejected_tokens["input_ids"] + [eos_token_id]
|
|
return chosen_tokens, rejected_tokens
|
|
|
|
trainer = _Trainer()
|
|
trainer.build_tokenized_answer = lambda prompt, answer: {
|
|
"prompt_input_ids": trainer.processing_class.tokenizer(prompt)["input_ids"],
|
|
"input_ids": trainer.processing_class.tokenizer(answer)["input_ids"],
|
|
}
|
|
fn = _exec_rewritten(
|
|
"tokenize_row",
|
|
source,
|
|
{
|
|
"add_bos_token_if_needed": add_bos_token_if_needed,
|
|
"add_eos_token_if_needed": add_eos_token_if_needed,
|
|
},
|
|
)
|
|
|
|
output = fn(trainer, {"prompt": "p", "chosen": "c", "rejected": "r"})
|
|
|
|
assert output["chosen_input_ids"][-1] == _Tokenizer.eos_token_id
|
|
assert [call[0] for call in trainer.processing_class.tokenizer.calls] == [
|
|
"p",
|
|
"p",
|
|
"c",
|
|
"p",
|
|
"r",
|
|
]
|
|
|
|
|
|
def test_orpo_init_pad_token_id_falls_back_to_tokenizer():
|
|
rewriter = _load_orpo_rewriter("orpo_trainer_processor_pad_token")
|
|
source = """
|
|
def __init__(self, processing_class):
|
|
data_collator = DPODataCollatorWithPadding(
|
|
pad_token_id=processing_class.pad_token_id,
|
|
)
|
|
self.padding_value = processing_class.pad_token_id
|
|
"""
|
|
|
|
rewritten = rewriter("__init__", source)
|
|
|
|
assert "processing_class.pad_token_id" not in rewritten
|
|
assert "getattr(processing_class, 'pad_token_id'" in rewritten
|
|
|
|
class _Processor:
|
|
# No pad_token_id at the processor level; only on the inner tokenizer.
|
|
class tokenizer:
|
|
pad_token_id = 17
|
|
|
|
captured = {}
|
|
|
|
def DPODataCollatorWithPadding(**kwargs):
|
|
captured["pad_token_id"] = kwargs["pad_token_id"]
|
|
return object()
|
|
|
|
ns = {"DPODataCollatorWithPadding": DPODataCollatorWithPadding}
|
|
exec(rewritten, ns)
|
|
|
|
class _Trainer:
|
|
pass
|
|
|
|
trainer = _Trainer()
|
|
ns["__init__"](trainer, _Processor())
|
|
|
|
assert captured["pad_token_id"] == 17
|
|
assert trainer.padding_value == 17
|
|
|
|
|
|
def test_orpo_init_pad_token_id_uses_processor_when_present():
|
|
rewriter = _load_orpo_rewriter("orpo_trainer_processor_pad_token")
|
|
source = """
|
|
def __init__(self, processing_class):
|
|
self.padding_value = processing_class.pad_token_id
|
|
"""
|
|
|
|
rewritten = rewriter("__init__", source)
|
|
|
|
class _Tokenizer:
|
|
# Inner tokenizer must NOT be consulted when the processor exposes
|
|
# pad_token_id itself.
|
|
pad_token_id = 999
|
|
|
|
class _Processor:
|
|
pad_token_id = 42
|
|
tokenizer = _Tokenizer()
|
|
|
|
ns = {}
|
|
exec(rewritten, ns)
|
|
|
|
class _Trainer:
|
|
pass
|
|
|
|
trainer = _Trainer()
|
|
ns["__init__"](trainer, _Processor())
|
|
assert trainer.padding_value == 42
|
|
|
|
|
|
def test_orpo_init_pad_token_id_noop_on_non_init():
|
|
rewriter = _load_orpo_rewriter("orpo_trainer_processor_pad_token")
|
|
source = "def tokenize_row(self):\n return processing_class.pad_token_id\n"
|
|
assert rewriter("tokenize_row", source) == source
|