Spaces:
Sleeping
Sleeping
File size: 2,473 Bytes
9fab08e |
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 |
from abc import abstractmethod
from models.tokenizer import TokenAligner
class PTCollator():
def __init__(self, tokenAligner: TokenAligner):
self.tokenAligner = tokenAligner
def collate(self, dataloader_batch, type = "train") -> dict:
if type == "train":
return self.collate_train(dataloader_batch)
elif type == "test":
return self.collate_test(dataloader_batch)
elif type == "correct":
return self.collate_correct(dataloader_batch)
@abstractmethod
def collate_train(self, dataloader_batch):
pass
@abstractmethod
def collate_test(self, dataloader_batch):
pass
@abstractmethod
def collate_correct(self, dataloader_batch):
pass
class DataCollatorForCharacterTransformer(PTCollator):
def __init__(self, tokenAligner: TokenAligner):
super().__init__(tokenAligner)
def collate_train(self, dataloader_batch):
noised, labels = [], []
for sample in dataloader_batch:
labels.append(sample[0])
noised.append(sample[1])
batch_srcs, batch_tgts, batch_lengths, batch_attention_masks = self.tokenAligner.tokenize_for_transformer_with_tokenization(noised, labels)
data = dict()
data['batch_src'] = batch_srcs
data['batch_tgt'] = batch_tgts
data['attn_masks'] = batch_attention_masks
data['lengths'] = batch_lengths
return data
def collate_test(self, dataloader_batch):
noised, labels = [], []
for sample in dataloader_batch:
labels.append(sample[0])
noised.append(sample[1])
batch_srcs, batch_attention_masks = self.tokenAligner.tokenize_for_transformer_with_tokenization(noised, None)
data = dict()
data['batch_src'] = batch_srcs
data['noised_texts'] = noised
data['label_texts'] = labels
data['attn_masks'] = batch_attention_masks
return data
def collate_correct(self, dataloader_batch):
noised, labels = [], []
for sample in dataloader_batch:
noised.append(sample[1])
batch_srcs, batch_attention_masks= self.tokenAligner.tokenize_for_transformer_with_tokenization(noised)
data = dict()
data['batch_src'] = batch_srcs
data['noised_texts'] = noised
data['attn_masks'] = batch_attention_masks
return data
|