| import torch |
| import torch.nn as nn |
| from transformers import AutoModel, AutoConfig, PreTrainedModel, PretrainedConfig |
|
|
|
|
| class CAPConfig(PretrainedConfig): |
| model_type = "cap" |
|
|
| def __init__(self, base_model_name="roberta-base", num_labels=3, dropout=0.1, **kwargs): |
| super().__init__(**kwargs) |
| self.base_model_name = base_model_name |
| self.num_labels = num_labels |
| self.dropout = dropout |
|
|
|
|
| class CAPModel(PreTrainedModel): |
| config_class = CAPConfig |
| base_model_prefix = "backbone" |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| backbone_config = AutoConfig.from_pretrained(config.base_model_name) |
| self.backbone = AutoModel.from_config(backbone_config) |
| hidden_size = self.backbone.config.hidden_size |
| self.num_labels = config.num_labels |
| self.dropout = nn.Dropout(config.dropout) |
| self.head = nn.Linear(hidden_size, config.num_labels) |
| self.post_init() |
|
|
| def forward(self, input_ids, attention_mask, token_valid_mask): |
| outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask) |
| subword_states = self.dropout(outputs.last_hidden_state) |
| token_logits = self.head(subword_states) |
| valid_mask = token_valid_mask.unsqueeze(-1) |
| masked_token_logits = token_logits * valid_mask |
| valid_counts = token_valid_mask.sum(dim=1, keepdim=True).clamp(min=1e-9) |
| sequence_logits = masked_token_logits.sum(dim=1) / valid_counts |
| return sequence_logits, token_logits |
