Accept processor in get_training_chat_template

trl/chat_template_utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from pathlib import Path
 from typing import TypeVar
 
@@ -552,18 +553,21 @@ def is_chat_template_prefix_preserving(processing_class: PreTrainedTokenizerBase
 qwen3_6_training_chat_template = (_CHAT_TEMPLATES_DIR / "qwen3_6_training.jinja").read_text()
 
 
-def get_training_chat_template(tokenizer: PreTrainedTokenizerBase) -> str | None:
+def get_training_chat_template(
+    processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None,
+    tokenizer: PreTrainedTokenizerBase | None = None,
+) -> str | None:
     r"""
     Get a training-compatible chat template, if needed.
 
     Returns a patched chat template that is prefix-preserving and includes `{%% generation %%}` / `{%% endgeneration
-    %%}` markers for assistant-only loss masking. Returns `None` if the tokenizer's template already satisfies both
-    requirements. Currently Cohere, DeepSeek-V3, Gemma, Gemma2, GLM-4-MoE, GPT-OSS, LLaMA 3, Phi-3, Qwen2.5, Qwen3, and
-    Qwen3.6 are supported.
+    %%}` markers for assistant-only loss masking. Returns `None` if the template already satisfies both requirements.
+    Currently Cohere, DeepSeek-V3, Gemma, Gemma2, GLM-4-MoE, GPT-OSS, LLaMA 3, Phi-3, Qwen2.5, Qwen3, and Qwen3.6 are
+    supported.
 
     Args:
-        tokenizer (`PreTrainedTokenizerBase`):
-            Tokenizer instance to check.
+        processing_class (`PreTrainedTokenizerBase` or `ProcessorMixin`):
+            Tokenizer or processor instance to check.
 
     Returns:
         `str` or `None`:
@@ -604,46 +608,60 @@ def get_training_chat_template(tokenizer: PreTrainedTokenizerBase) -> str | None
     '<|im_start|>user\nWhat is 2 * 3?<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n<tool_call>\n{"name": "multiply", "arguments": {"a": 2, "b": 3}}\n</tool_call><|im_end|>\n<|im_start|>user\n<tool_response>\n6\n</tool_response><|im_end|>\n<|im_start|>assistant\n'
     ```
     """
+    if tokenizer is not None:
+        if processing_class is not None:
+            raise TypeError(
+                "Pass only `processing_class`; `tokenizer` is a deprecated alias for backward compatibility."
+            )
+        warnings.warn(
+            "The `tokenizer` argument of `get_training_chat_template` is deprecated and will be removed in TRL 2.0. "
+            "Use `processing_class` instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        processing_class = tokenizer
+    if processing_class is None:
+        raise TypeError("get_training_chat_template() missing required argument: 'processing_class'")
+
     # First check if patching is needed. Prefix-preservation only matters when the template actually supports tools
     # (the check itself renders a tool message), so skip it otherwise.
-    prefix_ok = not supports_tool_calling(tokenizer) or is_chat_template_prefix_preserving(tokenizer)
-    if prefix_ok and "{% generation %}" in tokenizer.chat_template:
+    prefix_ok = not supports_tool_calling(processing_class) or is_chat_template_prefix_preserving(processing_class)
+    if prefix_ok and "{% generation %}" in processing_class.chat_template:
         return None  # No patching needed
 
-    if tokenizer.chat_template == cohere_chat_template:
+    if processing_class.chat_template == cohere_chat_template:
         return cohere_training_chat_template
 
-    if tokenizer.chat_template == deepseekv3_chat_template:
+    if processing_class.chat_template == deepseekv3_chat_template:
         return deepseekv3_training_chat_template
 
-    if tokenizer.chat_template == gemma_chat_template:
+    if processing_class.chat_template == gemma_chat_template:
         return gemma_training_chat_template
 
-    if tokenizer.chat_template == glm4moe_chat_template:
+    if processing_class.chat_template == glm4moe_chat_template:
         return glm4moe_training_chat_template
 
-    if tokenizer.chat_template == gptoss_chat_template:
+    if processing_class.chat_template == gptoss_chat_template:
         return gptoss_training_chat_template
 
-    if tokenizer.chat_template == llama3_chat_template:
+    if processing_class.chat_template == llama3_chat_template:
         return llama3_training_chat_template
 
-    if tokenizer.chat_template == phi3_chat_template:
+    if processing_class.chat_template == phi3_chat_template:
         return phi3_training_chat_template
 
-    if tokenizer.chat_template == qwen2_5_chat_template:
+    if processing_class.chat_template == qwen2_5_chat_template:
         return qwen2_5_training_chat_template
 
-    if tokenizer.chat_template == qwen3_chat_template:
+    if processing_class.chat_template == qwen3_chat_template:
         return qwen3_training_chat_template
 
-    if tokenizer.chat_template == qwen3_6_chat_template:
+    if processing_class.chat_template == qwen3_6_chat_template:
         return qwen3_6_training_chat_template
 
     raise ValueError(
-        "The tokenizer's chat template is not training-compatible (missing prefix-preservation or "
-        "`{% generation %}` markers) and patching is not supported for this template. "
-        "Please manually modify the tokenizer's chat template for training."
+        "The chat template is not training-compatible (missing prefix-preservation or `{% generation %}` markers) "
+        "and patching is not supported for this template. Please manually modify the chat template for training."
     )