| import torch | |
| import torch.nn as nn | |
| from transformers import AutoModel, PreTrainedModel, AutoConfig | |
| from transformers.modeling_outputs import SequenceClassifierOutput | |
| from .configuration_multitask_toxicity import MultiTaskToxicityConfig | |
| class MultiTaskToxicityEncoder(PreTrainedModel): | |
| config_class = MultiTaskToxicityConfig | |
| base_model_prefix = "encoder" | |
| def __init__(self, config): | |
| super().__init__(config) | |
| is_meta = False | |
| import inspect | |
| try: | |
| for frame in inspect.stack(): | |
| if 'from_pretrained' in frame.function: | |
| is_meta = True | |
| break | |
| except Exception: | |
| pass | |
| if is_meta: | |
| enc_config = AutoConfig.from_pretrained(config.encoder_name, token=False) | |
| self.encoder = AutoModel.from_config(enc_config) | |
| else: | |
| self.encoder = AutoModel.from_pretrained(config.encoder_name, token=False) | |
| for param in self.encoder.parameters(): | |
| param.requires_grad = False | |
| hidden_size = self.encoder.config.hidden_size | |
| self.dropout = nn.Dropout(config.dropout) | |
| self.heads = nn.ModuleDict({ | |
| label: nn.Linear(hidden_size, 1) for label in config.labels | |
| }) | |
| self.post_init() | |
| def forward(self, input_ids, attention_mask, labels=None, **kwargs): | |
| cls_embedding = self.encoder( | |
| input_ids=input_ids, attention_mask=attention_mask | |
| ).last_hidden_state[:, 0] | |
| logits = torch.cat([ | |
| self.heads[label](self.dropout(cls_embedding)) | |
| for label in self.config.labels | |
| ], dim=1) | |
| loss = None | |
| if labels is not None: | |
| loss = nn.functional.binary_cross_entropy_with_logits(logits, labels) | |
| return SequenceClassifierOutput(loss=loss, logits=logits) |