train_model.py

import pandas as pd
from sklearn.preprocessing import LabelEncoder
from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
import torch
import re
from torch.utils.data import Dataset
from sklearn.model_selection import train_test_split
import os
import argparse

def load_dictionary(file_path):
    """Load and preprocess dictionary data from Excel file."""

    # Load Weighted sheet for terms and weights
    df_weighted = pd.read_excel(file_path, sheet_name="Weighted")
    # Load Typology sheet for descriptions
    df_typology = pd.read_excel(file_path, sheet_name="Typology")
    
    themes = []
    subthemes = []
    terms = []
    weights = []
    descriptions = []
    
    current_theme = None
    for _, row in df_weighted.iterrows():
        if pd.notna(row['Theme']):
            current_theme = row['Theme']
        if pd.notna(row['Sub-theme']):
            subtheme = row['Sub-theme']
            for weight_col in ['Weight: +2 (Strongly Supports RW View)', 
                          'Weight: +1 (Moderately Supports RW View)',
                          'Weight: 0 (Neutral/Ambiguous)',
                          'Weight: -1 (Moderately Opposes RW View)',
                          'Weight: -2 (Strongly Opposes RW View)']:
                
                if pd.notna(row[weight_col]):
                    for term in row[weight_col].split(','):
                        term = term.strip()
                        if term:
                            themes.append(current_theme)
                            subthemes.append(subtheme)
                            terms.append(term)
                            weight = int(re.search(r'Weight: ([+-]?\d+)', weight_col).group(1))
                            weights.append(weight)
                            
                            # Get description from typology sheet
                            desc = df_typology[(df_typology['Sub-theme'] == subtheme)]['Description'].values
                            if len(desc) > 0:
                                descriptions.append(desc[0])
                            else:
                                descriptions.append("")
    
    # Create subtheme-weight mapping (highest weight terms for each subtheme)
    subtheme_df = pd.DataFrame({
        'theme': themes,
        'subtheme': subthemes,
        'term': terms,
        'weight': weights,
        'description': descriptions
    })
    
    # Get max weight for each subtheme
    max_weights = subtheme_df.groupby('subtheme')['weight'].max().reset_index()
    max_weights = max_weights.rename(columns={'weight': 'max_weight'})
    
    # Get terms with max weight for each subtheme
    max_terms = subtheme_df.merge(max_weights, on='subtheme')
    max_terms = max_terms[max_terms['weight'] == max_terms['max_weight']]
    
    return subtheme_df, max_terms

class TextDataset(Dataset):
    """Custom PyTorch Dataset for handling text data with BERT tokenization."""

    def __init__(self, texts, labels, tokenizer, max_length):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
        
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        encoding = self.tokenizer(
            text,
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

def save_model(model, tokenizer, label_encoder, subtheme_df, max_terms, save_path="model"):
    """Save trained model and related components to disk."""

    # Create directory if it doesn't exist
    os.makedirs(save_path, exist_ok=True)
    
    # Save model and tokenizer
    model.save_pretrained(save_path)
    tokenizer.save_pretrained(save_path)
    
    # Save other components
    torch.save({
        'label_encoder': label_encoder,
        'subtheme_df': subtheme_df,
        'max_terms': max_terms
    }, os.path.join(save_path, "classifier_info.pt"))

def train_bert_model(subtheme_df, model_name='bert-base-uncased', epochs=15, batch_size=16, max_length=128):
    """Train BERT model for subtheme classification."""

    # Encode labels
    label_encoder = LabelEncoder()
    subtheme_df['subtheme_encoded'] = label_encoder.fit_transform(subtheme_df['subtheme'])
    
    # Split data
    train_texts, val_texts, train_labels, val_labels = train_test_split(
        subtheme_df['term'], subtheme_df['subtheme_encoded'], test_size=0.2, random_state=42
    )
    
    # Initialize tokenizer and model
    tokenizer = BertTokenizer.from_pretrained(model_name)
    model = BertForSequenceClassification.from_pretrained(
        model_name, 
        num_labels=len(label_encoder.classes_)
    )

    # Create datasets
    train_dataset = TextDataset(train_texts.tolist(), train_labels.tolist(), tokenizer, max_length)
    val_dataset = TextDataset(val_texts.tolist(), val_labels.tolist(), tokenizer, max_length)
    
    # Training arguments
    training_args = TrainingArguments(
        output_dir='./results',
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        learning_rate=2e-5,
        warmup_steps=500,
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=10,
        eval_strategy="steps",
        eval_steps=50,
        save_strategy="steps",
        save_steps=1000,
        load_best_model_at_end=True,
        save_total_limit=1
    )
    
    # Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset
    )
    
    # Train
    trainer.train()
    
    return model, tokenizer, label_encoder

def parse_arguments():
    """Parse command line arguments for model training."""
    
    parser = argparse.ArgumentParser(description='Train BERT model for subtheme classification')
    parser.add_argument('--input_file', type=str, default='RWDictionary.xlsx', help='Path to Excel dictionary file')
    parser.add_argument('--model_name', type=str, default='bert-base-uncased', help='Pretrained BERT model name')
    parser.add_argument('--epochs', type=int, default=15, help='Number of training epochs')
    parser.add_argument('--batch_size', type=int, default=16, help='Training batch size')
    parser.add_argument('--max_length', type=int, default=128, help='Maximum token sequence length')
    parser.add_argument('--output_dir', type=str, default='rwd_classifier', help='Directory to save trained model')
    
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_arguments()
    
    # Load dictionary and create subtheme-weight mapping
    subtheme_df, max_terms = load_dictionary(args.input_file)
    
    # Train model with command line arguments
    model, tokenizer, label_encoder = train_bert_model(
        subtheme_df,
        model_name=args.model_name,
        epochs=args.epochs,
        batch_size=args.batch_size,
        max_length=args.max_length
    )
    
    # Save the entire model and related data
    save_model(model, tokenizer, label_encoder, subtheme_df, max_terms, save_path=args.output_dir)
    
    print(f"Model training complete and saved to '{args.output_dir}' directory")