File size: 1,496 Bytes
e42c877 | 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 | import re
from string import punctuation
import torch
from collate_functor import CollateFunctor
class Preprocessor:
def __init__(self, tokenizer):
self.punct_pattern = re.compile(r"([{}])".format(re.escape(punctuation)))
self.space_token = 'Ġ'
self.collator = CollateFunctor(pad_index=0)
self.tokenizer = tokenizer
def preprocess(self, sentences):
sentences = [re.sub(self.punct_pattern, r" \1 ", sentence).strip() for sentence in sentences]
encoding = self.tokenizer(sentences)
batch = []
for encoding in encoding.encodings:
batch.append({"subwords": encoding.ids, "tokens": encoding.tokens})
for sentence in batch:
alignment = [0] # starting with cls token
words = []
current_alignment = 1
current_word = []
for token in sentence["tokens"][1:]:
if (token == self.space_token) or (token == self.tokenizer.eos_token):
current_alignment += 1
words.append(self.tokenizer.convert_tokens_to_string(current_word))
current_word = []
else:
current_word.append(token)
alignment.append(current_alignment)
sentence["alignment"] = torch.LongTensor(alignment)
sentence["subwords"] = torch.LongTensor(sentence["subwords"])
sentence["words"] = words
return self.collator(batch)
|