INNER CODE UNIT · Python

get_geometric_data_standard

IBM/transition-amr-parser · src/ibm_neural_aligner/main.py:517

    def get_geometric_data_standard(self, amr):
        vocab = self.amr_tokenizer.vocab

        node_ids = get_node_ids(amr)
        node_TO_idx = {k: i for i, k in enumerate(node_ids)}

        edge_index = []
        for xin, label, xout in amr.edges:
            a = node_TO_idx[xin]
            b = node_TO_idx[xout]
            edge_index.append([a, b])
            edge_index.append([b, a])

        node_labels = [amr.nodes[k] for k in node_ids]
        node_tokens = torch.tensor([vocab[tok] for tok in node_labels], dtype=torch.long)
        edge_index = torch.tensor(edge_index, dtype=torch.long)
        data = Data(edge_index=edge_index.t().contiguous(), y=node_tokens, num_nodes=len(node_ids))
        return data

View source record →