Bert Model transformer with a sequence classification/regression head. This head is just a linear layer on top of the pooled output. Used for, e.g., GLUE tasks.
| 1247 | |
| 1248 | |
| 1249 | class FlexBertForSequenceClassification(FlexBertPreTrainedModel): |
| 1250 | """Bert Model transformer with a sequence classification/regression head. |
| 1251 | |
| 1252 | This head is just a linear layer on top of the pooled output. Used for, |
| 1253 | e.g., GLUE tasks. |
| 1254 | """ |
| 1255 | |
| 1256 | def __init__(self, config: FlexBertConfig): |
| 1257 | super().__init__(config) |
| 1258 | self.num_labels = config.num_labels |
| 1259 | self.config = config |
| 1260 | |
| 1261 | self.bert = FlexBertModel(config) |
| 1262 | self.head = FlexBertPoolingHead(config) |
| 1263 | self.classifier = nn.Linear(config.hidden_size, config.num_labels) |
| 1264 | |
| 1265 | # Initialize weights and apply final processing |
| 1266 | self._init_weights(reset_params=False) |
| 1267 | |
| 1268 | def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None): |
| 1269 | assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified" |
| 1270 | if module: |
| 1271 | self._init_module_weights(module) |
| 1272 | else: |
| 1273 | assert isinstance(reset_params, bool) |
| 1274 | self.bert._init_weights(reset_params=reset_params) |
| 1275 | self.head._init_weights(reset_params=reset_params) |
| 1276 | init_weights(self.config, self.classifier, self.config.hidden_size, type_of_module=ModuleType.final_out) |
| 1277 | |
| 1278 | @classmethod |
| 1279 | def from_composer( |
| 1280 | cls, |
| 1281 | pretrained_checkpoint, |
| 1282 | state_dict=None, |
| 1283 | cache_dir=None, |
| 1284 | from_tf=False, |
| 1285 | config=None, |
| 1286 | *inputs, |
| 1287 | **kwargs, |
| 1288 | ): |
| 1289 | """Load from pre-trained.""" |
| 1290 | model = cls(config, *inputs, **kwargs) |
| 1291 | if from_tf: |
| 1292 | raise ValueError("Mosaic BERT does not support loading TensorFlow weights.") |
| 1293 | |
| 1294 | state_dict = torch.load(pretrained_checkpoint) |
| 1295 | # If the state_dict was saved after wrapping with `composer.HuggingFaceModel`, it takes on the `model` prefix |
| 1296 | consume_prefix_in_state_dict_if_present(state_dict, prefix="model.") |
| 1297 | missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) |
| 1298 | |
| 1299 | if len(missing_keys) > 0: |
| 1300 | logger.warning(f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}") |
| 1301 | if len(unexpected_keys) > 0: |
| 1302 | logger.warning(f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}") |
| 1303 | |
| 1304 | return model |
| 1305 | |
| 1306 | def forward( |