Below is my SequenceClassificationClass
‘’’
class RobertaForSequenceClassification(RobertaPreTrainedModel):
def init(self, config):
super().init(config)
self.num_labels = config.num_labels
self.roberta = RobertaModel(config, add_pooling_layer=False) # initialize roberta with no pooling layer
self.dropout = nn.Dropout(config.hidden_dropout_prob)
self.linear = nn.Linear(config.hidden_size, config.num_labels)
self.post_init()
def forward(self, input_ids = None, attention_mask = None, token_type_ids=None, labels = None, **kwargs):
outputs = self.roberta(input_ids=input_ids, attention_mask = attention_mask,
token_type_ids = token_type_ids, **kwargs)
sequence_outputs = self.dropout(outputs[0][:, 0, :])
logits = self.linear(sequence_outputs)
loss = None
if labels is not None:
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(logits, labels)
return SequenceClassifierOutput(loss = loss, logits = logits,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions)
‘’’
pipeline results
When I switched to AutoModelForSequenceClassification, the pipeline gave me the right results so I guess the problem is from my custom head but I cant seem to debug it. I also got correct results using my custom head without a pipeline