data/field/anchor_field.py

#!/usr/bin/env python3
# coding=utf-8

import torch
from data.field.mini_torchtext.field import RawField


class AnchorField(RawField):
    def process(self, batch, device=None):
        tensors, masks = self.pad(batch, device)
        return tensors, masks

    def pad(self, anchors, device):
        tensor = torch.zeros(anchors[0], anchors[1], dtype=torch.long, device=device)
        for anchor in anchors[-1]:
            tensor[anchor[0], anchor[1]] = 1
        mask = tensor.sum(-1) == 0

        return tensor, mask