File size: 6,450 Bytes
6af52a6 5f987d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | import json
import os
from transformers import PreTrainedTokenizer
_RNA_VOCAB = {
"<cls>": 0, "<pad>": 1, "<eos>": 2, "<unk>": 3,
"A": 4, "C": 5, "G": 6, "U": 7,
"R": 8, "Y": 9, "K": 10, "M": 11,
"S": 12, "W": 13, "B": 14, "D": 15,
"H": 16, "V": 17, "N": 18, "-": 19,
"<null_1>": 20, "<null_2>": 21, "<null_3>": 22, "<null_4>": 23,
"<mask>": 24,
}
_MRNA_VOCAB = {
"<cls>": 0, "<pad>": 1, "<eos>": 2, "<unk>": 3,
"GAG": 4, "AAG": 5, "GAA": 6, "CUG": 7, "CAG": 8, "GAU": 9,
"AAA": 10, "GUG": 11, "GAC": 12, "AUG": 13, "GCC": 14, "AAC": 15,
"GCU": 16, "AAU": 17, "AUC": 18, "UUC": 19, "GGA": 20, "AUU": 21,
"GGC": 22, "UUU": 23, "CCA": 24, "AGC": 25, "GCA": 26, "UCU": 27,
"CUC": 28, "ACC": 29, "CAA": 30, "CCU": 31, "UCC": 32, "ACA": 33,
"UUG": 34, "GUU": 35, "CUU": 36, "UAC": 37, "ACU": 38, "CCC": 39,
"UCA": 40, "GUC": 41, "GGU": 42, "CAC": 43, "AGU": 44, "UAU": 45,
"AGA": 46, "CAU": 47, "GGG": 48, "UGG": 49, "UGC": 50, "AGG": 51,
"UGU": 52, "AUA": 53, "CGC": 54, "UUA": 55, "GCG": 56, "CGG": 57,
"CCG": 58, "GUA": 59, "CUA": 60, "ACG": 61, "UCG": 62, "CGA": 63,
"CGU": 64, "UGA": 65, "UAA": 66, "UAG": 67,
"<null_1>": 68, "<null_2>": 69, "<null_3>": 70, "<null_4>": 71,
"<mask>": 72,
}
class RnaFmTokenizer(PreTrainedTokenizer):
vocab_files_names = {"vocab_file": "vocab.json"}
model_input_names = ["input_ids", "attention_mask"]
def __init__(
self,
vocab_file=None,
k_mer: int = 1,
cls_token="<cls>",
pad_token="<pad>",
eos_token="<eos>",
unk_token="<unk>",
mask_token="<mask>",
**kwargs,
):
self.k_mer = k_mer
if vocab_file and os.path.isfile(vocab_file):
with open(vocab_file) as f:
self._vocab = json.load(f)
else:
self._vocab = dict(_MRNA_VOCAB if k_mer == 3 else _RNA_VOCAB)
self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
super().__init__(
cls_token=cls_token,
pad_token=pad_token,
eos_token=eos_token,
unk_token=unk_token,
mask_token=mask_token,
k_mer=k_mer,
**kwargs,
)
@property
def vocab_size(self):
return len(self._vocab)
def get_vocab(self):
return dict(self._vocab)
def _tokenize(self, text):
if self.k_mer == 1:
return list(text)
return [text[i:i + self.k_mer] for i in range(0, len(text), self.k_mer)]
def _convert_token_to_id(self, token):
return self._vocab.get(token, self._vocab["<unk>"])
def _convert_id_to_token(self, index):
return self._ids_to_tokens.get(index, "<unk>")
def save_vocabulary(self, save_directory, filename_prefix=None):
os.makedirs(save_directory, exist_ok=True)
fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
path = os.path.join(save_directory, fname)
with open(path, "w") as f:
json.dump(self._vocab, f, indent=2)
return (path,)
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
cls = [self.cls_token_id]
eos = [self.eos_token_id]
if token_ids_1 is None:
return cls + token_ids_0 + eos
return cls + token_ids_0 + eos + cls + token_ids_1 + eos
def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
if already_has_special_tokens:
return super().get_special_tokens_mask(token_ids_0, token_ids_1, already_has_special_tokens=True)
mask = [1] + [0] * len(token_ids_0) + [1]
if token_ids_1 is not None:
mask += [1] + [0] * len(token_ids_1) + [1]
return mask
def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
if token_ids_1 is None:
return [0] * (len(token_ids_0) + 2)
return [0] * (len(token_ids_0) + 2) + [0] * (len(token_ids_1) + 2)
@staticmethod
def _extract_cds(sequence, cds):
"""Extract CDS region from a sequence, trimmed to a multiple of 3."""
import numpy as np
if sum(cds) == 0:
return sequence[:len(sequence) - (len(sequence) % 3)]
first = int(np.argmax(cds == 1))
last = int(len(cds) - 1 - np.argmax(np.flip(cds) == 1)) + 2
region = sequence[first:last + 1]
if len(region) % 3 != 0:
region = region[:-(len(region) % 3)]
return region
def batch_encode_with_cds(self, sequences, cds, max_length=None, **kwargs):
"""Encode sequences with CDS extraction (k_mer=3 / mRNA-FM only).
Applies T->U, extracts the CDS region, chunks to max_length nucleotides
(aligned to codon boundaries), and encodes each chunk.
Args:
sequences: List of raw nucleotide strings (T or U).
cds: List of numpy arrays marking CDS codon start positions.
max_length: Nucleotide budget per chunk (defaults to
(model_max_length - 2) * k_mer).
**kwargs: Forwarded to batch_encode_plus (e.g. return_tensors,
padding, add_special_tokens).
Returns:
Tuple of (BatchEncoding, chunk_counts) where chunk_counts[i] is the
number of chunks produced for sequences[i].
"""
if self.k_mer != 3:
raise ValueError("batch_encode_with_cds requires k_mer=3 (mRNA-FM tokenizer)")
budget = max_length if max_length is not None else (self.model_max_length - 2) * self.k_mer
budget = (budget // self.k_mer) * self.k_mer
all_chunks = []
chunk_counts = []
for seq, c in zip(sequences, cds):
seq = seq.replace("T", "U").replace("t", "u")
seq = self._extract_cds(seq, c)
raw_chunks = [seq[i:i + budget] for i in range(0, max(len(seq), 1), budget)]
chunks = []
for chunk in raw_chunks:
if len(chunk) % self.k_mer != 0:
chunk = chunk[:-(len(chunk) % self.k_mer)]
if chunk:
chunks.append(chunk)
if not chunks:
chunks = ["AUG"]
all_chunks.extend(chunks)
chunk_counts.append(len(chunks))
enc = self.batch_encode_plus(all_chunks, **kwargs)
return enc, chunk_counts
|