| """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) |
|
|
| |
| 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 |
| |
| 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] |
|
|
