Blame for model.py

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import seaborn as sns
from torch.utils.data import Dataset
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import train_test_split
from transformers import(
    DebertaV2Tokenizer,
    DebertaV2ForSequenceClassification,
    TrainingArguments,
    Trainer,
    EarlyStoppingCallback
)
# Set seed for reproducibility
SEED = 42
torch.manual_seed(SEED)
np.random.seed(SEED)

class WeightedTrainer(Trainer):
    """Custom trainer that uses weighted loss function for imbaalance data."""

    def __init__(self, *args, class_weights=None, **kwargs):
        """Initialize trainer with class weights.

        Args:

        class_weights: Tensor of weights for each class

        *args, **kwargs: Arguments for parent Trainer class

        """
        super().__init__(*args, **kwargs)
        self.class_weights = class_weights.to(self.model.device)
        self.loss_fn = nn.CrossEntropyLoss(
                        weight=self.class_weights,
                        label_smoothing=0.1
                    )

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        """Compute loss using the weighted loss function



        args:

            model: Model to train

            inputs: Input to batch

            return_outputs: whether to return output along with loss



        Returns:

            Loss value or tuple of (loss, outputs)

        """

        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        loss = self.loss_fn(logits, labels)
        return (loss, outputs) if return_outputs else loss

class EmailClassification(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length):
        """Intialize dataset with texts, labels and tokenizer settings.

        Args:

            texts: List of email Text Strings

            labels: List of corresponding Labels

            tokenizer: Pretrained tokenizer

            max_length: Maximum sequence length

        """
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        """Return the number of sample in the dataset."""
        return len(self.texts)

    def __getitem__(self, idx):
        """Get a single item from the dataset



        Args:

            idx: Index of the Item

        Returns:

            Dict containing input_ids, attention_mask and labels

        """
        text = str(self.texts[idx])
        label = self.labels[idx]

        encoding = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt',
            padding='max_length'
        )

        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }


def compute_class_weights(labels):
    """Compute class weights inversly proportional to class frequencies.

    Args:

        labels: Array of labels



    return:

        Tensor of Class weights

    """
    class_counts = np.bincount(labels)
    total_sample = len(labels)
    class_weights = total_sample / (len(class_counts) * class_counts)
    return torch.tensor(class_weights, dtype=torch.float)



def compute_metrics(pred):
    """Compute metrics for evaluation.



    Args:

        pred: Prediction object from trainer



    Returns:

        Dict of metrics

    """
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)

    # Calculate metrics
    report = classification_report(
        labels,
        preds,
        target_names=list(label2id.keys()),
        output_dict=True
    )

    # Extract Metrics
    results = {
        'accuracy': report['accuracy'],
        'f1_macro': report['macro avg']['f1-score'],
        'f1_weighted': report['weighted avg']['f1-score'],
        'precision_weighted': report['weighted avg']['precision'],
        'recall_weighted': report['weighted avg']['recall']
    }

    # Add class specific f1 scores
    for cls_name, cls_id in label2id.items():
        results[f'f1_{cls_name}'] = report[cls_name]['f1-score']

    return results

def train(df, model_name, label2id, id2label, text_column='email_processed',

                        max_length=256, batch_size=16, epochs=5,seed=42):

    """Train and evaluate model using train/val/test split.

    Args:

        df: DataFrame containing the dataset

        model_name: Name of the pretrained model

        label2id: Dictionary mapping label names to IDs

        id2label: Dictionary mapping IDs to label names

        text_column: Column name containing processed email text

        max_length: Maximum sequence length for tokenizer

        batch_size: Batch size for training

        epochs: Number of training epochs

        seed: Random seed for reproducibility



    Returns:

        Tuple of (trained model, tokenizer)

    """
    # Map labels to IDs
    df['label_id'] = df['type'].map(label2id)

    # Load tokenizer
    tokenizer = DebertaV2Tokenizer.from_pretrained(model_name)

    # Prepare Data
    texts = df[text_column].values
    labels = df['label_id'].values

    # Compute class weights
    class_weights = compute_class_weights(labels)
    print("\nClass weights for handling imbalance: ")
    for label, weight in zip(label2id.keys(), class_weights):
        print(f"{label}: {weight.item():.4f}")

    # First, split data into train and temp (val+test combined)
    train_texts, temp_texts, train_labels, temp_labels = train_test_split(
        texts,
        labels,
        test_size=0.3,  # 70% train, 30% for val+test
        stratify=labels,  # Maintain class distribution
        random_state=seed
    )

    # Further split temp into val and test
    val_texts, test_texts, val_labels, test_labels = train_test_split(
        temp_texts,
        temp_labels,
        test_size=0.5,  # 15% val, 15% test (of the total dataset)
        stratify=temp_labels,  # Maintain class distribution
        random_state=seed
    )

    # Create Datasets
    train_dataset = EmailClassification(
        train_texts, train_labels, tokenizer, max_length
    )
    val_dataset = EmailClassification(
        val_texts, val_labels, tokenizer, max_length
    )
    test_dataset = EmailClassification(
        test_texts, test_labels, tokenizer, max_length
    )

    # Load Model
    model = DebertaV2ForSequenceClassification.from_pretrained(
        model_name,
        num_labels=len(label2id),
        id2label=id2label,
        label2id=label2id,
    )
    model.gradient_checkpointing_enable()
    # 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=5e-5,  # Starting learning rate
        lr_scheduler_type="cosine",
        warmup_steps=120,
        weight_decay=0.01,
        logging_dir="./logs",
        logging_steps=100,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1_weighted",
        greater_is_better=True,
        save_total_limit=1,
        seed=seed,
        report_to=None,
        fp16=True,  # Enable mixed precision training
        gradient_accumulation_steps=4,
    )

    # Define Trainer
    trainer = WeightedTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
        class_weights=class_weights,
    )

    # Train the model
    trainer.train()

    # Evaluate on validation set
    print("\nEvaluating on validation set:")
    val_results = trainer.evaluate()
    for metric_name, value in val_results.items():
        print(f"{metric_name}: {value:.4f}")

    # Evaluate on test set
    print("\nEvaluating on test set:")
    test_results = trainer.predict(test_dataset)
    test_metrics = compute_metrics(test_results)
    for metric_name, value in test_metrics.items():
        print(f"{metric_name}: {value:.4f}")

    # Get predictions for test set
    test_preds = np.argmax(test_results.predictions, axis=1)

    # Print classification report for test set
    print("\nTest Set Classification Report:")
    print(classification_report(
        test_labels,
        test_preds,
        target_names=list(label2id.keys())
    ))

    # Plot confusion matrix for test set
    plt.figure(figsize=(10, 8))
    cm = confusion_matrix(test_labels, test_preds)
    sns.heatmap(
        cm,
        annot=True,
        fmt='d',
        cmap='Blues',
        xticklabels=list(label2id.keys()),
        yticklabels=list(label2id.keys())
    )
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix - Test Set')
    plt.tight_layout()
    plt.savefig('confusion_matrix_test_set.png')
    plt.close()


    model.save_pretrained('./model')
    tokenizer.save_pretrained("./model")

    return model, tokenizer





if __name__ == "__main__":
    """Main function to run the complete pipeline."""
    print("Email Classification Training Pipeline")
    print(f"{'='*40}")
    # Label Mapping
    label2id = {
        'Incident': 0,
        'Request': 1,
        'Problem': 2,
        'Change': 3
    }
    id2label = {v: k for k, v in label2id.items()}
    # Model configuration
    MODEL_NAME = "microsoft/mdeberta-v3-base"
    MAX_LENGTH = 128
    BATCH_SIZE = 4
    EPOCHS = 5
    SEED = 42

    # Load your dataset (example)
    df = pd.read_csv("/content/processed_emails.csv")

    # Train and evaluate
    model, tokenizer = train(
        df=df,
        model_name=MODEL_NAME,
        label2id=label2id,
        id2label=id2label,
        text_column='email_processed',
        max_length=MAX_LENGTH,
        batch_size=BATCH_SIZE,
        epochs=EPOCHS,
        seed=SEED
    )