xiaomi-research/OneVL_visual_decoder_pt_ar1

1"""2Fast-loading Qwen3-VL tokenizer with 131k visual tokens.3 4Visual tokens live in model.vocab (fast BPE hash-map load) rather than5added_tokens (slow Aho-Corasick build).  A regex pre-split in the Python6wrapper ensures encode/call with visual token text produces single IDs.7 8Strategy: replace each <|visual token XXXXXX|> with a NUL byte (\x00)9before sending to the Rust backend, then swap the NUL-byte token ID (188)10with the real visual-token ID in the output.11"""12 13import re14from typing import List, Optional, Union15 16from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast17 18_VISUAL_RE = re.compile(r"<\|visual token (\d{6})\|>")19_VISUAL_TOKEN_START_ID = 15167420_PLACEHOLDER_CHAR = "\x00"21_PLACEHOLDER_TOKEN_ID = 18822 23 24class Qwen3VLVisualTokenizerFast(Qwen2TokenizerFast):25 26    # ---------- public encode() ----------27    def encode(self, text, text_pair=None, add_special_tokens=True, **kwargs):28        if isinstance(text, str) and _VISUAL_RE.search(text):29            replaced, vids = _replace_visual(text)30            pair_replaced, pair_vids = None, []31            if text_pair is not None and isinstance(text_pair, str):32                pair_replaced, pair_vids = _replace_visual(text_pair)33            ids = super().encode(34                replaced,35                text_pair=pair_replaced if pair_replaced is not None else text_pair,36                add_special_tokens=add_special_tokens,37                **kwargs,38            )39            _swap_ids(ids, vids + pair_vids)40            return ids41        return super().encode(text, text_pair, add_special_tokens=add_special_tokens, **kwargs)42 43    # ---------- batch path (powers __call__) ----------44    def _batch_encode_plus(self, batch_text_or_text_pairs, **kwargs):45        has_visual = any(46            _text_has_visual(item) for item in batch_text_or_text_pairs47        )48        if not has_visual:49            return super()._batch_encode_plus(batch_text_or_text_pairs, **kwargs)50 51        replaced_batch = []52        all_vids: list[list[int]] = []53 54        for item in batch_text_or_text_pairs:55            if isinstance(item, (tuple, list)):56                text, pair = item[0], (item[1] if len(item) > 1 else None)57            else:58                text, pair = item, None59 60            vids: list[int] = []61            if isinstance(text, str) and _VISUAL_RE.search(text):62                text, tvids = _replace_visual(text)63                vids.extend(tvids)64            if pair is not None and isinstance(pair, str) and _VISUAL_RE.search(pair):65                pair, pvids = _replace_visual(pair)66                vids.extend(pvids)67 68            replaced_batch.append((text, pair) if pair is not None else text)69            all_vids.append(vids)70 71        result = super()._batch_encode_plus(replaced_batch, **kwargs)72 73        for i, vids in enumerate(all_vids):74            if not vids:75                continue76            ids = result["input_ids"][i]77            tensor_type = None78            if hasattr(ids, "tolist"):79                tensor_type = type(ids)80                device = ids.device if hasattr(ids, "device") else None81                dtype = ids.dtype82                ids = ids.tolist()83            _swap_ids(ids, vids)84            if tensor_type is not None:85                import torch86                t = torch.tensor(ids, dtype=dtype)87                if device is not None:88                    t = t.to(device)89                result["input_ids"][i] = t90            else:91                result["input_ids"][i] = ids92 93        return result94 95 96def _text_has_visual(item) -> bool:97    t = item[0] if isinstance(item, (tuple, list)) else item98    return isinstance(t, str) and _VISUAL_RE.search(t) is not None99 100 101def _replace_visual(text: str):102    """Replace visual tokens with NUL bytes, return (new_text, ordered_visual_ids)."""103    vids: list[int] = []104 105    def _repl(m):106        vids.append(_VISUAL_TOKEN_START_ID + int(m.group(1)))107        return _PLACEHOLDER_CHAR108 109    new_text = _VISUAL_RE.sub(_repl, text)110    return new_text, vids111 112 113def _swap_ids(ids: list, vids: list[int]):114    """In-place replace placeholder token IDs with real visual-token IDs."""115    vi = 0116    for j in range(len(ids)):117        if ids[j] == _PLACEHOLDER_TOKEN_ID and vi < len(vids):118            ids[j] = vids[vi]119            vi += 1120