| import torch | |
| from data.field.mini_torchtext.field import RawField | |
| class AnchoredLabelField(RawField): | |
| def __init__(self): | |
| super(AnchoredLabelField, self).__init__() | |
| self.vocab = None | |
| def process(self, example, device=None): | |
| example = self.numericalize(example) | |
| tensor = self.pad(example, device) | |
| return tensor | |
| def pad(self, example, device): | |
| n_labels = len(self.vocab) | |
| n_nodes, n_tokens = len(example[1]), example[0] | |
| tensor = torch.full([n_nodes, n_tokens, n_labels + 1], 0, dtype=torch.long, device=device) | |
| for i_node, node in enumerate(example[1]): | |
| for anchor, rule in node: | |
| tensor[i_node, anchor, rule + 1] = 1 | |
| return tensor | |
| def numericalize(self, arr): | |
| def multi_map(array, function): | |
| if isinstance(array, tuple): | |
| return (array[0], function(array[1])) | |
| elif isinstance(array, list): | |
| return [multi_map(a, function) for a in array] | |
| else: | |
| return array | |
| if self.vocab is not None: | |
| arr = multi_map(arr, lambda x: self.vocab.stoi[x] if x in self.vocab.stoi else 0) | |
| return arr | |