def __init__(self, params, bert_layer, decoder_layer, name=None):
super(NHNet, self).__init__(params,
bert_layer,
decoder_layer,
name=name)
self.doc_attention = multi_channel_attention.VotingAttention(
num_heads=params.num_decoder_attn_heads,
head_size=params.hidden_size // params.num_decoder_attn_heads)
def test_doc_attention(self):
num_heads = 2
doc_attention = multi_channel_attention.VotingAttention(num_heads,
head_size=8)
num_docs = 3
inputs = np.zeros((2, num_docs, 10, 16), dtype=np.float32)
doc_mask = np.zeros((2, num_docs), dtype=np.float32)
outputs = doc_attention(inputs, doc_mask)
self.assertEqual(outputs.shape, (2, num_docs))
