Compute SQuAD end_logits from sequence hidden states and start token hidden state.
| 856 | |
| 857 | |
| 858 | class PoolerEndLogits(nn.Module): |
| 859 | """ Compute SQuAD end_logits from sequence hidden states and start token hidden state. |
| 860 | """ |
| 861 | |
| 862 | def __init__(self, config): |
| 863 | super().__init__() |
| 864 | self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) |
| 865 | self.activation = nn.Tanh() |
| 866 | self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) |
| 867 | self.dense_1 = nn.Linear(config.hidden_size, 1) |
| 868 | |
| 869 | def forward(self, hidden_states, start_states=None, start_positions=None, p_mask=None): |
| 870 | """ Args: |
| 871 | One of ``start_states``, ``start_positions`` should be not None. |
| 872 | If both are set, ``start_positions`` overrides ``start_states``. |
| 873 | |
| 874 | **start_states**: ``torch.LongTensor`` of shape identical to hidden_states |
| 875 | hidden states of the first tokens for the labeled span. |
| 876 | **start_positions**: ``torch.LongTensor`` of shape ``(batch_size,)`` |
| 877 | position of the first token for the labeled span: |
| 878 | **p_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, seq_len)`` |
| 879 | Mask of invalid position such as query and special symbols (PAD, SEP, CLS) |
| 880 | 1.0 means token should be masked. |
| 881 | """ |
| 882 | assert ( |
| 883 | start_states is not None or start_positions is not None |
| 884 | ), "One of start_states, start_positions should be not None" |
| 885 | if start_positions is not None: |
| 886 | slen, hsz = hidden_states.shape[-2:] |
| 887 | start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) |
| 888 | start_states = hidden_states.gather(-2, start_positions) # shape (bsz, 1, hsz) |
| 889 | start_states = start_states.expand(-1, slen, -1) # shape (bsz, slen, hsz) |
| 890 | |
| 891 | x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1)) |
| 892 | x = self.activation(x) |
| 893 | x = self.LayerNorm(x) |
| 894 | x = self.dense_1(x).squeeze(-1) |
| 895 | |
| 896 | if p_mask is not None: |
| 897 | if next(self.parameters()).dtype == torch.float16: |
| 898 | x = x * (1 - p_mask) - 65500 * p_mask |
| 899 | else: |
| 900 | x = x * (1 - p_mask) - 1e30 * p_mask |
| 901 | |
| 902 | return x |
| 903 | |
| 904 | |
| 905 | class PoolerAnswerClass(nn.Module): |