Add support for pushing audio models to HF hub, and cleanup tiktoken so it can have hf-hub/local-dir based bpe files for Tik

Author: rwightman

src/open_clip/factory.py

@@ -931,10 +931,25 @@ def get_tokenizer(
     if text_config.get('tokenizer_type', '') == 'tiktoken':
         # tiktoken-based tokenizer for generative (GenLIP) models.
         encoding_name = text_config.get('tiktoken_name', 'cl100k_base')
+        bpe_path = text_config.get('tiktoken_bpe_path')
+        encoding_config_path = text_config.get('tiktoken_config_path')
+        if bpe_path:
+            if schema == 'local-dir':
+                bpe_path = str(local_dir_path / bpe_path)
+                if encoding_config_path:
+                    encoding_config_path = str(local_dir_path / encoding_config_path)
+            elif schema == 'hf-hub':
+                bpe_path = download_pretrained_from_hf(identifier, filename=bpe_path, cache_dir=cache_dir)
+                if encoding_config_path:
+                    encoding_config_path = download_pretrained_from_hf(
+                        identifier, filename=encoding_config_path, cache_dir=cache_dir,
+                    )
         _logger.info(f"Using TikTokenTokenizer with encoding: '{encoding_name}'")
         tokenizer = TikTokenTokenizer(
             encoding_name=encoding_name,
             context_length=context_length,
+            bpe_path=bpe_path,
+            encoding_config_path=encoding_config_path,
             **{k: v for k, v in tokenizer_kwargs.items() if k in ('add_bos', 'add_eos', 'clean')},
         )
 

src/open_clip/push_to_hf_hub.py

@@ -1,5 +1,6 @@
 import argparse
 import json
+from copy import deepcopy
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import Optional, Tuple, Union
@@ -29,27 +30,65 @@
 
 from .constants import HF_WEIGHTS_NAME, HF_SAFE_WEIGHTS_NAME, HF_CONFIG_NAME
 from .factory import create_model_from_pretrained, get_model_config, get_tokenizer
-from .tokenizer import HFTokenizer, SigLipTokenizer
+from .tokenizer import HFTokenizer, SigLipTokenizer, TikTokenTokenizer
+
+
+_HF_TOKENIZER_FILE_PATTERNS = (
+    'tokenizer.json',
+    'tokenizer_config.json',
+    'special_tokens_map.json',
+    'added_tokens.json',
+    'vocab.json',
+    'merges.txt',
+    'spiece.model',
+    'sentencepiece.bpe.model',
+)
+
+
+def _get_text_config(model_config: Optional[dict]) -> dict:
+    if not model_config:
+        return {}
+    return model_config.get('text_cfg') or model_config.get('multimodal_cfg') or {}
+
+
+def _uses_tiktoken(tokenizer, model_config: Optional[dict]) -> bool:
+    text_cfg = _get_text_config(model_config)
+    return isinstance(tokenizer, TikTokenTokenizer) or text_cfg.get('tokenizer_type') == 'tiktoken'
+
+
+def _with_tiktoken_assets(model_config: Optional[dict], tokenizer_assets: Optional[dict]) -> Optional[dict]:
+    if not model_config or not isinstance(tokenizer_assets, dict):
+        return model_config
+    model_config = deepcopy(model_config)
+    text_cfg = model_config.get('text_cfg') or model_config.get('multimodal_cfg')
+    if text_cfg is not None:
+        text_cfg.update({
+            k: v for k, v in tokenizer_assets.items()
+            if k in ('tiktoken_bpe_path', 'tiktoken_config_path') and v
+        })
+    return model_config
 
 
 def save_config_for_hf(
         model,
         config_path: str,
         model_config: Optional[dict],
 ):
-    preprocess_cfg = {
-        'mean': model.visual.image_mean,
-        'std': model.visual.image_std,
-    }
-    other_pp = getattr(model.visual, 'preprocess_cfg', {})
-    if 'interpolation' in other_pp:
-        preprocess_cfg['interpolation'] = other_pp['interpolation']
-    if 'resize_mode' in other_pp:
-        preprocess_cfg['resize_mode'] = other_pp['resize_mode']
     hf_config = {
         'model_cfg': model_config,
-        'preprocess_cfg': preprocess_cfg,
     }
+    visual = getattr(model, 'visual', None)
+    if visual is not None:
+        preprocess_cfg = {
+            'mean': visual.image_mean,
+            'std': visual.image_std,
+        }
+        other_pp = getattr(visual, 'preprocess_cfg', {})
+        if 'interpolation' in other_pp:
+            preprocess_cfg['interpolation'] = other_pp['interpolation']
+        if 'resize_mode' in other_pp:
+            preprocess_cfg['resize_mode'] = other_pp['resize_mode']
+        hf_config['preprocess_cfg'] = preprocess_cfg
 
     with config_path.open('w') as f:
         json.dump(hf_config, f, indent=2)
@@ -76,10 +115,17 @@ def save_for_hf(
         if safe_serialization is False or safe_serialization == "both":
             torch.save(tensors, save_directory / HF_WEIGHTS_NAME)
 
-    tokenizer.save_pretrained(save_directory)
+    tokenizer_assets = None
+    save_pretrained = getattr(tokenizer, 'save_pretrained', None)
+    if callable(save_pretrained):
+        tokenizer_assets = save_pretrained(save_directory)
 
     config_path = save_directory / config_filename
-    save_config_for_hf(model, config_path, model_config=model_config)
+    save_config_for_hf(
+        model,
+        config_path,
+        model_config=_with_tiktoken_assets(model_config, tokenizer_assets),
+    )
 
 
 def push_to_hf_hub(
@@ -95,7 +141,9 @@ def push_to_hf_hub(
     model_card: Optional[dict] = None,
     safe_serialization: Union[bool, str] = 'both',
 ):
-    if not isinstance(tokenizer, (HFTokenizer, SigLipTokenizer)):
+    uses_tiktoken = _uses_tiktoken(tokenizer, model_config)
+    delete_patterns = _HF_TOKENIZER_FILE_PATTERNS if uses_tiktoken else None
+    if not isinstance(tokenizer, (HFTokenizer, SigLipTokenizer)) and not uses_tiktoken:
         # FIXME this makes it awkward to push models with new tokenizers, come up with better soln.
         # default CLIP tokenizers use https://huggingface.co/openai/clip-vit-large-patch14
         tokenizer = HFTokenizer('openai/clip-vit-large-patch14')
@@ -150,6 +198,7 @@ def push_to_hf_hub(
             revision=revision,
             create_pr=create_pr,
             commit_message=commit_message,
+            delete_patterns=delete_patterns,
         )
 
 

src/open_clip/tokenizer.py

@@ -4,10 +4,13 @@
 """
 import gzip
 import html
+import base64
+import json
 import os
 import random
 import string
 from functools import lru_cache, partial
+from pathlib import Path
 from typing import Callable, Dict, List, Optional, Tuple, Union
 import warnings
 
@@ -760,14 +763,27 @@ def __init__(
             add_bos: bool = True,
             add_eos: bool = True,
             clean: Optional[str] = None,
+            bpe_path: Optional[Union[str, os.PathLike]] = None,
+            encoding_config_path: Optional[Union[str, os.PathLike]] = None,
     ):
         try:
             import tiktoken
         except ImportError as e:
             raise ImportError("Please install tiktoken to use TikTokenTokenizer (`pip install tiktoken`).") from e
 
         self.encoding_name = encoding_name
-        self.enc = tiktoken.get_encoding(encoding_name)
+        if bpe_path is not None:
+            cfg = self._load_encoding_config(encoding_name, encoding_config_path)
+            self.enc = tiktoken.Encoding(
+                cfg.get("name", encoding_name),
+                pat_str=cfg["pat_str"],
+                mergeable_ranks=self._load_tiktoken_bpe(bpe_path),
+                special_tokens={str(k): int(v) for k, v in cfg.get("special_tokens", {}).items()},
+                explicit_n_vocab=cfg.get("explicit_n_vocab"),
+            )
+        else:
+            self.enc = tiktoken.get_encoding(encoding_name)
+        self.encoding_name = self.enc.name
         self.context_length = context_length
         self.add_bos = add_bos
         self.add_eos = add_eos
@@ -787,6 +803,59 @@ def __init__(
         self.all_special_ids = [self.eot_token_id, self.pad_token_id, self.bos_token_id]
         self.vocab_size = base + 3
 
+    @staticmethod
+    def _load_tiktoken_bpe(path: Union[str, os.PathLike]) -> Dict[bytes, int]:
+        ranks = {}
+        with open(path, "rb") as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                token, rank = line.split()
+                ranks[base64.b64decode(token)] = int(rank)
+        return ranks
+
+    @staticmethod
+    def _dump_tiktoken_bpe(ranks: Dict[bytes, int], path: Union[str, os.PathLike]) -> None:
+        with open(path, "wb") as f:
+            for token, rank in sorted(ranks.items(), key=lambda x: x[1]):
+                f.write(base64.b64encode(token) + b" " + str(rank).encode() + b"\n")
+
+    @staticmethod
+    def _load_encoding_config(encoding_name: str, path: Optional[Union[str, os.PathLike]]) -> Dict[str, object]:
+        if path is None:
+            raise ValueError(
+                f"encoding_config_path is required when loading tiktoken encoding {encoding_name!r} from bpe_path."
+            )
+        with open(path, "r", encoding="utf-8") as f:
+            return json.load(f)
+
+    def save_pretrained(self, dest: Union[str, os.PathLike]) -> Dict[str, str]:
+        dest = Path(dest)
+        asset_dir = dest / "open_clip_tiktoken"
+        asset_dir.mkdir(parents=True, exist_ok=True)
+
+        stem = self.encoding_name.replace("/", "_")
+        bpe_path = asset_dir / f"{stem}.tiktoken"
+        config_path = asset_dir / f"{stem}.json"
+
+        self._dump_tiktoken_bpe(self.enc._mergeable_ranks, bpe_path)
+        with config_path.open("w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    "name": self.enc.name,
+                    "pat_str": self.enc._pat_str,
+                    "special_tokens": self.enc._special_tokens,
+                    "explicit_n_vocab": self.enc.n_vocab,
+                },
+                f,
+                indent=2,
+            )
+
+        return {
+            "tiktoken_bpe_path": str(bpe_path.relative_to(dest)),
+            "tiktoken_config_path": str(config_path.relative_to(dest)),
+        }
+
     def encode(self, text: str) -> List[int]:
         # encode_ordinary ignores any special-token markup in the text, treating it as plain bytes.
         if self.clean_fn is not None: