the fsner model implements a fewshot named entity recognition method from the paper examplebased named entity recognition https arxiv orgabs2008 10570 by morteza ziyadi yuting sun abhishek goswami jade huang weizhu chen to identify entity spans in a new domain it uses a trainfree fewshot learning approach inspired by questionanswering find scores of each token being start and end token for an entity args wquery torch longtensor of shape batchsize sequencelength indices of query sequence tokens in the vocabulary wsupports torch longtensor of shape batchsize sequencelength indices of support sequence tokens in the vocabulary returns pstart torch floattensor of shape batchsize sequencelength scores of each token as being start token of an entity pend torch floattensor of shape batchsize sequencelength scores of each token as being end token of an entity the fsner model implements a few shot named entity recognition method from the paper example based named entity recognition https arxiv org abs 2008 10570 __ by morteza ziyadi yuting sun abhishek goswami jade huang weizhu chen to identify entity spans in a new domain it uses a train free few shot learning approach inspired by question answering find scores of each token being start and end token for an entity args w_query torch longtensor of shape batch_size sequence_length indices of query sequence tokens in the vocabulary w_supports torch longtensor of shape batch_size sequence_length indices of support sequence tokens in the vocabulary returns p_start torch floattensor of shape batch_size sequence_length scores of each token as being start token of an entity p_end torch floattensor of shape batch_size sequence_length scores of each token as being end token of an entity
import torch

from transformers import AutoModel


class FSNERModel(torch.nn.Module):
    

    def __init__(self, pretrained_model_name_or_path="sayef/fsner-bert-base-uncased"):
        super(FSNERModel, self).__init__()

        self.bert = AutoModel.from_pretrained(pretrained_model_name_or_path, return_dict=True)
        self.cos = torch.nn.CosineSimilarity(3, 1e-08)
        self.softmax = torch.nn.Softmax(dim=1)

    def BERT(self, **inputs):
        return self.bert(**inputs).last_hidden_state

    def VectorSum(self, token_embeddings):
        return token_embeddings.sum(2, keepdim=True)

    def Atten(self, q_rep, S_rep, T=1):
        return self.softmax(T * self.cos(q_rep, S_rep))

    def forward(self, W_query, W_supports):
        

        support_sizes = W_supports["sizes"].tolist()
        start_token_id = W_supports["start_token_id"].item()
        end_token_id = W_supports["end_token_id"].item()

        del W_supports["sizes"]
        del W_supports["start_token_id"]
        del W_supports["end_token_id"]

        q = self.BERT(**W_query)
        S = self.BERT(**W_supports)

        p_starts = None
        p_ends = None

        start_token_masks = W_supports["input_ids"] == start_token_id
        end_token_masks = W_supports["input_ids"] == end_token_id

        for i, size in enumerate(support_sizes):
            if i == 0:
                s = 0
            else:
                s = support_sizes[i - 1]

            s_start = S[s : s + size][start_token_masks[s : s + size]]
            s_end = S[s : s + size][end_token_masks[s : s + size]]

            p_start = torch.matmul(q[i], s_start.T).sum(1).softmax(0)
            p_end = torch.matmul(q[i], s_end.T).sum(1).softmax(0)

            if p_starts is not None:
                p_starts = torch.vstack((p_starts, p_start))
                p_ends = torch.vstack((p_ends, p_end))
            else:
                p_starts = p_start
                p_ends = p_end

        return p_starts, p_ends