import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer

class SentimentClassifier(nn.Module):
    def __init__(self, pretrained_model_name, num_classes=2):
        super(SentimentClassifier, self).__init__()
        self.bert = BertModel.from_pretrained(pretrained_model_name)
        self.dropout = nn.Dropout(0.1)
        self.linear = nn.Linear(self.bert.config.hidden_size, num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        logits = self.linear(pooled_output)
        return logits

# Initialize the BERT model and tokenizer
pretrained_model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(pretrained_model_name)

# Example data
text = "This movie is really good!"
labels = torch.tensor([1])  # 1 for positive sentiment

# Tokenize input
inputs = tokenizer(text, return_tensors='pt')
input_ids = inputs['input_ids']
attention_mask = inputs['attention_mask']

# Initialize and load the model
model = SentimentClassifier(pretrained_model_name, num_classes=1)  # Binary classification

# Forward pass
logits = model(input_ids, attention_mask)

# Loss and optimizer
criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

# Binary classification expects labels as float (0 or 1)
labels = labels.float()

# Compute loss
loss = criterion(logits.view(-1), labels)

# Backward pass and update
optimizer.zero_grad()
loss.backward()
optimizer.step()

# Evaluation
predictions = torch.sigmoid(logits) > 0.5  # Threshold at 0.5
predicted_labels = predictions.long()

print("Predicted label:", predicted_labels.item())