Fix ORPO text-only tokenization with processors (#5501)

* Fix ORPO text tokenization with processors

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Guard ORPO tokenizer rewrite anchor

* Resolve processor pad_token_id and preserve preference data collators for ORPO

Two follow-ups so the text-only ORPO + VL processor path works end to end on
top of the build_tokenized_answer and tokenize_row rewrites:

1. Add orpo_trainer_processor_pad_token to rewrite processing_class.pad_token_id
   in ORPOTrainer.__init__ to fall back to processing_class.tokenizer.pad_token_id
   when the processor itself has no pad_token_id (Qwen3-VL, Gemma-3, etc.).
   Without this, DPODataCollatorWithPadding(pad_token_id=processing_class.pad_token_id)
   raises AttributeError before training starts.

2. Stop the outer UnslothORPOTrainer.__init__ collator-swap from clobbering
   DPODataCollatorWithPadding when the tokenizer is a processor without .pad.
   The swap to TransformersDataCollatorForLanguageModeling is now only applied
   to LM-style collators, so ORPO/DPO/CPO/KTO keep their own prompt/chosen/
   rejected handling. Otherwise the collator can't pad ORPO rows and raises
   "You should supply an encoding ... that includes input_ids" at train time.

Verified with Qwen3-VL-2B-Instruct ORPO + text-only data (training completes
to max_steps, no AttributeError, no collator error) and Llama-3.2-1B-Instruct
ORPO (losses and grad-norms bit-exact identical to main, so the change is a
true no-op for plain text tokenizers).

Extends tests/python/test_orpo_processor_text_tokenizer.py with three new
unit tests covering the pad_token_id rewriter.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

tests/python/test_orpo_processor_text_tokenizer.py

@@ -0,0 +1,238 @@
+"""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