File size: 2,462 Bytes
7cdda9f |
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 |
# 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,)
|