train_depression_model.py.py

import pandas as pd
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from transformers import RobertaTokenizer, RobertaModel
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from tqdm import tqdm
import argparse

# 1. Dataset Class
class DepressionDataset(Dataset):
    def __init__(self, df, tokenizer, max_length=256):
        self.texts = df['clean_text'].values
        self.labels = df['is_depression'].values
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = str(self.texts[idx])
        label = self.labels[idx]
        
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )
        
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'label': torch.tensor(label, dtype=torch.long)
        }

# 2. Model Class
class DepressionClassifier(nn.Module):
    def __init__(self, dropout_rate=0.1):
        super(DepressionClassifier, self).__init__()
        self.roberta = RobertaModel.from_pretrained('roberta-base')
        self.dropout = nn.Dropout(dropout_rate)
        self.classifier = nn.Linear(768, 2)
        
    def forward(self, input_ids, attention_mask):
        outputs = self.roberta(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        
        sequence_output = outputs.last_hidden_state[:, 0, :]
        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)
        
        return logits

# 3. Prepare data loaders
def prepare_dataloaders(df, batch_size=16):
    # Split data
    train_df, temp_df = train_test_split(df, test_size=0.3, stratify=df['is_depression'], random_state=42)
    val_df, test_df = train_test_split(temp_df, test_size=0.5, stratify=temp_df['is_depression'], random_state=42)
    
    # Initialize tokenizer
    tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
    
    # Create datasets
    train_dataset = DepressionDataset(train_df, tokenizer)
    val_dataset = DepressionDataset(val_df, tokenizer)
    test_dataset = DepressionDataset(test_df, tokenizer)
    
    # Create dataloaders
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)
    
    return train_loader, val_loader, test_loader

# 4. Training function
def train_model(model, train_loader, val_loader, device, epochs=3, learning_rate=2e-5):
    # Move model to device
    model = model.to(device)
    
    # Initialize optimizer
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    
    # Initialize loss function
    loss_fn = nn.CrossEntropyLoss()
    
    # Training loop
    best_accuracy = 0
    
    for epoch in range(epochs):
        print(f'Epoch {epoch + 1}/{epochs}')
        
        # TRAINING
        model.train()
        train_loss = 0
        train_preds = []
        train_labels = []
        
        # Progress bar for training
        progress_bar = tqdm(train_loader, desc="Training")
        
        for batch in progress_bar:
            # Get batch data
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['label'].to(device)
            
            # Forward pass
            optimizer.zero_grad()
            outputs = model(input_ids, attention_mask)
            loss = loss_fn(outputs, labels)
            
            # Backward pass
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            
            # Track metrics
            train_loss += loss.item()
            _, preds = torch.max(outputs, dim=1)
            train_preds.extend(preds.cpu().tolist())
            train_labels.extend(labels.cpu().tolist())
            
            # Update progress bar
            progress_bar.set_postfix({'loss': loss.item()})
        
        # Calculate training metrics
        avg_train_loss = train_loss / len(train_loader)
        train_accuracy = accuracy_score(train_labels, train_preds)
        
        # VALIDATION
        model.eval()
        val_loss = 0
        val_preds = []
        val_labels = []
        
        with torch.no_grad():
            for batch in tqdm(val_loader, desc="Validation"):
                # Get batch data
                input_ids = batch['input_ids'].to(device)
                attention_mask = batch['attention_mask'].to(device)
                labels = batch['label'].to(device)
                
                # Forward pass
                outputs = model(input_ids, attention_mask)
                loss = loss_fn(outputs, labels)
                
                # Track metrics
                val_loss += loss.item()
                _, preds = torch.max(outputs, dim=1)
                val_preds.extend(preds.cpu().tolist())
                val_labels.extend(labels.cpu().tolist())
        
        # Calculate validation metrics
        avg_val_loss = val_loss / len(val_loader)
        val_accuracy = accuracy_score(val_labels, val_preds)
        
        # Print metrics
        print(f'Train Loss: {avg_train_loss:.4f} | Train Accuracy: {train_accuracy:.4f}')
        print(f'Val Loss: {avg_val_loss:.4f} | Val Accuracy: {val_accuracy:.4f}')
        
        # Save best model
        if val_accuracy > best_accuracy:
            torch.save(model.state_dict(), 'best_model.pt')
            best_accuracy = val_accuracy
            print(f'New best model saved with accuracy: {val_accuracy:.4f}')
        
        print('-' * 50)
    
    # Load best model
    model.load_state_dict(torch.load('best_model.pt'))
    return model

# 5. Evaluation function
def evaluate_model(model, test_loader, device):
    model.eval()
    test_preds = []
    test_labels = []
    
    with torch.no_grad():
        for batch in tqdm(test_loader, desc="Testing"):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['label'].to(device)
            
            outputs = model(input_ids, attention_mask)
            _, preds = torch.max(outputs, dim=1)
            
            test_preds.extend(preds.cpu().tolist())
            test_labels.extend(labels.cpu().tolist())
    
    # Calculate metrics
    accuracy = accuracy_score(test_labels, test_preds)
    precision, recall, f1, _ = precision_recall_fscore_support(
        test_labels, test_preds, average='binary'
    )
    
    return {
        'accuracy': accuracy,
        'precision': precision,
        'recall': recall,
        'f1': f1
    }

# 6. Main function
def main():
    parser = argparse.ArgumentParser(description='Train depression classifier')
    parser.add_argument('--data_path', type=str, default='depression_dataset_reddit_cleaned_final.csv', 
                        help='Path to the cleaned dataset')
    parser.add_argument('--batch_size', type=int, default=16, help='Batch size for training')
    parser.add_argument('--epochs', type=int, default=3, help='Number of training epochs')
    parser.add_argument('--learning_rate', type=float, default=2e-5, help='Learning rate')
    args = parser.parse_args()
    
    # Check for GPU
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Using device: {device}')
    
    # Load data
    df = pd.read_csv(args.data_path)
    print(f'Loaded dataset with {len(df)} examples')
    
    # Prepare data
    train_loader, val_loader, test_loader = prepare_dataloaders(
        df, batch_size=args.batch_size
    )
    print(f'Training samples: {len(train_loader.dataset)}')
    print(f'Validation samples: {len(val_loader.dataset)}')
    print(f'Testing samples: {len(test_loader.dataset)}')
    
    # Create model
    model = DepressionClassifier()
    print('Model created')
    
    # Train model
    print('Starting training...')
    trained_model = train_model(
        model, 
        train_loader, 
        val_loader, 
        device,
        epochs=args.epochs,
        learning_rate=args.learning_rate
    )
    
    # Evaluate model
    print('Evaluating model...')
    metrics = evaluate_model(trained_model, test_loader, device)
    
    # Print results
    print('\nTest Results:')
    for metric, value in metrics.items():
        print(f'{metric}: {value:.4f}')

if __name__ == '__main__':
    main()