Fixing encoder and tokenize functions.

tokenizer.py

@@ -1,5 +1,6 @@
 from transformers import PreTrainedTokenizer
 from huggingface_hub import hf_hub_download
+import torch
 import json
 import os
 from itertools import product
@@ -25,15 +26,24 @@ class KmerTokenizer(PreTrainedTokenizer):
         self.unk_token = "[UNK]"
         # self.pad_token = "[PAD]"
 
-    def _tokenize(self, text):
+    def _tokenize(self, text, **kwargs):
         splits = [text[i:i + self.k] for i in range(0, len(text) - self.k + 1, self.stride)]
-        return self.convert_tokens_to_ids(splits)
+        if kwargs.get('return_tensors') == 'pt':
+            return torch.tensor(splits)
+        return splits
+
+    def _encode(self, text, **kwargs):
+        tokens = self._tokenize(text, **kwargs)
+        token_ids = self.convert_tokens_to_ids(tokens)
+        if kwargs.get('return_tensors') == 'pt':
+            return torch.tensor(token_ids)
+        return token_ids
 
     def convert_tokens_to_ids(self, tokens):
         unk_id = self.vocab_dict.get(self.unk_token)
         return [self.vocab_dict[token] if token in self.vocab_dict else unk_id for token in tokens]
 
-    def convert_ids_to_tokens(self, ids):
+    def convert_ids_to_tokens(self, ids, **kwargs):
         id_to_token = {idx: token for token, idx in self.vocab_dict.items()}
         return [id_to_token.get(id_, self.unk_token) for id_ in ids]
 
@@ -58,21 +68,6 @@ class KmerTokenizer(PreTrainedTokenizer):
                     "k": self.k,
                     "stride": self.stride
                 },
-                # "post_processor": {
-                #     "type": "TemplateProcessing",
-                #     "single": [
-                #         {"SpecialToken": {"id": self.cls_token, "type_id": 0}},
-                #         {"Sequence": {"id": "A", "type_id": 0}},
-                #         {"SpecialToken": {"id": self.sep_token, "type_id": 0}}
-                #     ],
-                #     "pair": [
-                #         {"SpecialToken": {"id": self.cls_token, "type_id": 0}},
-                #         {"Sequence": {"id": "A", "type_id": 0}},
-                #         {"SpecialToken": {"id": self.sep_token, "type_id": 0}},
-                #         {"Sequence": {"id": "B", "type_id": 1}},
-                #         {"SpecialToken": {"id": self.sep_token, "type_id": 1}}
-                #     ]
-                # }
                 "model": {
                     "type": "k-mer",
                     "k": self.k,