diff --git a/README.md b/README.md index f03d59b7c..f983f0baf 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,7 @@ > - CoCa's exact-mask API adds `text_valid` after `text` in `encode_text`, `forward`, and `forward_intermediates`. This shifts the older trailing positional arguments (`normalize`, `image_latent`, `image_indices`, and so on); pass those arguments by keyword. For example, replace `model.encode_text(text, False)` with `model.encode_text(text, normalize=False)`. > - `CLIPTextCfg.eos_id` no longer defaults to `2` (that value is only correct for XLM-style vocabs). Configs using `pool_type="eos"` must set `eos_id` explicitly, and `get_tokenizer` now validates `eos_id`/`pad_id` against the resolved tokenizer, raising on mismatch instead of pooling/masking silently wrong positions. > - `HFTokenizer` no longer fabricates `pad_token_id=0` when the underlying tokenizer has no pad token (id 0 is a real token in most BPE vocabs); variable-text setups fail fast instead. It also forces `padding_side='right'`, which all OpenCLIP pooling/masking assumes. +> - Tokenizer wrappers now share special-token controls: `encode(..., add_special_tokens=False)` remains body-only by default, while model-facing `tokenizer(...)` defaults to `add_special_tokens=True`. `decode()` / `batch_decode()` default to `skip_special_tokens=False, stop_at_eos=True`; pass `stop_at_eos=False` to inspect tokens after the first EOS. This intentionally changes legacy SimpleTokenizer decode output by hiding post-EOS id-0 fill (`!`) and makes TikToken decode render its reserved control tokens unless `skip_special_tokens=True`. > - `MaxPooler` (`hf_pooler_type="max_pooler"`) mask polarity fixed — it previously max-pooled over the padding positions instead of the valid ones. > - `CoCa.__init__` no longer takes a `pad_id` argument — `model.pad_id` is derived from the text tower (the id it actually masks with: `text_cfg.pad_id` for native towers, the transformers config pad for HF towers). MaMMUT follows the same pattern. This fixes `coca_roberta-*`, which previously masked with roberta's pad (1) in the tower while the loss ignored 0 — its config now declares `pad_id: 1` and the caption loss no longer trains on padding. > - `CoCaTask` builds caption labels masked to `-100`; `CoCaLoss`'s cross-entropy uses `ignore_index=-100`. Its `pad_id` arg is retained for standalone callers passing raw labels (default `0` preserves the old value-based behavior; the task path passes `None`). Validation generative-loss metrics are likewise mask/pad-aware and will report different (correct) values for nonzero-pad models. diff --git a/src/open_clip/__init__.py b/src/open_clip/__init__.py index d92ab8497..1384905ff 100644 --- a/src/open_clip/__init__.py +++ b/src/open_clip/__init__.py @@ -52,7 +52,7 @@ download_pretrained, ) from .push_to_hf_hub import push_pretrained_to_hf_hub, push_to_hf_hub -from .tokenizer import SimpleTokenizer, tokenize, decode +from .tokenizer import Tokenizer, SimpleTokenizer, tokenize, decode, batch_decode from .transform import image_transform, AugmentationCfg from .zero_shot_classifier import build_zero_shot_classifier, build_zero_shot_classifier_legacy from .zero_shot_metadata import OPENAI_IMAGENET_TEMPLATES, SIMPLE_IMAGENET_TEMPLATES, IMAGENET_CLASSNAMES diff --git a/src/open_clip/factory.py b/src/open_clip/factory.py index 3ba9cd4e9..814b4b28a 100644 --- a/src/open_clip/factory.py +++ b/src/open_clip/factory.py @@ -33,7 +33,7 @@ merge_preprocess_kwargs, naflex_eval_transform_v2, ) -from .tokenizer import HFTokenizer, SimpleTokenizer, SigLipTokenizer, TikTokenTokenizer, DEFAULT_CONTEXT_LENGTH +from .tokenizer import HFTokenizer, SimpleTokenizer, SigLipTokenizer, TikTokenTokenizer, Tokenizer, DEFAULT_CONTEXT_LENGTH HF_HUB_PREFIX = 'hf-hub:' _MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"] @@ -835,7 +835,7 @@ def get_tokenizer( context_length: Optional[int] = None, cache_dir: Optional[str] = None, **kwargs, # Additional tokenizer kwargs passed to constructor -): +) -> Tokenizer: """ Gets the appropriate tokenizer based on the model identifier schema or name. diff --git a/src/open_clip/tokenizer.py b/src/open_clip/tokenizer.py index 4baf642e7..b79c01a7a 100644 --- a/src/open_clip/tokenizer.py +++ b/src/open_clip/tokenizer.py @@ -8,7 +8,7 @@ import random import string from functools import lru_cache, partial -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable, Iterable, List, Optional, Protocol, Sequence, Tuple, Union import warnings import ftfy @@ -22,6 +22,125 @@ DEFAULT_CONTEXT_LENGTH = 77 # default context length for OpenAI CLIP +TokenizerInput = Union[str, Sequence[str]] +TokenIds = Union[Sequence[int], np.ndarray, torch.Tensor] +BatchTokenIds = Union[Iterable[TokenIds], np.ndarray, torch.Tensor] +TokenizerOutput = Union[ + torch.Tensor, + List[torch.Tensor], + Tuple[torch.Tensor, torch.Tensor], +] + + +class Tokenizer(Protocol): + """Structural interface shared by OpenCLIP tokenizer implementations.""" + + context_length: Optional[int] + vocab_size: int + bos_token_id: Optional[int] + eos_token_id: Optional[int] + pad_token_id: Optional[int] + sot_token_id: Optional[int] + eot_token_id: Optional[int] + all_special_ids: List[int] + + def encode(self, text: str, add_special_tokens: bool = False) -> List[int]: ... + + def decode( + self, + tokens: TokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> str: ... + + def batch_decode( + self, + batch_tokens: BatchTokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> List[str]: ... + + def __call__( + self, + texts: TokenizerInput, + context_length: Optional[int] = None, + pad: bool = True, + output_mask: bool = False, + add_special_tokens: bool = True, + ) -> TokenizerOutput: ... + + +def _to_token_list(tokens: TokenIds) -> List[int]: + if isinstance(tokens, torch.Tensor): + tokens = tokens.detach().cpu().tolist() + elif isinstance(tokens, np.ndarray): + tokens = tokens.tolist() + return list(tokens) + + +def _to_token_batch(batch_tokens: BatchTokenIds) -> Iterable[TokenIds]: + if isinstance(batch_tokens, torch.Tensor): + return batch_tokens.detach().cpu().tolist() + if isinstance(batch_tokens, np.ndarray): + return batch_tokens.tolist() + return batch_tokens + + +def _truncate_at_eos(tokens: TokenIds, eos_token_id: Optional[int], stop_at_eos: bool) -> List[int]: + tokens = _to_token_list(tokens) + if stop_at_eos and eos_token_id is not None: + try: + tokens = tokens[:tokens.index(eos_token_id) + 1] + except ValueError: + pass + return tokens + + +def _decode_with_backend( + backend, + tokens: TokenIds, + eos_token_id: Optional[int], + skip_special_tokens: bool, + stop_at_eos: bool, +) -> str: + tokens = _truncate_at_eos(tokens, eos_token_id, stop_at_eos) + return backend.decode(tokens, skip_special_tokens=skip_special_tokens) + + +def _batch_decode_with_backend( + backend, + batch_tokens: BatchTokenIds, + eos_token_id: Optional[int], + skip_special_tokens: bool, + stop_at_eos: bool, +) -> List[str]: + batch_tokens = _to_token_batch(batch_tokens) + batch_tokens = [ + _truncate_at_eos(tokens, eos_token_id, stop_at_eos) + for tokens in batch_tokens + ] + return backend.batch_decode(batch_tokens, skip_special_tokens=skip_special_tokens) + + +def _get_pad_fill_id(pad_token_id: Optional[int]) -> int: + """Use the reserved pad id when present, otherwise preserve the historical id-0 fill.""" + return 0 if pad_token_id is None else pad_token_id + + +def _pad_token_sequences( + all_tokens: List[List[int]], + context_length: int, + pad_token_id: int = 0, + output_mask: bool = False, +) -> TokenizerOutput: + result = torch.full((len(all_tokens), context_length), pad_token_id, dtype=torch.long) + mask = torch.zeros_like(result, dtype=torch.bool) if output_mask else None + for i, tokens in enumerate(all_tokens): + result[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long) + if mask is not None: + mask[i, :len(tokens)] = True + return (result, mask) if mask is not None else result + @lru_cache() def default_bpe(): @@ -174,6 +293,9 @@ def __init__( self.all_special_ids = [self.encoder[t] for t in special_tokens] self.sot_token_id = self.all_special_ids[0] self.eot_token_id = self.all_special_ids[1] + self.bos_token_id = self.sot_token_id + self.eos_token_id = self.eot_token_id + self.pad_token_id = None self.context_length = context_length self.clean_fn = get_clean_fn(clean) self.reduction_fn = get_reduction_mask_fn(reduction_mask) if reduction_mask else None @@ -219,26 +341,49 @@ def bpe(self, token): self.cache[token] = word return word - def encode(self, text): + def encode(self, text: str, add_special_tokens: bool = False) -> List[int]: bpe_tokens = [] text = self.clean_fn(text) for token in re.findall(self.pat, text): token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8')) bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' ')) + if add_special_tokens: + bpe_tokens = [self.sot_token_id] + bpe_tokens + [self.eot_token_id] return bpe_tokens - def decode(self, tokens): + def decode( + self, + tokens: TokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> str: + tokens = _truncate_at_eos(tokens, self.eot_token_id, stop_at_eos) + if skip_special_tokens: + tokens = [token for token in tokens if token not in self.all_special_ids] text = ''.join([self.decoder[token] for token in tokens]) text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('', ' ') return text + def batch_decode( + self, + batch_tokens: BatchTokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> List[str]: + batch_tokens = _to_token_batch(batch_tokens) + return [ + self.decode(tokens, skip_special_tokens=skip_special_tokens, stop_at_eos=stop_at_eos) + for tokens in batch_tokens + ] + def __call__( self, - texts: Union[str, List[str]], + texts: TokenizerInput, context_length: Optional[int] = None, pad: bool = True, output_mask: bool = False, - ) -> Union[torch.LongTensor, Tuple[torch.LongTensor, torch.Tensor]]: + add_special_tokens: bool = True, + ) -> TokenizerOutput: """ Returns the tokenized representation of given input string(s) Parameters @@ -250,6 +395,8 @@ def __call__( output_mask : bool Also return a [B, L] bool attention mask (True = real token, HF polarity). Length-derived, so it stays exact even though this tokenizer pads with 0, a real vocab token. + add_special_tokens : bool + Add the start- and end-of-text tokens. Defaults to True for model-ready tokenization. Returns ------- @@ -277,50 +424,61 @@ def __call__( sot_token_id=self.sot_token_id, eot_token_id=self.eot_token_id, encode_fn=self.encode, + add_special_tokens=add_special_tokens, + output_mask=output_mask, ) - if output_mask: - # true lengths are not tracked through the reduction fns; positions through the first - # eot are valid by the right-padded contract (eot is a special id never emitted mid-text) - eot = result == self.eot_token_id - mask = eot.cumsum(dim=-1) - eot.long() == 0 - return result, mask return result - all_tokens = [[self.sot_token_id] + self.encode(text) + [self.eot_token_id] for text in texts] + all_tokens = [self.encode(text, add_special_tokens=add_special_tokens) for text in texts] truncated = [] for tokens in all_tokens: if len(tokens) > context_length: tokens = tokens[:context_length] # Truncate - tokens[-1] = self.eot_token_id + if add_special_tokens: + tokens[-1] = self.eot_token_id truncated.append(tokens) all_tokens = truncated - - result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) - - for i, tokens in enumerate(all_tokens): - result[i, :len(tokens)] = torch.tensor(tokens) - - if output_mask: - # exact length-based validity: this tokenizer has no reserved pad (fills with 0, a real - # vocab token), so a value-derived mask cannot distinguish pad fill from genuine id-0 tokens - mask = torch.zeros_like(result, dtype=torch.bool) - for i, tokens in enumerate(all_tokens): - mask[i, :len(tokens)] = True - return result, mask - - return result + # The length-derived mask remains exact even though id 0 is both fill and a valid body token. + return _pad_token_sequences(all_tokens, context_length, output_mask=output_mask) _tokenizer = SimpleTokenizer() -def decode(output_ids: torch.Tensor): - output_ids = output_ids.cpu().numpy() - return _tokenizer.decode(output_ids) - - -def tokenize(texts: Union[str, List[str]], context_length: int = DEFAULT_CONTEXT_LENGTH) -> torch.LongTensor: - return _tokenizer(texts, context_length=context_length) +def decode( + output_ids: TokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, +) -> str: + return _tokenizer.decode( + output_ids, + skip_special_tokens=skip_special_tokens, + stop_at_eos=stop_at_eos, + ) + + +def batch_decode( + output_ids: BatchTokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, +) -> List[str]: + return _tokenizer.batch_decode( + output_ids, + skip_special_tokens=skip_special_tokens, + stop_at_eos=stop_at_eos, + ) + + +def tokenize( + texts: TokenizerInput, + context_length: int = DEFAULT_CONTEXT_LENGTH, + add_special_tokens: bool = True, +) -> torch.LongTensor: + return _tokenizer( + texts, + context_length=context_length, + add_special_tokens=add_special_tokens, + ) def random_mask_tokenize( @@ -330,26 +488,29 @@ def random_mask_tokenize( eot_token_id: int, encode_fn: Callable, shuffle: bool = False, + add_special_tokens: bool = True, + output_mask: bool = False, ): all_tokens = [encode_fn(text) for text in texts] - result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + reduced_tokens = [] + num_special_tokens = 2 if add_special_tokens else 0 + num_keep = context_length - num_special_tokens - for i, tokens in enumerate(all_tokens): + for tokens in all_tokens: tokens = torch.tensor(tokens) num_tokens = len(tokens) - if num_tokens > context_length - 2: # 2 for sot and eot token - num_keep = context_length - 2 + if num_tokens > num_keep: indices = torch.randperm(len(tokens)) indices = indices[:num_keep] if not shuffle: indices = indices.msort() tokens = tokens[indices] - num_tokens = num_keep - result[i, 0] = sot_token_id - result[i, 1:num_tokens + 1] = tokens - result[i, num_tokens + 1] = eot_token_id + tokens = tokens.tolist() + if add_special_tokens: + tokens = [sot_token_id] + tokens + [eot_token_id] + reduced_tokens.append(tokens) - return result + return _pad_token_sequences(reduced_tokens, context_length, output_mask=output_mask) def simple_mask_tokenize( @@ -358,20 +519,24 @@ def simple_mask_tokenize( sot_token_id: int, eot_token_id: int, encode_fn: Callable, + add_special_tokens: bool = True, + output_mask: bool = False, ): all_tokens = [encode_fn(text) for text in texts] - result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + reduced_tokens = [] + num_special_tokens = 2 if add_special_tokens else 0 + num_keep = context_length - num_special_tokens - for i, tokens in enumerate(all_tokens): + for tokens in all_tokens: num_tokens = len(tokens) - if num_tokens > context_length - 2: # 2 for sot and eot token - num_keep = context_length - 2 + if num_tokens > num_keep: start_index = random.randint(0, num_tokens - num_keep) # high is incl tokens = tokens[start_index: start_index + num_keep] - tokens = [sot_token_id] + tokens + [eot_token_id] - result[i, :len(tokens)] = torch.tensor(tokens) + if add_special_tokens: + tokens = [sot_token_id] + tokens + [eot_token_id] + reduced_tokens.append(tokens) - return result + return _pad_token_sequences(reduced_tokens, context_length, output_mask=output_mask) def syntax_mask_tokenize( @@ -380,7 +545,9 @@ def syntax_mask_tokenize( sot_token_id: int, eot_token_id: int, encode_fn: Callable, -) -> torch.LongTensor: + add_special_tokens: bool = True, + output_mask: bool = False, +) -> Union[torch.LongTensor, Tuple[torch.LongTensor, torch.Tensor]]: """ Returns the tokenized representation of given input string(s). Apply syntax masking before tokenize. """ @@ -404,13 +571,14 @@ def get_order(x): # syntax masking new_texts = [] + num_special_tokens = 2 if add_special_tokens else 0 for text in texts: list_tokens = nltk.tokenize.word_tokenize(text) pos_tags = nltk.pos_tag(list_tokens) # sample the words by get_order method order_list = [get_order(tag) for _, tag in pos_tags] sorted_ids = np.argsort(np.array(order_list)) - sampled_ids = sorted(sorted_ids[:context_length - 2]) # need 2 slots for sot and eot tokens + sampled_ids = sorted(sorted_ids[:context_length - num_special_tokens]) sampled_tokens = np.take(np.array(list_tokens), sampled_ids, axis=0) # sample the tokens new_text = '' @@ -420,17 +588,20 @@ def get_order(x): new_texts.append(new_text) texts = new_texts - all_tokens = [[sot_token_id] + encode_fn(text) + [eot_token_id] for text in texts] - result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + all_tokens = [encode_fn(text) for text in texts] + truncated = [] - for i, tokens in enumerate(all_tokens): + for tokens in all_tokens: + if add_special_tokens: + tokens = [sot_token_id] + tokens + [eot_token_id] # still need first truncate because some words produces two tokens if len(tokens) > context_length: tokens = tokens[:context_length] # Truncate - tokens[-1] = eot_token_id - result[i, :len(tokens)] = torch.tensor(tokens) + if add_special_tokens: + tokens[-1] = eot_token_id + truncated.append(tokens) - return result + return _pad_token_sequences(truncated, context_length, output_mask=output_mask) def get_reduction_mask_fn(type: str): @@ -513,6 +684,9 @@ def __init__( self.sot_token_id = self.tokenizer.bos_token_id if self.sot_token_id is None: self.sot_token_id = self.tokenizer.cls_token_id + self.eos_token_id = self.eot_token_id + self.bos_token_id = self.sot_token_id + self.all_special_ids = self.tokenizer.all_special_ids self.vocab_size = len(self.tokenizer) # Set language function if available @@ -525,13 +699,59 @@ def __init__( def save_pretrained(self, dest): self.tokenizer.save_pretrained(dest) + def encode(self, text: str, add_special_tokens: bool = False) -> List[int]: + text = self.clean_fn(text) + if self.tokenizer_mode == 'clips': + tokens = self.tokenizer.encode(text, add_special_tokens=False) + if add_special_tokens: + tokens = [self.tokenizer.bos_token_id] + tokens + [ + self.tokenizer.eos_token_id, + self.tokenizer.cls_token_id, + ] + else: + tokens = self.tokenizer.encode(text, add_special_tokens=add_special_tokens) + + if self.strip_sep_token and self.tokenizer.sep_token_id in tokens: + fill_id = _get_pad_fill_id(self.pad_token_id) + tokens = [fill_id if token == self.tokenizer.sep_token_id else token for token in tokens] + return tokens + + def decode( + self, + tokens: TokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> str: + return _decode_with_backend( + self.tokenizer, + tokens, + self.eot_token_id, + skip_special_tokens, + stop_at_eos, + ) + + def batch_decode( + self, + batch_tokens: BatchTokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> List[str]: + return _batch_decode_with_backend( + self.tokenizer, + batch_tokens, + self.eot_token_id, + skip_special_tokens, + stop_at_eos, + ) + def __call__( self, - texts: Union[str, List[str]], + texts: TokenizerInput, context_length: Optional[int] = None, pad: bool = True, output_mask: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + add_special_tokens: bool = True, + ) -> TokenizerOutput: # same cleaning as for default tokenizer, except lowercasing # adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance if isinstance(texts, str): @@ -550,7 +770,12 @@ def __call__( # Handle different tokenization modes if self.tokenizer_mode == 'clips': - return self._clips_tokenize(texts, context_length, pad=pad) + return self._clips_tokenize( + texts, + context_length, + pad=pad, + add_special_tokens=add_special_tokens, + ) else: # Standard tokenization encoded = self.tokenizer( @@ -559,13 +784,13 @@ def __call__( padding='max_length' if pad else False, truncation=True, return_tensors='pt' if pad else None, + add_special_tokens=add_special_tokens, ) input_ids = encoded.input_ids if pad else encoded["input_ids"] attn_mask = encoded.attention_mask.bool() if (pad and output_mask) else None if self.strip_sep_token: - # pad_token_id can legitimately be None (no reserved pad token); fall back to the historical 0. - fill_id = 0 if self.pad_token_id is None else self.pad_token_id + fill_id = _get_pad_fill_id(self.pad_token_id) if pad: sep_positions = input_ids == self.tokenizer.sep_token_id input_ids = torch.where( @@ -601,6 +826,7 @@ def _clips_tokenize( texts: List[str], context_length: int, pad: bool = True, + add_special_tokens: bool = True, ) -> Union[torch.Tensor, List[torch.Tensor]]: """Use standard HF tokenizer but apply custom post-processing""" # Use standard tokenizer without special tokens - we'll add our own @@ -613,18 +839,23 @@ def _clips_tokenize( ) encoded = [] + num_special_tokens = 3 if add_special_tokens else 0 for tokens in encoded_outputs["input_ids"]: - tokens = tokens[:context_length - 3] # Leave room for special tokens - tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id] + tokens = tokens[:context_length - num_special_tokens] + if add_special_tokens: + tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id] encoded.append(tokens) if not pad: # Match the padded contract: the class token terminates the sequence. The body is truncated to # context_length - 3 above, so [bos] + body + [eos] + [cls] always fits within context_length. - return [ - torch.tensor(tokens + [self.tokenizer.cls_token_id], dtype=torch.long) - for tokens in encoded - ] + if add_special_tokens: + encoded = [tokens + [self.tokenizer.cls_token_id] for tokens in encoded] + return [torch.tensor(tokens, dtype=torch.long) for tokens in encoded] + + if not add_special_tokens: + fill_id = _get_pad_fill_id(self.pad_token_id) + return _pad_token_sequences(encoded, context_length, pad_token_id=fill_id) # Create result tensor and handle padding + class token result = torch.zeros(len(encoded), context_length, dtype=torch.long) @@ -703,19 +934,61 @@ def __init__( self.tokenizer.eos_token_id = 1 self.pad_token_id = self.tokenizer.pad_token_id self.eot_token_id = self.tokenizer.eos_token_id + self.eos_token_id = self.eot_token_id + self.sot_token_id = self.tokenizer.bos_token_id + self.bos_token_id = self.sot_token_id + self.all_special_ids = self.tokenizer.all_special_ids self.vocab_size = len(self.tokenizer) self.context_length = context_length def save_pretrained(self, dest): self.tokenizer.save_pretrained(dest) + def _clean(self, text: str) -> str: + return canonicalize_text(basic_clean(text)) + + def encode(self, text: str, add_special_tokens: bool = False) -> List[int]: + return self.tokenizer.encode( + self._clean(text), + add_special_tokens=add_special_tokens, + ) + + def decode( + self, + tokens: TokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> str: + return _decode_with_backend( + self.tokenizer, + tokens, + self.eot_token_id, + skip_special_tokens, + stop_at_eos, + ) + + def batch_decode( + self, + batch_tokens: BatchTokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> List[str]: + return _batch_decode_with_backend( + self.tokenizer, + batch_tokens, + self.eot_token_id, + skip_special_tokens, + stop_at_eos, + ) + def __call__( self, - texts: Union[str, List[str]], + texts: TokenizerInput, context_length: Optional[int] = None, pad: bool = True, output_mask: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + add_special_tokens: bool = True, + ) -> TokenizerOutput: # same cleaning as for default tokenizer, except lowercasing # adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance if output_mask: @@ -728,13 +1001,14 @@ def __call__( context_length = context_length or self.context_length assert context_length, 'Please set a valid context length in class init or call.' - texts = [canonicalize_text(basic_clean(text)) for text in texts] + texts = [self._clean(text) for text in texts] output = self.tokenizer( texts, return_tensors='pt' if pad else None, max_length=context_length, padding='max_length' if pad else False, truncation=True, + add_special_tokens=add_special_tokens, ) if not pad: return [torch.tensor(tokens, dtype=torch.long) for tokens in output.input_ids] @@ -785,17 +1059,55 @@ def __init__( self.bos_token_id = base + 2 self.sot_token_id = self.bos_token_id # alias for CLIP-style callers self.all_special_ids = [self.eot_token_id, self.pad_token_id, self.bos_token_id] + self.eos_token_id = self.eot_token_id self.vocab_size = base + 3 + self._special_token_text = { + self.bos_token_id: '<|bos|>', + self.eot_token_id: '<|eos|>', + self.pad_token_id: '<|pad|>', + } - def encode(self, text: str) -> List[int]: + def encode(self, text: str, add_special_tokens: bool = False) -> List[int]: # encode_ordinary ignores any special-token markup in the text, treating it as plain bytes. if self.clean_fn is not None: text = self.clean_fn(text) - return self.enc.encode_ordinary(text) + tokens = self.enc.encode_ordinary(text) + return self._wrap(tokens) if add_special_tokens else tokens - def decode(self, tokens: List[int]) -> str: - body = [t for t in tokens if t < self.enc.n_vocab] - return self.enc.decode(body) + def decode( + self, + tokens: TokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> str: + tokens = _truncate_at_eos(tokens, self.eot_token_id, stop_at_eos) + parts = [] + body = [] + for token in tokens: + if token in self._special_token_text: + if body: + parts.append(self.enc.decode(body)) + body = [] + if not skip_special_tokens: + parts.append(self._special_token_text[token]) + elif token < self.enc.n_vocab: + body.append(token) + # Preserve the legacy tolerance for unknown ids above the tiktoken vocabulary. + if body: + parts.append(self.enc.decode(body)) + return ''.join(parts) + + def batch_decode( + self, + batch_tokens: BatchTokenIds, + skip_special_tokens: bool = False, + stop_at_eos: bool = True, + ) -> List[str]: + batch_tokens = _to_token_batch(batch_tokens) + return [ + self.decode(tokens, skip_special_tokens=skip_special_tokens, stop_at_eos=stop_at_eos) + for tokens in batch_tokens + ] def _wrap(self, ids: List[int]) -> List[int]: if self.add_bos: @@ -806,11 +1118,12 @@ def _wrap(self, ids: List[int]) -> List[int]: def __call__( self, - texts: Union[str, List[str]], + texts: TokenizerInput, context_length: Optional[int] = None, pad: bool = True, output_mask: bool = False, - ) -> Union[torch.LongTensor, List[torch.LongTensor], Tuple[torch.LongTensor, torch.Tensor]]: + add_special_tokens: bool = True, + ) -> TokenizerOutput: """Tokenize text(s). Args: @@ -821,6 +1134,7 @@ def __call__( variable-length 1-D tensors. output_mask: Also return a [N, context_length] bool attention mask (True = real token, HF polarity). Requires ``pad=True``. Exact: the pad id is reserved above the vocab. + add_special_tokens: Apply the constructor-configured BOS/EOS template. Defaults to True. """ if isinstance(texts, str): texts = [texts] @@ -829,13 +1143,13 @@ def __call__( if output_mask and not pad: raise ValueError("output_mask=True requires pad=True (variable-length collation derives its own validity).") - all_tokens = [self._wrap(self.encode(text)) for text in texts] + all_tokens = [self.encode(text, add_special_tokens=add_special_tokens) for text in texts] if context_length is not None: truncated = [] for tokens in all_tokens: if len(tokens) > context_length: tokens = tokens[:context_length] - if self.add_eos: + if add_special_tokens and self.add_eos: tokens[-1] = self.eot_token_id truncated.append(tokens) all_tokens = truncated diff --git a/src/open_clip/zero_shot_classifier.py b/src/open_clip/zero_shot_classifier.py index 36787371f..7a6d6e800 100644 --- a/src/open_clip/zero_shot_classifier.py +++ b/src/open_clip/zero_shot_classifier.py @@ -1,9 +1,11 @@ from functools import partial from itertools import islice -from typing import Callable, List, Optional, Sequence, Union +from typing import Callable, Optional, Sequence, Union import torch +from .tokenizer import Tokenizer + def batched(iterable, n): """Batch data into lists of length *n*. The last batch may be shorter. @@ -19,7 +21,7 @@ def batched(iterable, n): def build_zero_shot_classifier( model, - tokenizer, + tokenizer: Tokenizer, classnames: Sequence[str], templates: Sequence[Union[Callable, str]], num_classes_per_batch: Optional[int] = 10, @@ -75,7 +77,7 @@ def _process_batch(batch_classnames): def build_zero_shot_classifier_legacy( model, - tokenizer, + tokenizer: Tokenizer, classnames: Sequence[str], templates: Sequence[Union[Callable, str]], device: Union[str, torch.device] = 'cpu', @@ -119,4 +121,3 @@ def build_zero_shot_classifier_legacy( zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device) return zeroshot_weights - diff --git a/src/open_clip_train/audio_data.py b/src/open_clip_train/audio_data.py index e5932fba7..f12b1294a 100644 --- a/src/open_clip_train/audio_data.py +++ b/src/open_clip_train/audio_data.py @@ -4,13 +4,16 @@ import random from dataclasses import asdict, is_dataclass from functools import partial -from typing import Dict, List, Optional +from typing import Dict, List, Optional, TYPE_CHECKING import torch import webdataset as wds from torch.utils.data import DataLoader, Dataset from torch.utils.data.distributed import DistributedSampler +if TYPE_CHECKING: + from open_clip.tokenizer import Tokenizer + from open_clip_train.data import ( DataInfo, RepeatedShardList, @@ -77,7 +80,7 @@ def _decode_audio_bytes(data): class _TokenizeAudioCaption: # Module-level callable (picklable for forkserver workers). - def __init__(self, tokenizer, variable: bool = False): + def __init__(self, tokenizer: "Tokenizer", variable: bool = False): self.tokenizer = tokenizer self.variable = variable @@ -95,7 +98,7 @@ class AudioCaptionTokenizer: unpadded sequence for per-batch text padding, mirroring ``TokenizeText``. """ - def __init__(self, tokenizer, variable: bool = False): + def __init__(self, tokenizer: "Tokenizer", variable: bool = False): self.tokenizer = tokenizer self.variable = variable diff --git a/src/open_clip_train/data.py b/src/open_clip_train/data.py index 2a95f6fdb..ce36c7ad9 100644 --- a/src/open_clip_train/data.py +++ b/src/open_clip_train/data.py @@ -9,7 +9,7 @@ import random import sys import warnings -from typing import Optional +from typing import Optional, TYPE_CHECKING import braceexpand from dataclasses import dataclass @@ -28,6 +28,9 @@ from webdataset.filters import _shuffle, pipelinefilter, reraise_exception from webdataset.tariterators import base_plus_ext, url_opener, tar_file_expander, valid_sample +if TYPE_CHECKING: + from open_clip.tokenizer import Tokenizer + # Finite backstop for PIL's decompression-bomb guard (warn > this, error > 2x), mainly for the side paths # (CsvDataset, ImageFolder, legacy); the WDS path is gated tighter/earlier by --max-image-pixels in decode_pil_rgb. Image.MAX_IMAGE_PIXELS = 128_000_000 @@ -37,7 +40,7 @@ class TokenizeText: # Module-level callable replaces inline lambdas in webdataset pipelines so # they survive pickling — required under forkserver multiprocessing # (Python 3.14+ default on POSIX). - def __init__(self, tokenizer, variable: bool = False, output_mask: bool = False): + def __init__(self, tokenizer: "Tokenizer", variable: bool = False, output_mask: bool = False): self.tokenizer = tokenizer self.variable = variable # output_mask: emit a per-sample bool validity mask (batch key "text_valid", True = real token), @@ -104,7 +107,7 @@ def _map_no_key(data, f, handler=reraise_exception): ) -def get_text_pad_id(tokenizer) -> int: +def get_text_pad_id(tokenizer: "Tokenizer") -> int: pad_id = getattr(tokenizer, "pad_token_id", None) if pad_id is None: raise ValueError("variable_text=True requires a tokenizer with a reserved `pad_token_id`.") diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py new file mode 100644 index 000000000..f14efb703 --- /dev/null +++ b/tests/test_tokenizer.py @@ -0,0 +1,113 @@ +import inspect + +import pytest +import torch + +from open_clip.tokenizer import HFTokenizer, SigLipTokenizer, SimpleTokenizer, TikTokenTokenizer, Tokenizer + + +@pytest.mark.parametrize( + "tokenizer_cls", + [Tokenizer, SimpleTokenizer, HFTokenizer, SigLipTokenizer, TikTokenTokenizer], +) +def test_tokenizer_special_token_defaults(tokenizer_cls): + encode_params = inspect.signature(tokenizer_cls.encode).parameters + call_params = inspect.signature(tokenizer_cls.__call__).parameters + decode_params = inspect.signature(tokenizer_cls.decode).parameters + batch_decode_params = inspect.signature(tokenizer_cls.batch_decode).parameters + + assert encode_params["add_special_tokens"].default is False + assert call_params["add_special_tokens"].default is True + for params in (decode_params, batch_decode_params): + assert params["skip_special_tokens"].default is False + assert params["stop_at_eos"].default is True + + +def test_simple_tokenizer_special_token_controls(): + tokenizer = SimpleTokenizer(context_length=8) + body = tokenizer.encode("hello") + wrapped = tokenizer.encode("hello", add_special_tokens=True) + + assert wrapped == [tokenizer.sot_token_id, *body, tokenizer.eot_token_id] + model_tokens = tokenizer("hello")[0] + assert model_tokens[:len(wrapped)].tolist() == wrapped + assert tokenizer.decode(model_tokens) == "hello " + assert tokenizer.decode(model_tokens, skip_special_tokens=True) == "hello " + assert tokenizer.decode(model_tokens, stop_at_eos=False).endswith("!!!!!") + assert tokenizer.batch_decode(model_tokens.unsqueeze(0)) == ["hello "] + + body_tokens, body_valid = tokenizer( + "hello", + add_special_tokens=False, + output_mask=True, + ) + assert body_tokens[0, :len(body)].tolist() == body + assert body_valid[0].sum().item() == len(body) + + +@pytest.mark.parametrize("reduction_mask", ["simple", "random", "shuffle"]) +def test_simple_tokenizer_body_only_reduction_mask(reduction_mask): + tokenizer = SimpleTokenizer(context_length=4, reduction_mask=reduction_mask) + tokens, valid = tokenizer( + "one two three four five", + add_special_tokens=False, + output_mask=True, + ) + + assert tokens.shape == valid.shape == (1, 4) + assert valid.all() + assert tokenizer.sot_token_id not in tokens + assert tokenizer.eot_token_id not in tokens + + +def test_tiktoken_special_token_controls(): + pytest.importorskip("tiktoken") + tokenizer = TikTokenTokenizer(context_length=8) + body = tokenizer.encode("hello") + wrapped = tokenizer.encode("hello", add_special_tokens=True) + + assert wrapped == [tokenizer.bos_token_id, *body, tokenizer.eot_token_id] + model_tokens = tokenizer("hello")[0] + assert tokenizer.decode(model_tokens) == "<|bos|>hello<|eos|>" + assert tokenizer.decode(model_tokens, skip_special_tokens=True) == "hello" + assert "<|pad|>" in tokenizer.decode(model_tokens, stop_at_eos=False) + assert tokenizer.batch_decode(model_tokens.unsqueeze(0)) == ["<|bos|>hello<|eos|>"] + unknown_id = tokenizer.vocab_size + 10 + assert tokenizer.decode([tokenizer.bos_token_id, *body, unknown_id, tokenizer.eot_token_id]) == ( + "<|bos|>hello<|eos|>" + ) + + +class _FakeHFBackend: + eos_token_id = 2 + + def __init__(self): + self.encode_add_special_tokens = None + + def encode(self, text, add_special_tokens=True): + self.encode_add_special_tokens = add_special_tokens + return [10, self.eos_token_id] if add_special_tokens else [10] + + def decode(self, tokens, skip_special_tokens=False): + return f"{list(tokens)}:{skip_special_tokens}" + + def batch_decode(self, batch_tokens, skip_special_tokens=False): + return [f"{list(tokens)}:{skip_special_tokens}" for tokens in batch_tokens] + + +def test_hf_tokenizer_encode_decode_controls_delegate_cleanly(): + tokenizer = HFTokenizer.__new__(HFTokenizer) + tokenizer.tokenizer = _FakeHFBackend() + tokenizer.tokenizer_mode = "" + tokenizer.clean_fn = lambda text: text + tokenizer.strip_sep_token = False + tokenizer.eot_token_id = tokenizer.tokenizer.eos_token_id + + assert tokenizer.encode("hello") == [10] + assert tokenizer.tokenizer.encode_add_special_tokens is False + assert tokenizer.encode("hello", add_special_tokens=True) == [10, 2] + assert tokenizer.decode(torch.tensor([10, 2, 99])) == "[10, 2]:False" + assert tokenizer.decode([10, 2, 99], skip_special_tokens=True) == "[10, 2]:True" + assert tokenizer.batch_decode([[10, 2, 99], [11]]) == ["[10, 2]:False", "[11]:False"] + batch_tokens = torch.tensor([[10, 2, 99], [11, 2, 98]]) + assert tokenizer.batch_decode(batch_tokens) == ["[10, 2]:False", "[11, 2]:False"]