from transformers import PreTrainedTokenizer, AddedToken from typing import List, Optional, Union, Dict, Sequence, Tuple from pathlib import Path import numpy as np import json import os class HybriDNATokenizer(PreTrainedTokenizer): model_input_names = ["input_ids", "attention_mask"] def __init__(self, model_max_length: int, bos_token="[BOS]", eos_token="[SEP]", sep_token="[SEP]", cls_token="[CLS]", pad_token="[PAD]", mask_token="[MASK]", unk_token="[UNK]", **kwargs): """Character tokenizer for Hugging Face transformers. Args: characters (Sequence[str]): List of desired characters. Any character which is not included in this list will be replaced by a special token called [UNK] with id=6. Following are list of all of the special tokens with their corresponding ids: "[CLS]": 0 "[SEP]": 1 "[BOS]": 2 "[MASK]": 3 "[PAD]": 4 "[RESERVED]": 5 "[UNK]": 6 an id (starting at 7) will be assigned to each character. model_max_length (int): Model maximum sequence length. """ self.characters = ('A', 'C', 'G', 'T', 'N') self.model_max_length = model_max_length self._vocab_str_to_int = { "[CLS]": 0, "[SEP]": 1, "[BOS]": 2, "[MASK]": 3, "[PAD]": 4, "[RESERVED]": 5, "[UNK]": 6, **{ch: i + 7 for i, ch in enumerate(self.characters)}, } self._vocab_int_to_str = {v: k for k, v in self._vocab_str_to_int.items()} self._bos_id = self._vocab_str_to_int["[BOS]"] self._eos_id = self._vocab_str_to_int["[SEP]"] self._pad_id = self._vocab_str_to_int["[PAD]"] self._unk_id = self._vocab_str_to_int["[UNK]"] self._bos_np = np.array([self._bos_id], dtype=np.uint16) self._eos_np = np.array([self._eos_id], dtype=np.uint16) self._numpy_lookup = np.full(256, self._unk_id, dtype=np.uint16) for ch in self.characters: self._numpy_lookup[ord(ch)] = self._vocab_str_to_int[ch] for special in ("[CLS]", "[SEP]", "[BOS]", "[MASK]", "[PAD]", "[RESERVED]", "[UNK]"): token_id = self._vocab_str_to_int[special] if special.startswith("[") and len(special) == 5: # Skip bracketed four-letter tokens from attempting ascii mapping. continue # Explicitly map special token string representations if they are single characters. if len(special) == 1: self._numpy_lookup[ord(special)] = token_id add_prefix_space = kwargs.pop("add_prefix_space", False) padding_side = kwargs.pop("padding_side", "left") super().__init__( bos_token=bos_token, eos_token=eos_token, sep_token=sep_token, cls_token=cls_token, pad_token=pad_token, mask_token=mask_token, unk_token=unk_token, add_prefix_space=add_prefix_space, model_max_length=model_max_length, padding_side=padding_side, **kwargs, ) @property def vocab_size(self) -> int: return len(self._vocab_str_to_int) def _tokenize(self, text: str) -> List[str]: return list(text) def _convert_token_to_id(self, token: str) -> int: return self._vocab_str_to_int.get(token, self._vocab_str_to_int["[UNK]"]) def _convert_id_to_token(self, index: int) -> str: return self._vocab_int_to_str[index] def convert_tokens_to_string(self, tokens): return "".join(tokens) def get_special_tokens_mask( self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False, ) -> List[int]: if already_has_special_tokens: return super().get_special_tokens_mask( token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True, ) result = ([0] * len(token_ids_0)) + [1] if token_ids_1 is not None: result += ([0] * len(token_ids_1)) + [1] return result def build_inputs_with_special_tokens( self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None ) -> List[int]: sep = [self.sep_token_id] bos = [self.bos_token_id] eos = [self.eos_token_id] result = bos + token_ids_0 + eos if token_ids_1 is not None: result += token_ids_1 + eos return result def create_attention_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None) -> List[int]: """Creates an attention mask to differentiate between padding and non-padding tokens. Args: token_ids_0 (List[int]): List of token IDs for the first sequence. token_ids_1 (Optional[List[int]]): List of token IDs for the second sequence if available. Returns: List[int]: A list where 1 represents non-padding tokens and 0 represents padding tokens. """ mask = [1] * len(token_ids_0) if token_ids_1 is not None: mask += [1] * len(token_ids_1) return mask def get_vocab(self) -> Dict[str, int]: return self._vocab_str_to_int def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple: vocab_file = os.path.join(save_directory, (filename_prefix or '') + 'vocab.json') with open(vocab_file, 'w') as f: json.dump(self._vocab_str_to_int, f) return (vocab_file,) def __call__( self, text: Union[str, List[str]], *, padding: bool = True, truncation: bool = True, max_length: Optional[int] = None, add_special_tokens: bool = True, return_tensors: Optional[str] = None, **kwargs, ): # ---------- detect batch vs single ---------- is_batch = not isinstance(text, str) seqs = text if is_batch else [text] # always work on a list internally max_len = max_length or self.model_max_length # ---------- encode every sequence ---------- batch_input_ids = [] for seq in seqs: seq_bytes = np.frombuffer(seq.encode("ascii", "ignore"), dtype=np.uint8) ids = self._numpy_lookup[seq_bytes] if add_special_tokens: ids = np.concatenate((self._bos_np, ids, self._eos_np)) if truncation and ids.size > max_len: ids = ids[:max_len] batch_input_ids.append(ids.astype(np.uint16, copy=False)) # ---------- pad ---------- if padding and batch_input_ids: if padding == "max_length": pad_len = max_len elif padding == "longest": pad_len = max(ids.size for ids in batch_input_ids) elif padding is True: pad_len = max(ids.size for ids in batch_input_ids) else: pad_len = None if pad_len is not None: pad_len = min(pad_len, max_len) padded_ids = [] for ids in batch_input_ids: if ids.size < pad_len: pad_width = pad_len - ids.size ids = np.pad(ids, (0, pad_width), constant_values=self._pad_id) elif ids.size > pad_len: ids = ids[:pad_len] ids = np.asarray(ids, dtype=np.uint16, order="C") padded_ids.append(ids[:pad_len]) batch_input_ids = padded_ids for ids in batch_input_ids: if not isinstance(ids, np.ndarray): # Fallback for any non-numpy path continue # ---------- masks ---------- batch_attention = [] for ids in batch_input_ids: if isinstance(ids, np.ndarray): mask = (ids != self._pad_id).astype(np.uint8, copy=False) else: mask = [0 if tok == self._pad_id else 1 for tok in ids] batch_attention.append(mask) # ---------- collapse back if it was a single example ---------- if not is_batch: batch_input_ids = batch_input_ids[0] batch_attention = batch_attention[0] result = { "input_ids": batch_input_ids, "attention_mask": batch_attention, } # ---------- convert to tensors if requested ---------- if return_tensors == "pt": import torch if is_batch: result["input_ids"] = torch.tensor(np.stack(result["input_ids"]), dtype=torch.long) result["attention_mask"] = torch.tensor(np.stack(result["attention_mask"]), dtype=torch.long) else: result["input_ids"] = torch.tensor(result["input_ids"], dtype=torch.long).unsqueeze(0) result["attention_mask"] = torch.tensor(result["attention_mask"], dtype=torch.long).unsqueeze(0) elif return_tensors == "np": if is_batch: result["input_ids"] = np.stack(result["input_ids"]) result["attention_mask"] = np.stack(result["attention_mask"]) else: result["input_ids"] = np.expand_dims(result["input_ids"], 0) result["attention_mask"] = np.expand_dims(result["attention_mask"], 0) return result