File size: 1,676 Bytes
12097aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from .configuration_tiny_log import TinyLogConfig
class TinyLogPreTrainedModel(PreTrainedModel):
config_class = TinyLogConfig
base_model_prefix = "tiny_log"
main_input_name = "input_ids"
class TinyLogForSequenceClassification(TinyLogPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.embedding = nn.Embedding(
config.vocab_size,
config.hidden_size,
padding_idx=config.pad_token_id,
)
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
self.post_init()
def forward(
self,
input_ids=None,
attention_mask=None,
labels=None,
return_dict=None,
**kwargs,
):
if input_ids is None:
raise ValueError("input_ids is required")
if attention_mask is None:
attention_mask = input_ids.ne(self.config.pad_token_id).long()
embeddings = self.embedding(input_ids)
mask = attention_mask.unsqueeze(-1).to(embeddings.dtype)
summed = (embeddings * mask).sum(dim=1)
denom = mask.sum(dim=1).clamp(min=1.0)
pooled = summed / denom
logits = self.classifier(pooled)
loss = None
if labels is not None:
loss = nn.CrossEntropyLoss()(logits, labels)
if return_dict is False:
output = (logits,)
return ((loss,) + output) if loss is not None else output
return SequenceClassifierOutput(loss=loss, logits=logits)
|