| import inspect |
|
|
| import torch.nn as nn |
| from transformers import AutoConfig, AutoModel, PreTrainedModel |
| from transformers.modeling_outputs import SequenceClassifierOutput |
|
|
| from custom_text_config import CustomTextConfig |
|
|
|
|
| class CustomTextForSequenceClassification(PreTrainedModel): |
| config_class = CustomTextConfig |
|
|
| def __init__(self, config: CustomTextConfig): |
| super().__init__(config) |
|
|
| if not getattr(config, "_name_or_path", None) and not getattr( |
| config, "_is_meta", False |
| ): |
| self.backbone = AutoModel.from_pretrained(config.backbone_name_or_path) |
| else: |
| backbone_config = AutoConfig.from_pretrained(config.backbone_name_or_path) |
| self.backbone = AutoModel.from_config(backbone_config) |
|
|
| hidden_size = self.backbone.config.hidden_size |
| self.classifier = nn.Linear(hidden_size, config.num_labels) |
| self.post_init() |
|
|
| def forward( |
| self, |
| input_ids, |
| attention_mask=None, |
| token_type_ids=None, |
| labels=None, |
| **kwargs, |
| ): |
| backbone_inputs = { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| "token_type_ids": token_type_ids, |
| **kwargs, |
| } |
| accepted_args = set(inspect.signature(self.backbone.forward).parameters) |
| backbone_inputs = { |
| key: value |
| for key, value in backbone_inputs.items() |
| if value is not None and key in accepted_args |
| } |
|
|
| outputs = self.backbone(**backbone_inputs) |
| cls_output = outputs.last_hidden_state[:, 0, :] |
| logits = self.classifier(cls_output) |
|
|
| loss = None |
| if labels is not None: |
| loss = nn.CrossEntropyLoss()(logits, labels) |
|
|
| return SequenceClassifierOutput(loss=loss, logits=logits) |
|
|
|
|
| CustomTextForSequenceClassification.register_for_auto_class( |
| "AutoModelForSequenceClassification" |
| ) |
|
|
