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.
| 1378 | |
| 1379 | |
| 1380 | class FlexBertForMultipleChoice(FlexBertPreTrainedModel): |
| 1381 | """ |
| 1382 | Bert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a |
| 1383 | softmax) e.g. for RocStories/SWAG tasks. |
| 1384 | """ |
| 1385 | |
| 1386 | def __init__(self, config: FlexBertConfig): |
| 1387 | super().__init__(config) |
| 1388 | self.num_labels = config.num_labels |
| 1389 | self.config = config |
| 1390 | |
| 1391 | self.bert = FlexBertModel(config) |
| 1392 | self.head = FlexBertPoolingHead(config) |
| 1393 | |
| 1394 | # In multiple choice tasks, all choices are submitted in a batch, and |
| 1395 | # we compute a logit for each option independently. The logits are then |
| 1396 | # normalized in the forward pass to get a probability distribution over |
| 1397 | # the choices. |
| 1398 | self.classifier = nn.Linear(config.hidden_size, 1) |
| 1399 | |
| 1400 | # Initialize weights and apply final processing |
| 1401 | self._init_weights(reset_params=False) |
| 1402 | |
| 1403 | def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None): |
| 1404 | assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified" |
| 1405 | if module: |
| 1406 | self._init_module_weights(module) |
| 1407 | else: |
| 1408 | assert isinstance(reset_params, bool) |
| 1409 | self.bert._init_weights(reset_params=reset_params) |
| 1410 | self.head._init_weights(reset_params=reset_params) |
| 1411 | init_weights(self.config, self.classifier, self.config.hidden_size, type_of_module=ModuleType.final_out) |
| 1412 | |
| 1413 | @classmethod |
| 1414 | def from_composer( |
| 1415 | cls, |
| 1416 | pretrained_checkpoint, |
| 1417 | state_dict=None, |
| 1418 | cache_dir=None, |
| 1419 | from_tf=False, |
| 1420 | config=None, |
| 1421 | *inputs, |
| 1422 | **kwargs, |
| 1423 | ): |
| 1424 | """Load from pre-trained.""" |
| 1425 | model = cls(config, *inputs, **kwargs) |
| 1426 | if from_tf: |
| 1427 | raise ValueError("Mosaic BERT does not support loading TensorFlow weights.") |
| 1428 | |
| 1429 | state_dict = torch.load(pretrained_checkpoint) |
| 1430 | # If the state_dict was saved after wrapping with `composer.HuggingFaceModel`, it takes on the `model` prefix |
| 1431 | consume_prefix_in_state_dict_if_present(state_dict, prefix="model.") |
| 1432 | missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) |
| 1433 | |
| 1434 | if len(missing_keys) > 0: |
| 1435 | logger.warning(f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}") |
| 1436 | if len(unexpected_keys) > 0: |
| 1437 | logger.warning(f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}") |
nothing calls this directly
no outgoing calls
no test coverage detected