| | import torch
|
| | import torch.nn as nn
|
| | from torch.utils.data import DataLoader
|
| | from transformers import BertTokenizer
|
| | import torch.optim as optim
|
| | from sklearn.metrics import accuracy_score
|
| | from tqdm import tqdm
|
| |
|
| | from data.imdb_data import IMDBDataset
|
| | from model.sentiment_model import SentimentAnalysisModel
|
| |
|
| |
|
| | def evaluate(model, dataloader, device):
|
| | model.eval()
|
| | preds, labels = [], []
|
| |
|
| | with torch.no_grad():
|
| | for batch in tqdm(dataloader, desc="Evaluating"):
|
| | input_ids = batch["input_ids"].to(device)
|
| | attention_mask = batch["attention_mask"].to(device)
|
| | batch_labels = batch["labels"].to(device)
|
| |
|
| | outputs = model(input_ids, attention_mask)
|
| | predictions = torch.argmax(outputs, dim=1)
|
| |
|
| | preds.extend(predictions.cpu().tolist())
|
| | labels.extend(batch_labels.cpu().tolist())
|
| |
|
| | acc = accuracy_score(labels, preds)
|
| | print(f"Validation Accuracy: {acc:.4f}")
|
| | return acc
|
| |
|
| |
|
| | def train():
|
| |
|
| | model_name = "bert-base-uncased"
|
| | batch_size = 8
|
| | max_length = 256
|
| | lr = 2e-5
|
| | epochs = 3
|
| |
|
| |
|
| | device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| | print("Using device:", device)
|
| |
|
| | tokenizer = BertTokenizer.from_pretrained(model_name)
|
| |
|
| | train_dataset = IMDBDataset("train", tokenizer, max_length)
|
| | test_dataset = IMDBDataset("test", tokenizer, max_length)
|
| |
|
| | train_loader = DataLoader(
|
| | train_dataset,
|
| | batch_size=batch_size,
|
| | shuffle=True,
|
| | num_workers=4,
|
| | pin_memory=True
|
| | )
|
| |
|
| | test_loader = DataLoader(
|
| | test_dataset,
|
| | batch_size=batch_size,
|
| | shuffle=False,
|
| | num_workers=4,
|
| | pin_memory=True
|
| | )
|
| |
|
| | model = SentimentAnalysisModel(model_name).to(device)
|
| | print("Model device:", next(model.parameters()).device)
|
| |
|
| | optimizer = optim.AdamW(model.parameters(), lr=lr)
|
| | criterion = nn.CrossEntropyLoss()
|
| |
|
| |
|
| | for epoch in range(epochs):
|
| | model.train()
|
| | total_loss = 0
|
| |
|
| | loop = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}")
|
| | for step, batch in enumerate(loop):
|
| | input_ids = batch["input_ids"].to(device)
|
| | attention_mask = batch["attention_mask"].to(device)
|
| | labels = batch["labels"].to(device)
|
| |
|
| | optimizer.zero_grad()
|
| | outputs = model(input_ids, attention_mask)
|
| | loss = criterion(outputs, labels)
|
| |
|
| | loss.backward()
|
| | optimizer.step()
|
| |
|
| | total_loss += loss.item()
|
| | loop.set_postfix(loss=loss.item())
|
| |
|
| | avg_loss = total_loss / len(train_loader)
|
| | print(f"\nEpoch {epoch + 1} Training Loss: {avg_loss:.4f}")
|
| |
|
| | evaluate(model, test_loader, device)
|
| |
|
| |
|
| | torch.save(model.state_dict(), "bert_imdb_sentiment.pth")
|
| | print("Model saved.")
|
| |
|
| |
|
| | if __name__ == "__main__":
|
| | train()
|
| |
|
