|
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 |
|
|