| import pickle | |
| from pathlib import Path | |
| import numpy as np | |
| from utils.train_util import pad_sequence | |
| class DictTokenizer: | |
| def __init__(self, | |
| tokenizer_path: str = None, | |
| max_length: int = 20) -> None: | |
| self.word2idx = {} | |
| self.idx2word = {} | |
| self.idx = 0 | |
| self.add_word("<pad>") | |
| self.add_word("<start>") | |
| self.add_word("<end>") | |
| self.add_word("<unk>") | |
| if tokenizer_path is not None and Path(tokenizer_path).exists(): | |
| state_dict = pickle.load(open(tokenizer_path, "rb")) | |
| self.load_state_dict(state_dict) | |
| self.loaded = True | |
| else: | |
| self.loaded = False | |
| self.bos, self.eos = self.word2idx["<start>"], self.word2idx["<end>"] | |
| self.pad = self.word2idx["<pad>"] | |
| self.max_length = max_length | |
| def add_word(self, word): | |
| if not word in self.word2idx: | |
| self.word2idx[word] = self.idx | |
| self.idx2word[self.idx] = word | |
| self.idx += 1 | |
| def encode_word(self, word): | |
| if word in self.word2idx: | |
| return self.word2idx[word] | |
| else: | |
| return self.word2idx["<unk>"] | |
| def __call__(self, texts): | |
| assert isinstance(texts, list), "the input must be List[str]" | |
| batch_tokens = [] | |
| for text in texts: | |
| tokens = [self.encode_word(token) for token in text.split()][:self.max_length] | |
| tokens = [self.bos] + tokens + [self.eos] | |
| tokens = np.array(tokens) | |
| batch_tokens.append(tokens) | |
| caps, cap_lens = pad_sequence(batch_tokens, self.pad) | |
| return { | |
| "cap": caps, | |
| "cap_len": cap_lens | |
| } | |
| def decode(self, batch_token_ids): | |
| output = [] | |
| for token_ids in batch_token_ids: | |
| tokens = [] | |
| for token_id in token_ids: | |
| if token_id == self.eos: | |
| break | |
| elif token_id == self.bos: | |
| continue | |
| tokens.append(self.idx2word[token_id]) | |
| output.append(" ".join(tokens)) | |
| return output | |
| def __len__(self): | |
| return len(self.word2idx) | |
| def state_dict(self): | |
| return self.word2idx | |
| def load_state_dict(self, state_dict): | |
| self.word2idx = state_dict | |
| self.idx2word = {idx: word for word, idx in self.word2idx.items()} | |
| self.idx = len(self.word2idx) | |
| class HuggingfaceTokenizer: | |
| def __init__(self, | |
| model_name_or_path, | |
| max_length) -> None: | |
| from transformers import AutoTokenizer | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) | |
| self.max_length = max_length | |
| self.bos, self.eos = self.tokenizer.bos_token_id, self.tokenizer.eos_token_id | |
| self.pad = self.tokenizer.pad_token_id | |
| self.loaded = True | |
| def __call__(self, texts): | |
| assert isinstance(texts, list), "the input must be List[str]" | |
| batch_token_dict = self.tokenizer(texts, | |
| padding=True, | |
| truncation=True, | |
| max_length=self.max_length, | |
| return_tensors="pt") | |
| batch_token_dict["cap"] = batch_token_dict["input_ids"] | |
| cap_lens = batch_token_dict["attention_mask"].sum(dim=1) | |
| cap_lens = cap_lens.numpy().astype(np.int32) | |
| batch_token_dict["cap_len"] = cap_lens | |
| return batch_token_dict | |
| def decode(self, batch_token_ids): | |
| return self.tokenizer.batch_decode(batch_token_ids, skip_special_tokens=True) | |