Bert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a softmax) e.g. for RocStories/SWAG tasks.
| 581 | |
| 582 | |
| 583 | class BertForMultipleChoice(BertPreTrainedModel): |
| 584 | """ |
| 585 | Bert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a |
| 586 | softmax) e.g. for RocStories/SWAG tasks. |
| 587 | """ |
| 588 | |
| 589 | def __init__(self, config): |
| 590 | super().__init__(config) |
| 591 | self.num_labels = config.num_labels |
| 592 | self.config = config |
| 593 | |
| 594 | self.bert = BertModel(config) |
| 595 | classifier_dropout = ( |
| 596 | config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob |
| 597 | ) |
| 598 | self.dropout = nn.Dropout(classifier_dropout) |
| 599 | |
| 600 | # In multiple choice tasks, all choices are submitted in a batch, and |
| 601 | # we compute a logit for each option independently. The logits are then |
| 602 | # normalized in the forward pass to get a probability distribution over |
| 603 | # the choices. |
| 604 | self.classifier = nn.Linear(config.hidden_size, 1) |
| 605 | |
| 606 | # Initialize weights and apply final processing |
| 607 | self.post_init() |
| 608 | |
| 609 | @classmethod |
| 610 | def from_composer( |
| 611 | cls, |
| 612 | pretrained_checkpoint, |
| 613 | state_dict=None, |
| 614 | cache_dir=None, |
| 615 | from_tf=False, |
| 616 | config=None, |
| 617 | *inputs, |
| 618 | **kwargs, |
| 619 | ): |
| 620 | """Load from pre-trained.""" |
| 621 | model = cls(config, *inputs, **kwargs) |
| 622 | if from_tf: |
| 623 | raise ValueError("Mosaic BERT does not support loading TensorFlow weights.") |
| 624 | |
| 625 | state_dict = torch.load(pretrained_checkpoint) |
| 626 | # If the state_dict was saved after wrapping with `composer.HuggingFaceModel`, it takes on the `model` prefix |
| 627 | consume_prefix_in_state_dict_if_present(state_dict, prefix="model.") |
| 628 | missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) |
| 629 | |
| 630 | if len(missing_keys) > 0: |
| 631 | logger.warning(f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}") |
| 632 | if len(unexpected_keys) > 0: |
| 633 | logger.warning(f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}") |
| 634 | |
| 635 | return model |
| 636 | |
| 637 | def forward( |
| 638 | self, |
| 639 | input_ids: Optional[torch.Tensor] = None, |
| 640 | attention_mask: Optional[torch.Tensor] = None, |
nothing calls this directly
no outgoing calls
no test coverage detected