Compute SQuAD 2.0 answer class from classification and start tokens hidden states.
| 903 | |
| 904 | |
| 905 | class PoolerAnswerClass(nn.Module): |
| 906 | """ Compute SQuAD 2.0 answer class from classification and start tokens hidden states. """ |
| 907 | |
| 908 | def __init__(self, config): |
| 909 | super().__init__() |
| 910 | self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) |
| 911 | self.activation = nn.Tanh() |
| 912 | self.dense_1 = nn.Linear(config.hidden_size, 1, bias=False) |
| 913 | |
| 914 | def forward(self, hidden_states, start_states=None, start_positions=None, cls_index=None): |
| 915 | """ |
| 916 | Args: |
| 917 | One of ``start_states``, ``start_positions`` should be not None. |
| 918 | If both are set, ``start_positions`` overrides ``start_states``. |
| 919 | |
| 920 | **start_states**: ``torch.LongTensor`` of shape identical to ``hidden_states``. |
| 921 | hidden states of the first tokens for the labeled span. |
| 922 | **start_positions**: ``torch.LongTensor`` of shape ``(batch_size,)`` |
| 923 | position of the first token for the labeled span. |
| 924 | **cls_index**: torch.LongTensor of shape ``(batch_size,)`` |
| 925 | position of the CLS token. If None, take the last token. |
| 926 | |
| 927 | note(Original repo): |
| 928 | no dependency on end_feature so that we can obtain one single `cls_logits` |
| 929 | for each sample |
| 930 | """ |
| 931 | hsz = hidden_states.shape[-1] |
| 932 | assert ( |
| 933 | start_states is not None or start_positions is not None |
| 934 | ), "One of start_states, start_positions should be not None" |
| 935 | if start_positions is not None: |
| 936 | start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) |
| 937 | start_states = hidden_states.gather(-2, start_positions).squeeze(-2) # shape (bsz, hsz) |
| 938 | |
| 939 | if cls_index is not None: |
| 940 | cls_index = cls_index[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) |
| 941 | cls_token_state = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, hsz) |
| 942 | else: |
| 943 | cls_token_state = hidden_states[:, -1, :] # shape (bsz, hsz) |
| 944 | |
| 945 | x = self.dense_0(torch.cat([start_states, cls_token_state], dim=-1)) |
| 946 | x = self.activation(x) |
| 947 | x = self.dense_1(x).squeeze(-1) |
| 948 | |
| 949 | return x |
| 950 | |
| 951 | |
| 952 | class SQuADHead(nn.Module): |