Create tokenizer_i3.py

tokenizer_i3.py

@@ -0,0 +1,76 @@
+# tokenizer_i3.py
+import os
+import json
+from transformers import PreTrainedTokenizer
+
+from i3_model import ChunkTokenizer
+
+# ======================================================================
+# HuggingFace Tokenizer Wrapper for ChunkTokenizer
+# ======================================================================
+class I3Tokenizer(PreTrainedTokenizer):
+    """
+    HuggingFace-compatible tokenizer for i3 model using ChunkTokenizer.
+    """
+
+    vocab_files_names = {"vocab_file": "chunk_vocab_combined.json"}
+    pretrained_vocab_files_map = {}
+    max_model_input_sizes = {"i3": 512}
+
+    def __init__(self, vocab_file=None, **kwargs):
+        """
+        Args:
+            vocab_file: Path to chunk_vocab_combined.json
+        """
+        super().__init__(**kwargs)
+        self.chunk_tokenizer = ChunkTokenizer()
+        if vocab_file:
+            self.chunk_tokenizer.load(vocab_file)
+        self.vocab_file = vocab_file
+
+    @property
+    def vocab_size(self):
+        return self.chunk_tokenizer.vocab_size
+
+    def _tokenize(self, text, **kwargs):
+        """
+        Convert text string to list of token strings (chunks).
+        """
+        # Encode to indices, then convert back to chunk strings
+        indices = self.chunk_tokenizer.encode(text)
+        tokens = [self.chunk_tokenizer.idx_to_chunk[i] for i in indices]
+        return tokens
+
+    def _convert_token_to_id(self, token):
+        """
+        Convert chunk string to integer ID.
+        """
+        return self.chunk_tokenizer.chunk_to_idx.get(token, self.chunk_tokenizer.unk_idx)
+
+    def _convert_id_to_token(self, index):
+        """
+        Convert integer ID to chunk string.
+        """
+        return self.chunk_tokenizer.idx_to_chunk.get(int(index), self.chunk_tokenizer.unk_token)
+
+    def encode(self, text, **kwargs):
+        """
+        Convert text string to list of indices.
+        """
+        return self.chunk_tokenizer.encode(text)
+
+    def decode(self, token_ids, **kwargs):
+        """
+        Convert list of indices back to text string.
+        """
+        return self.chunk_tokenizer.decode(token_ids)
+
+    def save_vocabulary(self, save_directory):
+        """
+        Save the vocabulary to a directory.
+        """
+        if not os.path.exists(save_directory):
+            os.makedirs(save_directory)
+        save_path = os.path.join(save_directory, "chunk_vocab_combined.json")
+        self.chunk_tokenizer.save(save_path)
+        return (save_path,)