| 41 | return (start_logits, end_logits, relevance_logits) |
| 42 | |
| 43 | class DPRTypeMixin(BaseMixin): |
| 44 | def __init__(self, num_types, hidden_size): |
| 45 | super().__init__() |
| 46 | self.type_embeddings = nn.Embedding(num_types, hidden_size) |
| 47 | |
| 48 | def word_embedding_forward(self, input_ids, **kwargs): |
| 49 | print("DPRTypeMixin word_embedding_forward") |
| 50 | if "token_type_ids" in kwargs: |
| 51 | token_type_ids = kwargs["token_type_ids"] |
| 52 | else: |
| 53 | token_type_ids = torch.zeros(input_ids.shape, dtype=torch.long).to(input_ids.device) |
| 54 | return self.transformer.word_embeddings(input_ids) + self.type_embeddings(token_type_ids) |
| 55 | |
| 56 | class DPRQuestionEncoder(BaseModel): |
| 57 | def __init__(self, args, transformer=None, **kwargs): |