File size: 4,297 Bytes
9801f83 | 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 | """Wrapper for the canonical DNAMotifTokenizer (HF: Anonymous-843q0u4q08/DNAMotifTokenizer).
Exposes a `tokenizers.Tokenizer`-compatible interface (.encode, .get_vocab) so
motif_coverage_eval.py can consume it like the other tokenizers.
"""
import itertools
import pickle
import random
from os.path import join
from pathlib import Path
class TrieNode:
def __init__(self):
self.children = {}
self.is_end_of_word = False
self.features = []
class Trie:
def __init__(self):
self.root = TrieNode()
self.lookup_table = {}
def insert(self, word, features=None):
current_node = self.root
for char in word:
if char not in current_node.children:
current_node.children[char] = TrieNode()
current_node = current_node.children[char]
current_node.is_end_of_word = True
if features:
current_node.features.append(features)
def search(self, word):
current_node = self.root
for char in word:
if char not in current_node.children:
return False
current_node = current_node.children[char]
if current_node.is_end_of_word:
if len(current_node.features) > 0:
return current_node.features
return True
return False
class _Encoding:
def __init__(self, tokens):
self.tokens = tokens
self.offsets = None
class DNAMotifTokenizer:
"""Canonical greedy-longest-match tokenizer with +1/+2 lookahead."""
MAXLEN = 12
def __init__(self, tokenizer_dir, seed=0):
tokenizer_dir = Path(tokenizer_dir)
with open(tokenizer_dir / "motifs_hardcode_trie.pkl", "rb") as f:
self.trie = pickle.load(f)
self.k1 = ["A", "T", "C", "G", "N"]
self.k3 = ["".join(t) for t in itertools.product("ATCG", repeat=3)]
self.lookup = {}
with open(tokenizer_dir / "motifs_dedup.txt") as f:
for line in f:
seg, name = line.strip().split(maxsplit=1)
self.lookup[seg] = name
self.vocab = {}
with open(tokenizer_dir / "vocab_dedup.txt") as f:
for i, line in enumerate(f):
tok = line.strip()
if tok:
self.vocab[tok] = i
self._rng = random.Random(seed)
def get_vocab(self):
return dict(self.vocab)
def get_vocab_size(self):
return len(self.vocab)
def _tokenize_at(self, seg, i):
candidates = []
for l in range(4, self.MAXLEN + 1):
segment = seg[i:i + l]
if self.trie.search(segment):
candidates.append(segment)
if candidates:
best = self._rng.choice(candidates)
return best, len(best)
# Fallback: 3-mer then 1-mer
for l in range(3, 0, -1):
segment = seg[i:i + l]
if l == 3 and segment in self.k3:
return segment, 3
if l == 1 and segment in self.k1:
return segment, 1
# Last resort: emit single char even if unknown (e.g. 'N' outside k1 — but k1 has N)
return seg[i:i + 1], 1
def _tokenize_seq(self, seg):
tokens = []
i = 0
while i < len(seg):
best_token, best_score = self._tokenize_at(seg, i)
best_i = i
next_pos = i + len(best_token)
offsets = [1, 2] if len(best_token) > 1 else []
for shift in offsets:
i_shifted = i + shift
if i_shifted < len(seg):
tok_, score_ = self._tokenize_at(seg, i_shifted)
next_pos_ = i_shifted + len(tok_)
if score_ > best_score:
best_token, best_score, best_i, next_pos = tok_, score_, i_shifted, next_pos_
for skip in range(best_i - i):
tokens.append(seg[i + skip])
tokens.append(best_token)
i = next_pos
return tokens
def encode(self, text, add_special_tokens=False):
return _Encoding(self._tokenize_seq(text))
def encode_batch(self, texts, add_special_tokens=False):
return [_Encoding(self._tokenize_seq(t)) for t in texts]
|