| | import torch |
| | from transformers import BertTokenizerFast |
| | from pt_variety_identifier.src.data import Data as BaseData |
| | from torch.utils.data import DataLoader |
| |
|
| |
|
| | class Data(BaseData): |
| | def __init__(self, dataset_name, tokenizer_name, batch_size, test_set_list): |
| | super().__init__(dataset_name=dataset_name, test_set_list=test_set_list) |
| |
|
| | self.tokenizer_name = tokenizer_name |
| | self.tokenizer = BertTokenizerFast.from_pretrained(self.tokenizer_name) |
| | self.batch_size = batch_size |
| |
|
| | def _tokenize(self, example): |
| | return self.tokenizer(example['text'], padding='max_length', truncation=True, max_length=512) |
| |
|
| | def _adapt_dataset(self, dataset): |
| | dataset = dataset.map(self._tokenize, batched=True) |
| |
|
| | |
| | dataset.set_format(type='torch', columns=[ |
| | 'input_ids', 'attention_mask', 'label']) |
| |
|
| | return DataLoader(dataset, batch_size=self.batch_size, shuffle=True) |
| |
|
| | def load_domain(self, domain, balance, pos_prob, ner_prob, sample_size=None): |
| | dataset = super().load_domain(domain=domain, balance=balance, |
| | pos_prob=pos_prob, ner_prob=ner_prob, sample_size=sample_size) |
| |
|
| | return self._adapt_dataset(dataset) |
| |
|
| | def load_validation_set(self): |
| | dataset_dict = super().load_validation_set() |
| |
|
| | for domain in dataset_dict.keys(): |
| | dataset_dict[domain] = self._adapt_dataset(dataset_dict[domain]) |
| |
|
| | return dataset_dict |
| |
|
| | def load_test_set(self, filter_label_2=False): |
| | dataset_dict = super().load_test_set(filter_label_2) |
| |
|
| | for test_set in dataset_dict.keys(): |
| | dataset_dict[test_set] = self._adapt_dataset( |
| | dataset_dict[test_set]) |
| |
|
| | validation_dataset_dict = self.load_validation_set() |
| |
|
| | for val_set in validation_dataset_dict.keys(): |
| | dataset_dict[val_set] = validation_dataset_dict[val_set] |
| |
|
| | return dataset_dict |
| |
|
