Hugging Face's logo

HuggingFaceBio
/
Carbon-X

Text Generation
Transformers
Safetensors
English
qwen3
dna
genomics
biology
variant-effect-prediction
conversational
text-generation-inference
Model card Files
xet
 Community
2
Carbon-X / tokenizer.py
lewtun's picture
lewtun HF Staff
loubnabnl's picture
loubnabnl HF Staff
danaaubakirova's picture
danaaubakirova HF Staff
Carbon-X
1fac52c
Raw History Blame Contribute Delete
33.9 kB
"""
HybridDNATokenizer: Combines Qwen3 BPE tokenization with DNA 6-mer tokenization.
DNA spans are detected automatically or supplied in <dna>...</dna> tags and
tokenized as 6-mers. All remaining text uses Qwen3's BPE tokenization.
Supports token_mask for Fine-grained Nucleotide Supervision (FNS):
-2: padding token
-1: text token (BPE)
0: DNA special token (<dna>, </dna>, <oov>)
1-5: partial 6-mer token — valid_length real bases at positions [0, valid_length),
right-padded with 'A' at positions [valid_length, k) so loss can supervise
positions 0..valid_len-1 via pos_mask = (valid_len > pos)
6: full 6-mer
"""
import os
import json
import itertools
import re
from typing import List, Optional, Tuple, Dict, Union, Any
from transformers import PreTrainedTokenizer, AutoTokenizer, BatchEncoding
# The BPE vocabulary comes from this tokenizer, and its size sets where the DNA ids start. Pin it so that a change on the
# Hub cannot silently shift the ids that checkpoints were trained with.
BASE_TOKENIZER = "Qwen/Qwen3-4B-Base"
BASE_TOKENIZER_REVISION = "906bfd4b4dc7f14ee4320094d8b41684abff8539"
class HybridDNATokenizer(PreTrainedTokenizer):
"""
Hybrid tokenizer combining Qwen3 BPE with DNA 6-mer tokenization.
Explicit <dna>...</dna> regions are always tokenized as 6-mers. By default,
untagged DNA spans embedded in text are detected and tagged automatically;
pass auto_dna_tags=False to encode() or __call__() to disable detection for
one call.
"""
model_input_names = ["input_ids", "attention_mask"]
# Single-nucleotide alphabet for tail="char": standard bases (A C T G U) +
# IUPAC ambiguity codes (R Y S W K M B D H V N). These 16 chars match
# microsoft/NatureLM-8x7B nlm_tokenizer.py's single-char DNA set and the
# NatureLM paper's stated DNA vocabulary size of 16 (arXiv:2502.07527, S2.1).
SINGLE_BASE_ALPHABET = "ACTGURYSWKMBDHVN"
IUPAC_DNA_ALPHABET = "ACGTRYSWKMBDHVN"
STANDARD_DNA_ALPHABET = frozenset("ACGT")
IUPAC_RUN_RE = re.compile(f"[{IUPAC_DNA_ALPHABET}]+", re.IGNORECASE | re.ASCII)
SPACED_DNA_RE = re.compile(
r"(?<![A-Za-z])(?:[ACGT]{1,5}[ \t]+){5,}[ACGT]{1,5}(?![A-Za-z])"
)
def __init__(
self,
base_tokenizer_path: Optional[str] = None,
k: int = 6,
auto_dna_tags: bool = True,
tail: str = "pad_a",
**kwargs
):
if tail not in ("pad_a", "pad_n", "char"):
raise ValueError(f"tail must be 'pad_a', 'pad_n' or 'char', got {tail!r}")
self.tail = tail
self.pad_char = {"pad_a": "A", "pad_n": "N"}.get(tail)
self.k = k
# Load base tokenizer (Qwen3-4B-Base) at the pinned revision
self._base_tokenizer = AutoTokenizer.from_pretrained(BASE_TOKENIZER, revision=BASE_TOKENIZER_REVISION)
# Get base vocabulary
self._base_vocab = self._base_tokenizer.get_vocab()
self._base_vocab_size = len(self._base_vocab)
# Initialize DNA vocabulary
self._init_dna_vocab()
# Build combined vocabulary
self._build_combined_vocab()
# Set special tokens
self._eos_token = kwargs.pop('eos_token', None) or "<|endoftext|>"
self._pad_token = kwargs.pop('pad_token', None) or self._base_tokenizer.pad_token or "<|endoftext|>"
# Initialize parent class
super().__init__(
eos_token=self._eos_token,
pad_token=self._pad_token,
**kwargs
)
# Preserve the base tokenizer's AddedToken metadata. These tokens are
# already present in the combined vocabulary at their original IDs, so
# this only registers their metadata and does not grow the vocabulary.
self.add_tokens(list(self._base_tokenizer.added_tokens_decoder.values()))
self.special_tokens = self.dna_special_tokens + [self._eos_token, self._pad_token]
self.auto_dna_tags = auto_dna_tags
def _init_dna_vocab(self):
"""Initialize DNA vocabulary (special tokens + k-mers + padding for 128 alignment)."""
bases = ['A', 'T', 'C', 'G']
# DNA special tokens
self.dna_special_tokens = ["<dna>", "</dna>", "<oov>"]
# Generate all k-mer combinations (4^k = 4096 for k=6)
self.kmers = [''.join(kmer) for kmer in itertools.product(bases, repeat=self.k)]
# tail="pad_n": dedicated tokens for the N-padded tail ("CT" -> "CTNNNN"),
# one per real prefix of length 1..k-1, i.e. 4^1+...+4^(k-1) = 1364.
self.npad_tokens = []
if self.tail == "pad_n":
for prefix_len in range(1, self.k):
for prefix in itertools.product(bases, repeat=prefix_len):
self.npad_tokens.append(''.join(prefix) + 'N' * (self.k - prefix_len))
# tail="char": single-base tokens for the trailing < k bases (NatureLM
# IUPAC alphabet), so N / ambiguity codes fall back to a token, not <oov>.
self.base_tokens = []
self.dna_base_token = {} # base -> token, e.g. "A" -> "<b_A>"
self.dna_base_token_to_base = {} # token -> base
if self.tail == "char":
for b in self.SINGLE_BASE_ALPHABET:
tok = f"<b_{b}>"
self.base_tokens.append(tok)
self.dna_base_token[b] = tok
self.dna_base_token_to_base[tok] = b
# DNA tokens start after base vocabulary
self.dna_start_id = self._base_vocab_size
# All DNA tokens get new IDs (no reuse of base vocab IDs, even for
# overlapping tokens like CCCCCC — they have different semantics in
# DNA context vs BPE context, per Qiuyi's recommendation). Tail-variant
# tokens are listed AFTER the 4^k k-mers so the k-mer ids are unchanged.
base_dna_tokens = self.dna_special_tokens + self.kmers + self.npad_tokens + self.base_tokens
# Calculate padding for 128 alignment
total_vocab_unpadded = self._base_vocab_size + len(base_dna_tokens)
target_vocab_size = ((total_vocab_unpadded + 127) // 128) * 128
num_padding_tokens = target_vocab_size - total_vocab_unpadded
# Add unused padding tokens
self.padding_tokens = [f"<unused_{i}>" for i in range(num_padding_tokens)]
# Create DNA token mappings — all get sequential new IDs
self.dna_token_to_id = {}
self.dna_id_to_token = {}
current_id = self.dna_start_id
for token in base_dna_tokens:
self.dna_token_to_id[token] = current_id
self.dna_id_to_token[current_id] = token
current_id += 1
# Add padding tokens
for token in self.padding_tokens:
self.dna_token_to_id[token] = current_id
self.dna_id_to_token[current_id] = token
current_id += 1
self.dna_vocab_size = len(base_dna_tokens) + len(self.padding_tokens)
# Set DNA special token IDs
self.dna_begin_token_id = self.dna_token_to_id["<dna>"]
self.dna_end_token_id = self.dna_token_to_id["</dna>"]
self.oov_token_id = self.dna_token_to_id["<oov>"]
def _build_combined_vocab(self):
"""Build combined vocabulary (base + DNA)."""
self._vocab = self._base_vocab.copy()
for token, token_id in self.dna_token_to_id.items():
if token not in self._vocab:
self._vocab[token] = token_id
self._id_to_token = {v: k for k, v in self._vocab.items()}
for token_id, token in self.dna_id_to_token.items():
if token_id not in self._id_to_token:
self._id_to_token[token_id] = token
@property
def vocab_size(self) -> int:
return max(self._vocab.values()) + 1
def get_vocab(self) -> Dict[str, int]:
return self._vocab.copy()
def __len__(self):
# Override default (len(get_vocab())) because get_vocab() deduplicates
# CCCCCC which exists as both BPE (ID 91443) and DNA 6-mer (ID 154402).
return self.vocab_size
@classmethod
def _auto_tag_text_segment(cls, text: str) -> str:
"""Tag conservative DNA candidates in a text-only segment.
Each maximal IUPAC nucleotide run is considered independently. Pure
A/C/G/T runs need at least six bases. Runs containing IUPAC ambiguity
codes need at least twelve bases and may contain at most 20% ambiguity.
Rejected ambiguous runs stay wholly on the BPE path rather than having
any A/C/G/T substring tagged.
"""
parts = []
previous_end = 0
for match in cls.IUPAC_RUN_RE.finditer(text):
candidate = match.group(0)
normalized = candidate.upper()
ambiguous_count = sum(base not in cls.STANDARD_DNA_ALPHABET for base in normalized)
is_standard_dna = ambiguous_count == 0 and len(normalized) >= 6
is_conservative_iupac = (
ambiguous_count > 0
and len(normalized) >= 12
and ambiguous_count / len(normalized) <= 0.2
)
if not (is_standard_dna or is_conservative_iupac):
continue
parts.append(text[previous_end:match.start()])
parts.append(f"<dna>{normalized}</dna>")
previous_end = match.end()
if not parts:
return text
parts.append(text[previous_end:])
return "".join(parts)
@classmethod
def _normalize_spaced_dna(cls, text: str) -> str:
"""Join uppercase DNA displays with at least six chunks and twelve bases.
Chunks contain one to five A/C/G/T bases, separated only by spaces or
tabs. Alphabetic boundaries prevent joining fragments of words.
"""
def normalize(match):
sequence = "".join(match.group(0).split())
return sequence if len(sequence) >= 12 else match.group(0)
return cls.SPACED_DNA_RE.sub(normalize, text)
def _add_auto_dna_tags(self, text: str) -> str:
"""Tag DNA outside existing <dna>...</dna> regions."""
return "".join(
segment if is_dna else self._auto_tag_text_segment(self._normalize_spaced_dna(segment))
for segment, is_dna in self._split_by_dna_tags(text)
)
def _split_by_dna_tags(self, text: str) -> List[Tuple[str, bool]]:
segments = []
i = 0
n = len(text)
while i < n:
start_pos = text.find('<dna>', i)
end_pos = text.find('</dna>', i)
if start_pos == -1 and end_pos == -1:
remaining = text[i:]
if remaining:
segments.append((remaining, False))
break
if start_pos == -1 and end_pos != -1:
dna_region = text[i:end_pos + 6]
if dna_region:
segments.append((dna_region, True))
i = end_pos + 6
continue
if start_pos != -1 and end_pos == -1:
if i < start_pos:
normal_text = text[i:start_pos]
if normal_text:
segments.append((normal_text, False))
dna_region = text[start_pos:]
if dna_region:
segments.append((dna_region, True))
break
if start_pos < end_pos:
if i < start_pos:
normal_text = text[i:start_pos]
if normal_text:
segments.append((normal_text, False))
dna_region = text[start_pos:end_pos + 6]
if dna_region:
segments.append((dna_region, True))
i = end_pos + 6
else:
dna_region = text[i:end_pos + 6]
if dna_region:
segments.append((dna_region, True))
i = end_pos + 6
return segments
def _parse_dna_region(self, dna_region: str) -> Tuple[str, bool, bool]:
if dna_region == '<dna>':
return '', True, False
elif dna_region == '</dna>':
return '', False, True
has_start = dna_region.startswith('<dna>')
has_end = dna_region.endswith('</dna>')
content = dna_region
if has_start:
content = content[5:]
if has_end and content.endswith('</dna>'):
content = content[:-6]
return content.strip(), has_start, has_end
def _process_dna_sequence(self, dna_seq: str) -> Dict:
k = self.k
dna_seq = dna_seq.upper()
kmer_tokens = []
valid_bases = set('ATCG')
def is_valid_kmer(kmer):
return len(kmer) == k and all(base in valid_bases for base in kmer)
if self.tail == "char":
# Greedy NatureLM tokenization (NatureLM, arXiv:2502.07527), ported
# from microsoft/NatureLM-8x7B nlm_tokenizer.py::_tokenize_entity:
# take a full k-mer when the next k chars form one, else emit a single
# base token and advance by 1 (also re-aligns after any non-base char).
i = 0
n = len(dna_seq)
while i < n:
window = dna_seq[i:i + k]
if is_valid_kmer(window):
kmer_tokens.append(window)
i += k
else:
base = dna_seq[i]
kmer_tokens.append(self.dna_base_token.get(base, "<oov>"))
i += 1
return {"kmer_tokens": kmer_tokens, "padding_length": 0, "valid_length": k}
for i in range(0, len(dna_seq) - k + 1, k):
kmer = dna_seq[i:i+k]
if is_valid_kmer(kmer):
kmer_tokens.append(kmer)
else:
kmer_tokens.append("<oov>")
processed_length = len(kmer_tokens) * k
remaining = dna_seq[processed_length:]
padding_length = 0
valid_length = k
if remaining:
padding_needed = k - len(remaining)
# Right-pad up to a full k-mer: real bases occupy positions
# [0, valid_length). The hybrid BP loss supervises positions
# 0..valid_len-1 via pos_mask = (valid_len > pos), so padding must be
# at the END. pad_char="A" reuses a real k-mer id ("CT"->"CTAAAA");
# "N" maps to a dedicated tail token ("CT"->"CTNNNN").
padded = remaining + self.pad_char * padding_needed
if all(base in valid_bases for base in remaining):
kmer_tokens.append(padded)
else:
kmer_tokens.append("<oov>")
padding_length = padding_needed
valid_length = len(remaining)
return {
"kmer_tokens": kmer_tokens,
"padding_length": padding_length,
"valid_length": valid_length,
}
def _tokenize(self, text: str, **kwargs) -> List[str]:
return list(text)
def _convert_token_to_id(self, token: str) -> int:
if token in self.dna_token_to_id:
return self.dna_token_to_id[token]
return self._base_vocab.get(token, self._base_tokenizer.unk_token_id or 0)
def _convert_id_to_token(self, index: int) -> str:
if index in self.dna_id_to_token:
return self.dna_id_to_token[index]
return self._id_to_token.get(index, "<oov>")
def convert_tokens_to_string(self, tokens: List[str]) -> str:
ids = [self._convert_token_to_id(t) for t in tokens]
return self.decode(ids, skip_special_tokens=False)
def encode(
self,
text: str,
add_special_tokens: bool = False,
return_token_mask: bool = False,
auto_dna_tags: Optional[bool] = None,
**kwargs
) -> Union[List[int], Tuple[List[int], List[int]]]:
use_auto = self.auto_dna_tags if auto_dna_tags is None else auto_dna_tags
if use_auto:
text = self._add_auto_dna_tags(text)
segments = self._split_by_dna_tags(text)
token_ids = []
token_mask = [] if return_token_mask else None
for segment_content, is_dna in segments:
if is_dna:
dna_content, has_start, has_end = self._parse_dna_region(segment_content)
if has_start:
token_ids.append(self.dna_begin_token_id)
if return_token_mask:
token_mask.append(0)
if dna_content:
result = self._process_dna_sequence(dna_content)
for idx, kmer in enumerate(result["kmer_tokens"]):
token_id = self.dna_token_to_id.get(kmer, self.oov_token_id)
token_ids.append(token_id)
if return_token_mask:
if kmer == "<oov>":
token_mask.append(0)
elif kmer in self.dna_base_token_to_base:
token_mask.append(1) # char-fallback single base
elif idx == len(result["kmer_tokens"]) - 1 and result["padding_length"] > 0:
token_mask.append(result["valid_length"])
else:
token_mask.append(self.k)
if has_end:
token_ids.append(self.dna_end_token_id)
if return_token_mask:
token_mask.append(0)
else:
base_ids = self._base_tokenizer.encode(
segment_content,
add_special_tokens=False
)
token_ids.extend(base_ids)
if return_token_mask:
token_mask.extend([-1] * len(base_ids))
# Do NOT append EOS when add_special_tokens=True. Qwen3 doesn't add
# BOS/EOS either, and appending EOS here breaks lighteval's
# tok_encode_pair: it relies on
# len(encode(ctx)) + len(encode(answer)) == len(encode(ctx + answer))
# which the extra EOS violates by shifting the split by 1.
if return_token_mask:
return token_ids, token_mask
return token_ids
def _decode_dna_token(self, token_str: str) -> str:
# Surface text for a DNA-vocab token. Baseline (pad_a) returns it
# unchanged; char maps "<b_A>"->"A"; pad_n strips the N-pad tail fill
# ("CTNNNN"->"CT"). Real k-mers contain no N, so rstrip is a no-op there.
if token_str in self.dna_base_token_to_base:
return self.dna_base_token_to_base[token_str]
if self.tail == "pad_n" and len(token_str) == self.k and token_str[0] in "ATCG":
return token_str.rstrip("N")
return token_str
def decode(
self,
token_ids: Union[int, List[int]],
skip_special_tokens: bool = False,
**kwargs
) -> str:
# A raw tensor/ndarray must become plain ints: dna_id_to_token is keyed by
# python ints and a 0-dim tensor scalar hashes by identity, so DNA lookups
# silently miss and the bases are dropped (e.g. decode(generate(...)[0])).
if hasattr(token_ids, "tolist"):
token_ids = token_ids.tolist()
if isinstance(token_ids, int):
token_ids = [token_ids]
if skip_special_tokens:
special_ids = {self.eos_token_id, self.pad_token_id}
token_ids = [tid for tid in token_ids if tid not in special_ids]
parts = []
i = 0
while i < len(token_ids):
tid = token_ids[i]
if tid == self.dna_begin_token_id:
dna_tokens = []
i += 1
while i < len(token_ids) and token_ids[i] != self.dna_end_token_id:
if token_ids[i] in self.dna_id_to_token:
dna_tokens.append(self._decode_dna_token(self.dna_id_to_token[token_ids[i]]))
i += 1
dna_seq = ''.join(dna_tokens)
if skip_special_tokens:
parts.append(dna_seq)
else:
parts.append(f"<dna>{dna_seq}")
if i < len(token_ids) and token_ids[i] == self.dna_end_token_id:
parts.append("</dna>")
i += 1
elif tid in self.dna_id_to_token:
# This branch handles k-mer tokens that appear without a <dna>
# wrapper — the common generation case where <dna> was in the
# prompt but only the generated portion is being decoded.
# K-mer tokens are content, not special tokens, so always decode
# them. Only drop true DNA special tokens (<dna>, </dna>, <oov>)
# when skip_special_tokens=True.
is_dna_special = tid in (self.dna_begin_token_id, self.dna_end_token_id, self.oov_token_id)
if not (skip_special_tokens and is_dna_special):
parts.append(self._decode_dna_token(self.dna_id_to_token[tid]))
i += 1
else:
text_ids = []
while i < len(token_ids):
curr_id = token_ids[i]
if curr_id in self.dna_id_to_token or curr_id == self.dna_begin_token_id:
break
text_ids.append(curr_id)
i += 1
if text_ids:
decoded = self._base_tokenizer.decode(text_ids, skip_special_tokens=skip_special_tokens)
parts.append(decoded)
return ''.join(parts)
def batch_decode(
self,
sequences: Union[List[int], List[List[int]], "torch.Tensor"],
skip_special_tokens: bool = False,
**kwargs
) -> List[str]:
return [
self.decode(
seq.tolist() if hasattr(seq, 'tolist') else list(seq),
skip_special_tokens=skip_special_tokens,
**kwargs
)
for seq in sequences
]
def __call__(
self,
text: Union[str, List[str]],
add_special_tokens: bool = False,
padding: bool = False,
truncation: bool = False,
max_length: Optional[int] = None,
return_tensors: Optional[str] = None,
return_token_mask: bool = False,
auto_dna_tags: Optional[bool] = None,
**kwargs
) -> Dict[str, Any]:
is_batch = isinstance(text, list)
texts = text if is_batch else [text]
all_ids = []
all_masks = [] if return_token_mask else None
for t in texts:
if return_token_mask:
ids, mask = self.encode(t, add_special_tokens=add_special_tokens, return_token_mask=True, auto_dna_tags=auto_dna_tags)
all_ids.append(ids)
all_masks.append(mask)
else:
ids = self.encode(t, add_special_tokens=add_special_tokens, return_token_mask=False, auto_dna_tags=auto_dna_tags)
all_ids.append(ids)
if padding:
max_len = max(len(ids) for ids in all_ids)
if max_length:
max_len = min(max_len, max_length)
padded_ids = []
attention_masks = []
padded_token_masks = [] if return_token_mask else None
for idx, ids in enumerate(all_ids):
pad_len = max_len - len(ids)
if pad_len > 0:
ids = ids + [self.pad_token_id] * pad_len
attn = [1] * (max_len - pad_len) + [0] * pad_len
if return_token_mask:
mask = all_masks[idx] + [-2] * pad_len
else:
ids = ids[:max_len]
attn = [1] * max_len
if return_token_mask:
mask = all_masks[idx][:max_len]
padded_ids.append(ids)
attention_masks.append(attn)
if return_token_mask:
padded_token_masks.append(mask)
all_ids = padded_ids
all_masks = padded_token_masks
else:
attention_masks = [[1] * len(ids) for ids in all_ids]
result = {
"input_ids": all_ids if is_batch else all_ids[0],
"attention_mask": attention_masks if is_batch else attention_masks[0],
}
if return_token_mask:
result["token_mask"] = all_masks if is_batch else all_masks[0]
if return_tensors == "pt":
import torch
if is_batch:
result["input_ids"] = torch.tensor(result["input_ids"])
result["attention_mask"] = torch.tensor(result["attention_mask"])
if return_token_mask:
result["token_mask"] = torch.tensor(result["token_mask"])
else:
result["input_ids"] = torch.tensor([result["input_ids"]])
result["attention_mask"] = torch.tensor([result["attention_mask"]])
if return_token_mask:
result["token_mask"] = torch.tensor([result["token_mask"]])
return BatchEncoding(result, tensor_type=return_tensors)
# Markers delimiting an assistant turn, used to build the assistant-token mask.
# Qwen chat format; override on the instance for other bases (e.g. Gemma uses
# "<start_of_turn>model\n" / "<end_of_turn>").
ASSISTANT_HEADER = "<|im_start|>assistant\n"
TURN_END = "<|im_end|>\n"
def apply_chat_template(self, conversation, *args, return_assistant_tokens_mask: bool = False, **kwargs):
"""Completion-only assistant mask WITHOUT a fast tokenizer.
HF derives the assistant mask from ``{% generation %}`` char-spans via
``BatchEncoding.char_to_token()``, which only fast (Rust) tokenizers implement.
HybridDNATokenizer is a slow Python tokenizer (the DNA k-mer logic can't be a
``tokenizer.json``), so that path raises
``char_to_token() is not available when using Python based tokenizers``.
Instead we tokenize normally and mark the tokens of every assistant turn -- from
just after ``ASSISTANT_HEADER`` up to and including the next ``TURN_END``. This
makes ``return_assistant_tokens_mask=True`` (and thus TRL ``assistant_only_loss``)
work unchanged, and needs no ``{% generation %}`` markers in the template.
"""
auto_dna_tags = kwargs.pop("auto_dna_tags", None)
if auto_dna_tags is not None:
tokenizer_kwargs = dict(kwargs.pop("tokenizer_kwargs", None) or {})
tokenizer_kwargs["auto_dna_tags"] = auto_dna_tags
kwargs["tokenizer_kwargs"] = tokenizer_kwargs
if not return_assistant_tokens_mask:
return super().apply_chat_template(conversation, *args, **kwargs)
kwargs.pop("return_assistant_tokens_mask", None)
want_tensors = kwargs.pop("return_tensors", None)
kwargs["tokenize"] = True
kwargs["return_dict"] = True
enc = super().apply_chat_template(conversation, *args, return_tensors=None, **kwargs)
seqs = enc["input_ids"]
batched = len(seqs) > 0 and isinstance(seqs[0], (list, tuple))
batch = seqs if batched else [seqs]
hdr = self(self.ASSISTANT_HEADER, add_special_tokens=False)["input_ids"]
end = self(self.TURN_END, add_special_tokens=False)["input_ids"]
masks = []
for ids in batch:
ids = list(ids)
m = [0] * len(ids)
i = 0
while i <= len(ids) - len(hdr):
if ids[i:i + len(hdr)] == hdr:
j = i + len(hdr)
while j < len(ids) and ids[j:j + len(end)] != end:
m[j] = 1
j += 1
for k in range(len(end)): # include the turn-end token(s)
if j + k < len(ids):
m[j + k] = 1
i = j + len(end)
else:
i += 1
masks.append(m)
enc["assistant_masks"] = masks if batched else masks[0]
if want_tensors is not None:
enc = enc.convert_to_tensors(tensor_type=want_tensors)
return enc
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
vocab_file = os.path.join(
save_directory,
(filename_prefix + "-" if filename_prefix else "") + "vocab.json"
)
with open(vocab_file, "w", encoding="utf-8") as f:
json.dump(self._vocab, f, ensure_ascii=False, indent=2)
return (vocab_file,)
def save_pretrained(self, save_directory: str, **kwargs):
os.makedirs(save_directory, exist_ok=True)
# Save base tokenizer files with the wrapper's effective chat template.
self._base_tokenizer.chat_template = self.chat_template
self._base_tokenizer.save_pretrained(save_directory)
# Save DNA config
dna_config = {
"k": self.k,
"dna_start_id": self.dna_start_id,
"dna_vocab_size": self.dna_vocab_size,
"dna_special_tokens": self.dna_special_tokens,
"auto_dna_tags": self.auto_dna_tags,
"tail": self.tail,
}
dna_config_path = os.path.join(save_directory, "dna_config.json")
with open(dna_config_path, "w", encoding="utf-8") as f:
json.dump(dna_config, f, indent=2)
# Update tokenizer_config.json with auto_map
config_path = os.path.join(save_directory, "tokenizer_config.json")
if os.path.exists(config_path):
with open(config_path, "r") as f:
config = json.load(f)
else:
config = {}
config.update({
"tokenizer_class": "HybridDNATokenizer",
"auto_map": {
"AutoTokenizer": ["tokenizer.HybridDNATokenizer", None]
},
"k": self.k,
"auto_dna_tags": self.auto_dna_tags,
"tail": self.tail,
})
with open(config_path, "w", encoding="utf-8") as f:
json.dump(config, f, indent=2, ensure_ascii=False)
# Copy this tokenizer.py to save directory
import shutil
src_py = os.path.abspath(__file__)
dst_py = os.path.join(save_directory, "tokenizer.py")
if os.path.exists(src_py) and src_py != dst_py:
shutil.copy2(src_py, dst_py)
return (save_directory,)
@classmethod
def from_pretrained(
cls,
pretrained_model_name_or_path,
*args,
cache_dir=None,
force_download=False,
local_files_only=False,
token=None,
revision="main",
trust_remote_code=False,
**kwargs,
):
# Custom loader: PreTrainedTokenizer.from_pretrained is incompatible with this
# tokenizer's __init__ (it re-passes bos_token etc.), so construct directly from
# bio_config/dna_config. That bypasses standard metadata loading, so also read
# chat_template from tokenizer_config.json (vLLM / apply_chat_template need it).
from transformers.utils.hub import cached_file
proxies = kwargs.pop("proxies", None)
subfolder = kwargs.pop("subfolder", None)
commit_hash = kwargs.pop("_commit_hash", None)
hub_kwargs = {
"cache_dir": cache_dir,
"force_download": force_download,
"local_files_only": local_files_only,
"proxies": proxies,
"revision": revision,
"subfolder": subfolder,
"token": token,
"_commit_hash": commit_hash,
"_raise_exceptions_for_missing_entries": False,
}
cfg = {}
for _name in ("bio_config.json", "dna_config.json"):
_p = cached_file(pretrained_model_name_or_path, _name, **hub_kwargs)
if _p is not None:
with open(_p, encoding="utf-8") as _f:
cfg = json.load(_f)
break
_init = {"base_tokenizer_path": pretrained_model_name_or_path}
for _key in ("k", "tail", "auto_dna_tags"):
if _key in cfg:
_init[_key] = cfg[_key]
tokenizer_config_path = cached_file(pretrained_model_name_or_path, "tokenizer_config.json", **hub_kwargs)
if tokenizer_config_path is not None:
with open(tokenizer_config_path, encoding="utf-8") as _f:
tokenizer_config = json.load(_f)
if tokenizer_config.get("chat_template"):
_init["chat_template"] = tokenizer_config["chat_template"]
chat_template_path = cached_file(pretrained_model_name_or_path, "chat_template.jinja", **hub_kwargs)
if chat_template_path is not None:
with open(chat_template_path, encoding="utf-8") as _f:
_init["chat_template"] = _f.read()
_init.update(kwargs)
tokenizer = cls(**_init)
# DNA ids start right after the base vocabulary, so a layout that differs from the saved one would misalign every
# DNA embedding. Fail instead of tokenizing with shifted ids.
for _key in ("dna_start_id", "dna_vocab_size"):
if _key in cfg and cfg[_key] != getattr(tokenizer, _key):
raise ValueError(
f"{pretrained_model_name_or_path} was saved with {_key}={cfg[_key]}, but the base tokenizer gives "
f"{getattr(tokenizer, _key)}, so its DNA token ids would not match the model."
)
return tokenizer