Compute SQuAD start_logits from sequence hidden states.
| 832 | |
| 833 | |
| 834 | class PoolerStartLogits(nn.Module): |
| 835 | """ Compute SQuAD start_logits from sequence hidden states. """ |
| 836 | |
| 837 | def __init__(self, config): |
| 838 | super().__init__() |
| 839 | self.dense = nn.Linear(config.hidden_size, 1) |
| 840 | |
| 841 | def forward(self, hidden_states, p_mask=None): |
| 842 | """ Args: |
| 843 | **p_mask**: (`optional`) ``torch.FloatTensor`` of shape `(batch_size, seq_len)` |
| 844 | invalid position mask such as query and special symbols (PAD, SEP, CLS) |
| 845 | 1.0 means token should be masked. |
| 846 | """ |
| 847 | x = self.dense(hidden_states).squeeze(-1) |
| 848 | |
| 849 | if p_mask is not None: |
| 850 | if next(self.parameters()).dtype == torch.float16: |
| 851 | x = x * (1 - p_mask) - 65500 * p_mask |
| 852 | else: |
| 853 | x = x * (1 - p_mask) - 1e30 * p_mask |
| 854 | |
| 855 | return x |
| 856 | |
| 857 | |
| 858 | class PoolerEndLogits(nn.Module): |