Create tokenizer.py

tokenizer.py

@@ -0,0 +1,46 @@
+# tokenizer.py
+
+from transformers import PreTrainedTokenizer
+from itertools import product
+import json
+import os
+
+class CodonTokenizer(PreTrainedTokenizer):
+    def __init__(self, **kwargs):
+        bases = ['A', 'T', 'G', 'C']
+        codons = [''.join(p) for p in product(bases, repeat=3)]
+        special_tokens = ['[PAD]', '[BOS]', '[EOS]']
+        self.vocab_list = special_tokens + codons
+
+        self.codon2id = {token: idx for idx, token in enumerate(self.vocab_list)}
+        self.id2codon = {idx: token for token, idx in self.codon2id.items()}
+
+        kwargs['bos_token'] = '[BOS]'
+        kwargs['eos_token'] = '[EOS]'
+        kwargs['pad_token'] = '[PAD]'
+        super().__init__(**kwargs)
+
+    def _tokenize(self, text):
+        return [text[i:i+3] for i in range(0, len(text), 3)]
+
+    def _convert_token_to_id(self, token):
+        return self.codon2id.get(token, self.codon2id['[PAD]'])
+
+    def _convert_id_to_token(self, idx):
+        return self.id2codon.get(idx, '[PAD]')
+
+    def convert_tokens_to_string(self, tokens):
+        return ''.join(tokens)
+
+    def get_vocab(self):
+        return self.codon2id
+
+    @property
+    def vocab_size(self):
+        return len(self.codon2id)
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        path = os.path.join(save_directory, (filename_prefix or "") + "vocab.json")
+        with open(path, "w") as f:
+            json.dump(self.codon2id, f)
+        return (path,)