binuser007/Toxic_comment_classification_using_Bert

1import torch
2from transformers import BertTokenizer, AdamW
3from src.models.toxic_classifier import ToxicClassifier
4from src.models.trainer import ModelTrainer
5from src.data.data_loader import load_toxic_data, create_data_loaders
6import logging
7import os
8from torch.cuda.amp import GradScaler, autocast  # For mixed precision training
9
10# Setup logging
11logging.basicConfig(level=logging.INFO)
12logger = logging.getLogger(__name__)
13
14def train_model(
15    data_path: str,
16    model_save_path: str,
17    num_epochs: int = 5,
18    batch_size: int = 64,  # Increased for RTX 3060
19    learning_rate: float = 2e-5,
20    max_grad_norm: float = 1.0
21):
22    # Set device and enable CUDA optimizations
23    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
24    if device.type == 'cuda':
25        torch.backends.cudnn.benchmark = True
26    logger.info(f"Using device: {device}")
27
28    # Load tokenizer
29    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
30
31    # Load data
32    logger.info("Loading dataset...")
33    texts, labels = load_toxic_data(data_path)
34    train_loader, val_loader = create_data_loaders(
35        texts, 
36        labels, 
37        tokenizer, 
38        batch_size=batch_size
39    )
40
41    # Initialize model
42    logger.info("Initializing model...")
43    model = ToxicClassifier().to(device)
44    
45    # Initialize optimizer with weight decay
46    optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
47    
48    # Initialize gradient scaler for mixed precision training
49    scaler = GradScaler()
50
51    # Initialize trainer with mixed precision support
52    trainer = ModelTrainer(model, optimizer, criterion=torch.nn.BCELoss(), device=device, scaler=scaler)
53
54    # Training loop
55    logger.info("Starting training...")
56    best_val_loss = float('inf')
57    
58    for epoch in range(num_epochs):
59        # Train
60        train_metrics = trainer.train_epoch(train_loader)
61        logger.info(f"Epoch {epoch+1}/{num_epochs}")
62        logger.info(f"Training Loss: {train_metrics['loss']:.4f}")
63
64        # Evaluate
65        val_metrics = trainer.evaluate(val_loader)
66        val_loss = val_metrics['loss']
67        logger.info(f"Validation Loss: {val_loss:.4f}")
68
69        # Save best model
70        if val_loss < best_val_loss:
71            best_val_loss = val_loss
72            torch.save({
73                'epoch': epoch,
74                'model_state_dict': model.state_dict(),
75                'optimizer_state_dict': optimizer.state_dict(),
76                'loss': best_val_loss,
77            }, os.path.join(model_save_path, 'best_model.pt'))
78            logger.info("Saved best model checkpoint")
79
80    logger.info("Training completed!")
81
82if __name__ == "__main__":
83    DATA_PATH = os.path.join("data", "raw", "train.csv")
84    MODEL_SAVE_PATH = os.path.join("models", "saved")
85    
86    # Create model save directory if it doesn't exist
87    os.makedirs(MODEL_SAVE_PATH, exist_ok=True)
88    
89    train_model(DATA_PATH, MODEL_SAVE_PATH)