import torch
import torch.nn as nn
from pytorch_transformers import RobertaForSequenceClassification

# defining our model architecture
class RobertaForSequenceClassificationModel(nn.Module):
    def __init__(self,num_labels=2):
        super(RobertaForSequenceClassificationModel,self).__init__()
        self.num_labels = num_labels
        self.roberta = RobertaForSequenceClassification.from_pretrained("roberta-base",num_labels= self.num_labels)

def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
        outputs = self.roberta(input_ids, token_type_ids, attention_mask)
        logits = outputs[0]
        return logits