|
import torch, json |
|
|
|
class CharTokenizer: |
|
def __init__(self, corpus=None, vocab=None): |
|
if vocab is not None: |
|
self.vocab = vocab |
|
elif corpus is not None: |
|
self.vocab = self._build_vocab(corpus) |
|
else: |
|
raise Exception("Either corpus or vocab has to be supplied") |
|
self.id2vocab = [char for char, index in sorted(self.vocab.items(), key=lambda item: item[1])] |
|
|
|
def _tokenize(self, text): |
|
return list(text) |
|
|
|
def __call__(self, prompt, text=None, add_eos_token=False): |
|
token_ids = [self.vocab.get(token, 0) for token in self._tokenize(prompt)] |
|
if text is not None: |
|
text_token_ids = [self.vocab.get(token, 0) for token in self._tokenize(text)] |
|
token_ids = token_ids + [self.vocab["<bos>"]] + text_token_ids |
|
if add_eos_token: |
|
token_ids = token_ids + [self.vocab["<eos>"]] |
|
input_ids_tensor = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0) |
|
attention_masks = torch.ones_like(input_ids_tensor) |
|
return {"input_ids": input_ids_tensor, "attention_mask": attention_masks} |
|
|
|
def _build_vocab(self, corpus): |
|
vocab = {"<pad>": 0} |
|
for verse_lengths in range(3, 10): |
|
vocab[str(verse_lengths)] = len(vocab) |
|
for doc in corpus: |
|
chars = self._tokenize(doc) |
|
for char in chars: |
|
if char not in vocab: |
|
vocab[char] = len(vocab) |
|
vocab["<bos>"] = len(vocab) |
|
vocab["<eos>"] = len(vocab) |
|
return vocab |
|
|
|
def decode(self, token_ids): |
|
chars = [self.id2vocab[token_id] for token_id in token_ids.flatten().tolist()] |
|
filtered_chars = [char for char in chars if char not in ["<eos>", "<bos>", "<pad>"]] |
|
return "".join(filtered_chars) |
|
|
|
def save(self, filepath): |
|
with open(filepath, "w") as f: |
|
json.dump(self.vocab, f) |
|
|
|
@classmethod |
|
def load(cls, filepath): |
|
with open(filepath) as f: |
|
vocab = json.load(f) |
|
return cls(vocab=vocab) |
|
|
