File size: 2,127 Bytes
c6ee1cf | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import re
from collections import Counter
class Tokenizer:
def __init__(self, min_freq = 5):
self.min_freq = min_freq
self.word_to_idx = {}
self.idx_to_word = {}
self.special_tokens = {
"<pad>" : 0,
"<unk>" : 1,
"<bos>" : 2,
"<eos>" : 3,
}
def __len__(self):
return len(self.word_to_idx)
def tokenize(self, text):
# raw text -> list of words
text = text.lower()
# split words and punctuations
tokens = re.findall(r"\w+|[^\w\s]", text)
return tokens
def build_vocab(self, texts):
# create word -> id mapping
counter = Counter()
for text in texts:
tokens = self.tokenize(text)
counter.update(tokens)
self.word_to_idx = self.special_tokens.copy()
for word, freq in sorted(counter.items()):
if freq >= self.min_freq:
if word not in self.word_to_idx:
self.word_to_idx[word] = len(self.word_to_idx)
self.idx_to_word = {
idx: word for word, idx in self.word_to_idx.items()
}
def encode(self, text):
# text -> ids
text = text.lower()
tokens = self.tokenize(text)
ids = [
self.word_to_idx.get(
token,
self.word_to_idx["<unk>"]
)
for token in tokens
]
ids = (
[self.word_to_idx["<bos>"]]
+ ids
+[self.word_to_idx["<eos>"]]
)
return ids
def decode(self, ids):
words = []
for idx in ids:
word = self.idx_to_word.get(
idx,
"<unk>"
)
if word not in [
"<pad>",
"<bos>",
"<eos>",
"<unk>"
]:
words.append(word)
return " ".join(words) |