[levanter] Add MarinTokenizer.as_hf_tokenizer(), fix kitoken find-links (#4451)

with_tokenizer_padded_to_match_model() crashed when given a
MarinTokenizer
because it called add_tokens(), an HF-only API. Add as_hf_tokenizer() to
the MarinTokenizer Protocol so any backend can produce an HF tokenizer
on
demand. Update padding, save_pretrained, and LoRA export to use it.

Also add find-links to levanter's pyproject.toml so uv can resolve
kitoken>=0.10.2 without the workspace root.

Author: rjpower

lib/levanter/pyproject.toml

@@ -150,6 +150,9 @@ dev = [
 
 
 [tool.uv]
+find-links = [
+  "https://github.com/marin-community/kitoken/releases/expanded_assets/kitoken-0.10.2-a3012f4",
+]
 conflicts = [
   [
     { extra = "gpu" },

lib/levanter/src/levanter/compat/hf_checkpoints.py

@@ -515,14 +515,15 @@ def with_tokenizer_padded_to_match_model(
             f"(adding {num_to_add} dummy tokens) to match model vocab size."
         )
 
-        # Add dummy tokens to the tokenizer
+        # Add dummy tokens to the tokenizer. MarinTokenizer is read-only,
+        # so we convert to an HF tokenizer which supports add_tokens.
         dummy_tokens = [f"<|padding_{i}|>" for i in range(num_to_add)]
-        self.tokenizer.add_tokens(dummy_tokens)
+        tokenizer = self.tokenizer
+        if isinstance(tokenizer, MarinTokenizer):
+            tokenizer = tokenizer.as_hf_tokenizer()
+        tokenizer.add_tokens(dummy_tokens)
 
-        # Return a new converter with the modified tokenizer
-        # Note: We modify self.tokenizer in place, but since the Vocab property is cached,
-        # we need to return a new converter to get a fresh Vocab
-        return dataclasses.replace(self, tokenizer=self.tokenizer)  # type: ignore
+        return dataclasses.replace(self, tokenizer=tokenizer)  # type: ignore
 
     def with_config_overrides(self, config_overrides: dict, merge: bool = True) -> "HFCheckpointConverter":
         if self.config_overrides is not None and merge:
@@ -1098,9 +1099,7 @@ def _list_relative_files(directory: str) -> set[str]:
                 logger.info("Saving tokenizer")
                 tokenizer = self.tokenizer
                 if isinstance(tokenizer, MarinTokenizer):
-                    # MarinTokenizer doesn't have save_pretrained; load the
-                    # underlying HF tokenizer so we can serialize it.
-                    tokenizer = load_tokenizer(tokenizer.name_or_path)
+                    tokenizer = tokenizer.as_hf_tokenizer()
                 tokenizer.save_pretrained(local_path)
 
             if save_feature_extractor and self.feature_extractor is not None:

lib/levanter/src/levanter/data/passthrough_tokenizer.py

@@ -95,3 +95,6 @@ def apply_chat_template_with_masks(
         **kwargs,
     ) -> dict[str, list[list[int]]]:
         raise ValueError("PassthroughTokenizer does not support chat templates")
+
+    def as_hf_tokenizer(self):
+        raise ValueError("PassthroughTokenizer cannot be converted to an HF tokenizer")

lib/levanter/src/levanter/lora.py

@@ -77,6 +77,8 @@
 silence_transformer_nag()
 from transformers import PreTrainedTokenizerBase  # noqa: E402
 
+from levanter.tokenizers import MarinTokenizer  # noqa: E402
+
 
 logger = logging.getLogger(__name__)
 
@@ -320,7 +322,7 @@ def save_peft_pretrained(
     config: LoraConfig,
     base_model_name_or_path,
     path: str,
-    tokenizer: Optional[PreTrainedTokenizerBase] = None,
+    tokenizer: Optional[PreTrainedTokenizerBase | MarinTokenizer] = None,
     *,
     prefix: Optional[str] = DEFAULT_DICT_PREFIX,
     upload_to: Optional[Union[bool, str, RepoRef]] = None,
@@ -349,6 +351,8 @@ def save_peft_pretrained(
             json.dump(hf_config, f)
 
         if tokenizer is not None:
+            if isinstance(tokenizer, MarinTokenizer):
+                tokenizer = tokenizer.as_hf_tokenizer()
             tokenizer.save_pretrained(local_path)
 
         if upload_to is True:
@@ -363,7 +367,7 @@ def save_peft_checkpoint_callback(
     base_path,
     config: LoraConfig,
     base_model_name_or_path,
-    tokenizer: Optional[PreTrainedTokenizerBase] = None,
+    tokenizer: Optional[PreTrainedTokenizerBase | MarinTokenizer] = None,
     upload_to_hf: Optional[Union[bool, str, RepoRef]] = False,
     **hf_upload_kwargs,
 ):

lib/levanter/src/levanter/tokenizers.py

@@ -89,6 +89,14 @@ def apply_chat_template_with_masks(
         **kwargs,
     ) -> dict[str, list[list[int]]]: ...
 
+    def as_hf_tokenizer(self) -> Any:
+        """Return a HuggingFace PreTrainedTokenizerFast for this tokenizer.
+
+        Useful for operations that require the HF API (save_pretrained,
+        add_tokens, generation config, etc.).
+        """
+        ...
+
 
 # Sentinel used to mark generation (assistant) boundaries in rendered templates.
 _GENERATION_SENTINEL_START = "__MARIN_GEN_START_7f3a9c__"
@@ -315,6 +323,11 @@ def apply_chat_template_with_masks(
     ) -> dict[str, list[list[int]]]:
         return _apply_chat_template_with_masks(self, conversations, chat_template=chat_template, **kwargs)
 
+    def as_hf_tokenizer(self) -> Any:
+        from transformers import AutoTokenizer
+
+        return AutoTokenizer.from_pretrained(self._name_or_path, trust_remote_code=True)
+
 
 @dataclasses.dataclass(frozen=True)
 class KitokenMarinTokenizer:
@@ -449,6 +462,11 @@ def apply_chat_template_with_masks(
     ) -> dict[str, list[list[int]]]:
         return _apply_chat_template_with_masks(self, conversations, chat_template=chat_template, **kwargs)
 
+    def as_hf_tokenizer(self) -> Any:
+        from transformers import AutoTokenizer
+
+        return AutoTokenizer.from_pretrained(self._name_or_path, trust_remote_code=True)
+
 
 class TokenizerBackend(StrEnum):
     HF = "hf"