!pip install transformers==4.4.2

import torch from transformers import BertTokenizer, BertForSequenceClassification, AdamW

Load the tokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)

Load the model

model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=3, output_attentions=False, output_hidden_states=False)

Define optimizer and learning rate

optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8)

Load the data

from datasets import load_dataset

dataset = load_dataset('imdb')

train_data = dataset['train'] test_data = dataset['test']

Define data preprocessing function

def preprocess_data(data): return tokenizer(data['text'], padding=True, truncation=True, max_length=512)

Preprocess the data

train_data = train_data.map(preprocess_data, batched=True) test_data = test_data.map(preprocess_data, batched=True)

Convert data to PyTorch tensors

train_input = torch.tensor(train_data['input_ids']) train_labels = torch.tensor(train_data['label']) train_attention_mask = torch.tensor(train_data['attention_mask'])

test_input = torch.tensor(test_data['input_ids']) test_labels = torch.tensor(test_data['label']) test_attention_mask = torch.tensor(test_data['attention_mask'])

Define data loader

from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler

batch_size = 32

train_data = TensorDataset(train_input, train_attention_mask, train_labels) train_sampler = RandomSampler(train_data) train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)

test_data = TensorDataset(test_input, test_attention_mask, test_labels) test_sampler = SequentialSampler(test_data) test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)

Define training function

def train(model, dataloader, optimizer, epochs=2): model.train()

for epoch in range(epochs):
    for batch in dataloader:
        input_ids = batch[0].to(device)
        attention_mask = batch[1].to(device)
        labels = batch[2].to(device)

        model.zero_grad()

        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)

        loss = outputs[0]

        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

        optimizer.step()

Train the model

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model.to(device)

train(model, train_dataloader, optimizer)

Evaluate the model

model.eval()

total = 0 correct = 0

for batch in test_dataloader: input_ids = batch[0].to(device) attention_mask = batch[1].to(device) labels = batch[2].to(device)

with torch.no_grad():
    outputs = model(input_ids, attention_mask=attention_mask)

logits = outputs[0]
predictions = torch.argmax(logits, dim=1)

total += labels.size(0)
correct += (predictions == labels).sum().item()

accuracy = correct / total

print('Accuracy: {:.2f}%'.format(accuracy * 100))