Centralize tokenizer defaults and preserve sampled inference history

dev/trainer_rank.py

@@ -5,9 +5,9 @@
 import torch
 import torch.distributed as dist
 from trainer_rank_support import load_random_checkpoints
-from transformers import AutoTokenizer
 import typer
 
+from art import get_tokenizer
 from art.trainer_rank import AdamParams, ForwardInput, TrainerRank
 
 
@@ -34,9 +34,7 @@ def main(
 
         from art.megatron import train as megatron_train
 
-        tokenizer = cast(
-            Any, AutoTokenizer.from_pretrained(model, trust_remote_code=True)
-        )
+        tokenizer = cast(Any, get_tokenizer(model, trust_remote_code=True))
         inputs: list[ForwardInput[torch.Tensor, None, None, None]] = []
         rows = load_dataset("roneneldan/TinyStories", split="train", streaming=True)
         for row in islice(rows, samples):

dev/trainer_rank_landing_acceptance.py

@@ -197,13 +197,11 @@ def _check_corpus_tokenizer(corpus: dict[str, Any], model: str) -> dict[str, Any
     compares vocabulary sizes; it fails the cell loudly on any difference.
     """
 
-    from transformers import AutoTokenizer
+    from art import get_tokenizer
 
     corpus_model = str(corpus.get("tokenizer_model") or "")
-    model_tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
-    corpus_tokenizer = AutoTokenizer.from_pretrained(
-        corpus_model, trust_remote_code=True
-    )
+    model_tokenizer = get_tokenizer(model, trust_remote_code=True)
+    corpus_tokenizer = get_tokenizer(corpus_model, trust_remote_code=True)
     sample = corpus["groups"][0]["histories"][0]["tokens"][:256]
     problems = []
     if len(model_tokenizer) != len(corpus_tokenizer):

examples/hn_title_generator/train.py

@@ -8,10 +8,10 @@
 import openai
 from openai.types.chat import ChatCompletionMessageParam
 from openpipe import AsyncOpenPipe
-from transformers.models.auto.tokenization_auto import AutoTokenizer
 from utils import cache, prompt_for_title, pull_data, score_title
 
 import art
+from art import get_tokenizer
 from art.local import LocalBackend
 from art.utils import iterate_dataset, limit_concurrency
 
@@ -37,7 +37,7 @@ def filter_on_length(data: Dataset, max_length: int, tokenizer_name: str) -> Dat
     print(
         f"Filtering dataset for max prompt length: {max_length} using tokenizer: {tokenizer_name}"
     )
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    tokenizer = get_tokenizer(tokenizer_name)
 
     def check_length(x):
         # Ensure 'prompt' is a list of dicts

pyproject.toml

@@ -152,7 +152,7 @@ build-backend = "hatchling.build"
 allow-direct-references = true
 
 [tool.hatch.build.targets.wheel]
-packages = ["src/art", "src/mp_actors"]
+packages = ["src/art", "src/art_inference", "src/mp_actors"]
 
 [tool.hatch.build.targets.wheel.force-include]
 ".agents/skills" = "art/skills"

src/art/__init__.py

@@ -71,6 +71,7 @@
 from .model import Model, TrainableModel
 from .pipeline_tuner import PipelineAutotuneConfig, PipelineRuntimeConfig
 from .serverless import ServerlessBackend
+from .tokenizer import get_tokenizer
 from .trajectories import (
     Trajectory,
     TrajectoryGroup,
@@ -116,6 +117,7 @@
     "PIPELINE_RL_METRIC_DEFINITIONS",
     "PIPELINE_RL_SCORE_METRICS",
     "get_megatron_runtime_config",
+    "get_tokenizer",
     "init_megatron_runtime_config",
     "ServerlessBackend",
     "ServerlessTrainResult",

src/art/local/backend.py

@@ -16,9 +16,9 @@
 from typing import TYPE_CHECKING, Any, AsyncIterator, Iterable, Literal, cast
 import warnings
 
+from art.tokenizer import get_tokenizer
 from art.utils.chat_template import (
     chat_template_with_preserved_thinking,
-    configure_preserved_thinking_chat_template,
 )
 from art.utils.lifecycle import (
     PROCESS_SHUTDOWN_TIMEOUT_SECONDS,
@@ -40,7 +40,6 @@
 from pydantic import BaseModel, ConfigDict
 import torch
 from tqdm import auto as tqdm
-from transformers import AutoTokenizer
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 from typing_extensions import Self
 
@@ -196,7 +195,9 @@ def _apply_configured_chat_template(
 ) -> None:
     chat_template = _configured_chat_template_value(internal_config)
     if chat_template is not None:
-        tokenizer.chat_template = chat_template
+        tokenizer.chat_template = cast(
+            str, chat_template_with_preserved_thinking(chat_template)
+        )
 
 
 def _model_support_handler(
@@ -241,9 +242,9 @@ def _apply_configured_chat_template_server_args(
         chat_template = _model_support_default_chat_template(
             base_model, internal_config
         )
-    if chat_template is None and _should_probe_preserve_thinking_template(base_model):
+    if chat_template is None and base_model is not None:
         try:
-            tokenizer = AutoTokenizer.from_pretrained(base_model)
+            tokenizer = get_tokenizer(base_model)
         except (OSError, ValueError) as error:
             warnings.warn(
                 f"Could not load {base_model!r} to configure prior-thinking "
@@ -252,13 +253,14 @@ def _apply_configured_chat_template_server_args(
                 stacklevel=2,
             )
         else:
-            default = getattr(tokenizer, "chat_template", None)
-            preserved = chat_template_with_preserved_thinking(default)
-            if preserved != default:
-                chat_template = cast(str, preserved)
+            template = getattr(tokenizer, "chat_template", None)
+            if isinstance(template, str) and ("{{" in template or "{%" in template):
+                chat_template = template
     if chat_template is None:
         return
-    server_args.setdefault("chat_template", chat_template)
+    server_args.setdefault(
+        "chat_template", chat_template_with_preserved_thinking(chat_template)
+    )
     if chat_template_content_format := internal_config.get(
         "chat_template_content_format"
     ):
@@ -269,13 +271,6 @@ def _apply_configured_chat_template_server_args(
     config_dict["server_args"] = server_args
 
 
-def _should_probe_preserve_thinking_template(base_model: str | None) -> bool:
-    if base_model is None:
-        return False
-    model_name = base_model.rstrip("/").rsplit("/", 1)[-1]
-    return model_name.startswith(("Qwen3-", "Qwen3.5-"))
-
-
 def _tokenizer_cache_key(
     base_model: str,
     internal_config: dev.InternalModelConfig,
@@ -287,12 +282,7 @@ def _tokenizer_cache_key(
 
 
 def _load_training_tokenizer(base_model: str) -> PreTrainedTokenizerBase:
-    return cast(
-        PreTrainedTokenizerBase,
-        configure_preserved_thinking_chat_template(
-            AutoTokenizer.from_pretrained(base_model)
-        ),
-    )
+    return get_tokenizer(base_model)
 
 
 class LocalBackend:

src/art/megatron/dsv4/tokenizer.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
 
+from contextvars import ContextVar
 import copy
 from typing import Any
 
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
-from art.megatron.dsv4.encoding import encode_messages
+from art.megatron.dsv4 import encoding
+from art_inference.append_only import patch_deepseek_renderer, preserves_history
+
+_PRESERVE_HISTORY: ContextVar[bool] = ContextVar(
+    "art_dsv4_preserve_history", default=True
+)
+patch_deepseek_renderer(encoding, _PRESERVE_HISTORY.get)
 
 DSV4_CHAT_TEMPLATE_MARKER = "deepseek_v4_python_encoder enable_thinking"
 
@@ -35,12 +42,28 @@ def apply_chat_template(
             tools: list[dict[str, Any]] | None = None,
             **kwargs: Any,
         ) -> str | list[int]:
+            chat_template = kwargs.get("chat_template")
+            if chat_template is None:
+                chat_template = self.chat_template
+            if chat_template != DSV4_CHAT_TEMPLATE_MARKER:
+                return super().apply_chat_template(messages, tools=tools, **kwargs)
             thinking = bool(kwargs.get("thinking", False)) or bool(
                 kwargs.get("enable_thinking", False)
             )
             thinking_mode = "thinking" if thinking else "chat"
             conversation = kwargs.get("conversation", messages)
-            rendered_messages = list(conversation)
+            rendered_messages = [
+                {
+                    **message,
+                    "reasoning": message.get(
+                        "reasoning",
+                        message.get("reasoning_content", message.get("thinking")),
+                    ),
+                }
+                if message.get("role") == "assistant"
+                else message
+                for message in conversation
+            ]
             if tools:
                 rendered_messages.insert(0, {"role": "system", "tools": tools})
 
@@ -55,12 +78,17 @@ def apply_chat_template(
             else:
                 reasoning_effort = "high"
 
-            prompt = encode_messages(
-                rendered_messages,
-                thinking_mode=thinking_mode,
-                drop_thinking=kwargs.get("drop_thinking", True),
-                reasoning_effort=reasoning_effort,
-            )
+            preserve = preserves_history(kwargs)
+            token = _PRESERVE_HISTORY.set(preserve)
+            try:
+                prompt = encoding.encode_messages(
+                    rendered_messages,
+                    thinking_mode=thinking_mode,
+                    drop_thinking=not preserve,
+                    reasoning_effort=reasoning_effort,
+                )
+            finally:
+                _PRESERVE_HISTORY.reset(token)
             if not kwargs.get("tokenize", True):
                 return prompt
             tokenizer_kwargs = {