INNER CODE UNIT · Python
__init__
thunlp/Few-NERD · model/nnshot.py:11
def __init__(self,word_encoder, dot=False, ignore_index=-1):
util.framework.FewShotNERModel.__init__(self, word_encoder, ignore_index=ignore_index)
self.drop = nn.Dropout()
self.dot = dot
def __dist__(self, x, y, dim):
if self.dot:
return (x * y).sum(dim)
else:
return -(torch.pow(x - y, 2)).sum(dim)
def __batch_dist__(self, S, Q, q_mask):
# S [class, embed_dim], Q [num_of_sent, num_of_tokens, embed_dim]
assert Q.size()[:2] == q_mask.size()
Q = Q[q_mask==1].view(-1, Q.size(-1))
return self.__dist__(S.unsqueeze(0), Q.unsqueeze(1), 2)
def __get_nearest_dist__(self, embedding, tag, mask, query, q_mask):