INNER CODE UNIT · Python
S
thunlp/Few-NERD · model/nnshot.py:30
S = embedding[mask==1].view(-1, embedding.size(-1))
tag = torch.cat(tag, 0)
assert tag.size(0) == S.size(0)
dist = self.__batch_dist__(S, query, q_mask) # [num_of_query_tokens, num_of_support_tokens]
for label in range(torch.max(tag)+1):
nearest_dist.append(torch.max(dist[:,tag==label], 1)[0])
nearest_dist = torch.stack(nearest_dist, dim=1) # [num_of_query_tokens, class_num]
return nearest_dist
def forward(self, support, query):
'''
support: Inputs of the support set.
query: Inputs of the query set.
N: Num of classes
K: Num of instances for each class in the support set
Q: Num of instances in the query set
'''
support_emb = self.word_encoder(support['word'], support['mask']) # [num_sent, number_of_tokens, 768]