# Install required libraries
!pip install torch torchvision transformers datasets
import torch
from torch.utils.data import DataLoader
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from datasets import load_dataset
Load dataset
dataset = load_dataset("imdb")
train_data = dataset["train"]
test_data = dataset["test"]
Initialize tokenizer and model
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForSequenceClassification.from_pretrained("bert-base-uncased")
Tokenize dataset
def tokenize_data(data):
return tokenizer(data["text"], padding=True, truncation=True, max_length=512)
train_data = train_data.map(tokenize_data, batched=True)
test_data = test_data.map(tokenize_data, batched=True)
Convert to PyTorch data format
train_data.set_format("torch", columns=["input_ids", "attention_mask", "label"])
test_data.set_format("torch", columns=["input_ids", "attention_mask", "label"])
train_loader = DataLoader(train_data, batch_size=8, shuffle=True)
test_loader = DataLoader(test_data, batch_size=8, shuffle=False)
Set up optimizer
optimizer = AdamW(model.parameters(), lr=1e-5)
Training loop
for epoch in range(3): # Number of training epochs
model.train()
for batch in train_loader:
optimizer.zero_grad()
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["label"]
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss
loss.backward()
optimizer.step()
print(f"Epoch {epoch + 1} completed")
Evaluation
model.eval()
correct = 0
total = 0
for batch in test_loader:
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["label"]
with torch.no_grad():
outputs = model(input_ids, attention_mask=attention_mask)
predictions = torch.argmax(outputs.logits, dim=-1)
correct += (predictions == labels).sum().item()
total += labels.size(0)
accuracy = correct / total
print(f"Test Accuracy: {accuracy:.2f}")
