Utils/model_utils.py

import torch
import torch.nn as nn
from sklearn.preprocessing import LabelEncoder, StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
from sklearn.utils.class_weight import compute_class_weight
from torch.amp import autocast, GradScaler
from torch.utils.data import TensorDataset, DataLoader
from torch.nn.utils import clip_grad_norm_
from collections import Counter
import torch
import torch.nn as nn
import os
import torch.optim as optim

# Define the ImprovedTagClassifier class for tag prediction
class ImprovedTagClassifier(nn.Module):
    def __init__(self, input_size, output_size, dropout_rate=0.4):
        super(ImprovedTagClassifier, self).__init__()
        
        # First hidden layer: transforms input features to 512 dimensions
        self.fc1 = nn.Linear(input_size, 512)
        self.bn1 = nn.BatchNorm1d(512)  # Normalizes the output
        
        # Second hidden layer: reduces from 512 to 256 dimensions
        self.fc2 = nn.Linear(512, 256)
        self.bn2 = nn.BatchNorm1d(256)  # Normalizes again
        
        # Third hidden layer: further reduces to 128 dimensions
        self.fc3 = nn.Linear(256, 128)
        self.bn3 = nn.BatchNorm1d(128)  # Another normalization
        
        # Output layer: maps 128 dimensions to the number of classes
        self.fc4 = nn.Linear(128, output_size)
        
        # Tools to prevent overfitting and improve learning
        self.dropout = nn.Dropout(dropout_rate)  # Randomly drops some data
        self.leaky_relu = nn.LeakyReLU(0.1)  # Activation function with a small slope
        
        # Skip connection: connects layer 1 directly to layer 3
        self.skip1_3 = nn.Linear(512, 128)
        
        # Set up the initial weights for better training
        self._initialize_weights()
    
    def _initialize_weights(self):
        # Loop through all parts of the model
        for m in self.modules():
            if isinstance(m, nn.Linear):
                # Use a special method to set weights for linear layers
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu')
                if m.bias is not None:
                    # Set biases to zero
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm1d):
                # Set batch norm weights to 1 and biases to 0
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        # First block: process input through the first layer
        x1 = self.fc1(x)
        x1 = self.bn1(x1)  # Normalize
        x1 = self.leaky_relu(x1)  # Activate
        x1 = self.dropout(x1)  # Drop some data to prevent overfitting
        
        # Second block: process through the second layer
        x2 = self.fc2(x1)
        x2 = self.bn2(x2)  # Normalize
        x2 = self.leaky_relu(x2)  # Activate
        x2 = self.dropout(x2)  # Drop some data
        
        # Third block: process with a skip connection
        x3 = self.fc3(x2)
        skip_x1 = self.skip1_3(x1)  # Skip connection from first layer
        x3 = x3 + skip_x1  # Add the skip connection
        x3 = self.bn3(x3)  # Normalize
        x3 = self.leaky_relu(x3)  # Activate
        x3 = self.dropout(x3)  # Drop some data
        
        # Final output: get the class predictions
        output = self.fc4(x3)
        return output

class FocalLoss(nn.Module):
    """Focal Loss for handling class imbalance"""
    def __init__(self, weight=None, gamma=2.0, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.weight = weight  # Weights for each class
        self.gamma = gamma    # Focus on hard examples
        self.reduction = reduction
        self.ce_loss = nn.CrossEntropyLoss(weight=weight, reduction='none')
    
    def forward(self, inputs, targets):
        # Calculate basic cross-entropy loss
        ce_loss = self.ce_loss(inputs, targets)
        pt = torch.exp(-ce_loss)  # Probability of correct class
        focal_loss = ((1 - pt) ** self.gamma) * ce_loss  # Adjust loss
        
        # Combine losses based on reduction type
        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss
    
class MultiLevelTagClassifier:
    def __init__(self, device='cuda'):
        # Use GPU
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
        self.models = {}           # Store models for each parent tag
        self.preprocessors = {}    # Store preprocessing tools
        self.label_encoders = {}   # Store label encoders
        
        # Define tag groups
        self.tag_hierarchy = {
            'DIV': ['DIV', 'LIST', 'CARD'],
            'P': ['P', 'LI'],
            'INPUT': ['INPUT', 'DROPDOWN'],
            'ICON': ['ICON', 'CHECKBOX', 'RADIO'],
        }
        print(f"Using device: {self.device}")
    
    def prepare_data_for_subtask(self, df, parent_tag, subtags):
        # Get only the data for this parent tag’s subtags
        filtered_df = df[df['tag'].isin(subtags)].copy()
        print(f"\n=== Preparing data for {parent_tag} sub-classification ===")
        print(f"Subtags: {subtags}")
        print(f"Total samples: {len(filtered_df)}")
        print(f"Distribution: \n{filtered_df['tag'].value_counts()}")
        
        if len(filtered_df) == 0:
            print(f"No data found for {parent_tag} subtags!")
            return None, None, None, None, None, None
        
        y = filtered_df["tag"]  # Target tags
        X = filtered_df.drop(columns=["tag"])  # Features
        
        # Define which columns are categories and numerical features
        categorical_cols = ['type', 'prev_sibling_html_tag', 'child_1_html_tag', 'child_2_html_tag', 'parent_tag_html']
        continuous_cols = [col for col in X.columns if col not in categorical_cols]
        
        # Add missing columns with default values
        missing_cols = [col for col in categorical_cols + continuous_cols if col not in X.columns]
        if missing_cols:
            print(f"Warning: Missing columns {missing_cols} in data for {parent_tag}")
            for col in missing_cols:
                X[col] = 'unknown' if col in categorical_cols else 0
        
        # Process categories
        X[categorical_cols] = X[categorical_cols].astype(str).fillna('unknown')
        ohe = OneHotEncoder(sparse_output=False, handle_unknown='ignore')
        X_cat_encoded = ohe.fit_transform(X[categorical_cols])
        
        # Process continous features
        imputer = SimpleImputer(strategy='median')
        X_continuous_imputed = imputer.fit_transform(X[continuous_cols])
        scaler = StandardScaler()
        X_continuous_scaled = scaler.fit_transform(X_continuous_imputed)
        X_processed = np.concatenate([X_cat_encoded, X_continuous_scaled], axis=1)
        
        # Encode target tags
        label_encoder = LabelEncoder()
        y_encoded = label_encoder.fit_transform(y)
        
        # Boost rare classes by copying them
        class_counts = Counter(y_encoded)
        min_samples_threshold = max(10, len(subtags) * 3)
        rare_classes = [cls for cls, count in class_counts.items() if count < min_samples_threshold]
        
        for cls in rare_classes:
            idx = np.where(y_encoded == cls)[0]
            original_class_name = label_encoder.inverse_transform([cls])[0]
            samples_needed = min_samples_threshold - len(idx)
            print(f"Adding {samples_needed} copies to class '{original_class_name}'")
            for _ in range(samples_needed):
                sample_idx = np.random.choice(idx)
                new_sample = X_processed[sample_idx].copy()
                continuous_start = X_cat_encoded.shape[1]
                noise = np.random.normal(0, 0.05, size=X_continuous_scaled.shape[1])
                new_sample[continuous_start:] += noise
                X_processed = np.vstack([X_processed, new_sample])
                y_encoded = np.append(y_encoded, cls)
        
        # Bundle up preprocessing models
        preprocessors = {
            'ohe': ohe,
            'imputer': imputer,
            'scaler': scaler,
            'label_encoder': label_encoder,
            'categorical_cols': categorical_cols,
            'continuous_cols': continuous_cols
        }
        return X_processed, y_encoded, preprocessors, categorical_cols, continuous_cols, label_encoder
    
    def train_subtask_model(self, X, y, preprocessors, parent_tag, epochs=100):
        # Split data into train, validation, and test sets
        print(f"\n=== Training {parent_tag} sub-classifier ===")
        X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.15, random_state=42, stratify=y)
        X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.15, random_state=42, stratify=y_temp)
        print(f"Training set size: {X_train.shape[0]}")
        print(f"Validation set size: {X_val.shape[0]}")
        print(f"Test set size: {X_test.shape[0]}")
        
        # Balance classes
        class_weights = compute_class_weight('balanced', classes=np.unique(y_train), y=y_train)
        
        # Turn data into tensors
        X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
        y_train_tensor = torch.tensor(y_train, dtype=torch.long)
        X_val_tensor = torch.tensor(X_val, dtype=torch.float32)
        y_val_tensor = torch.tensor(y_val, dtype=torch.long)
        X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
        y_test_tensor = torch.tensor(y_test, dtype=torch.long)
        class_weights_tensor = torch.tensor(class_weights, dtype=torch.float32).to(self.device)
        
        # Set up datasets and loaders
        train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
        val_dataset = TensorDataset(X_val_tensor, y_val_tensor)
        test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
        train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=2)
        val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False, num_workers=2)
        test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=2)
        
        # Create and set up the model
        input_size = X_train.shape[1]
        output_size = len(np.unique(y))
        model = ImprovedTagClassifier(input_size, output_size).to(self.device)
        criterion = FocalLoss(weight=class_weights_tensor, gamma=2.0)
        optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-4)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5, verbose=True)
        scaler = GradScaler()
        
        # Training loop
        best_val_loss = float('inf')
        patience = 15
        counter = 0
        train_losses = []
        val_losses = []
        val_accuracies = []
        
        for epoch in range(epochs):
            model.train()
            running_loss = 0.0
            for batch_X, batch_y in train_loader:
                batch_X, batch_y = batch_X.to(self.device), batch_y.to(self.device)
                optimizer.zero_grad()
                with autocast(device_type=self.device.type):
                    outputs = model(batch_X)
                    loss = criterion(outputs, batch_y)
                scaler.scale(loss).backward()
                clip_grad_norm_(model.parameters(), max_norm=1.0)
                scaler.step(optimizer)
                scaler.update()
                running_loss += loss.item()
            
            train_loss = running_loss / len(train_loader)
            model.eval()
            val_running_loss = 0.0
            all_preds = []
            all_labels = []
            
            with torch.no_grad():
                for batch_X, batch_y in val_loader:
                    batch_X, batch_y = batch_X.to(self.device), batch_y.to(self.device)
                    with autocast(device_type=self.device.type):
                        outputs = model(batch_X)
                        loss = criterion(outputs, batch_y)
                    val_running_loss += loss.item()
                    _, preds = torch.max(outputs, 1)
                    all_preds.extend(preds.cpu().numpy())
                    all_labels.extend(batch_y.cpu().numpy())
            
            val_loss = val_running_loss / len(val_loader)
            val_accuracy = accuracy_score(all_labels, all_preds)
            scheduler.step(val_loss)
            
            # Track progress
            train_losses.append(train_loss)
            val_losses.append(val_loss)
            val_accuracies.append(val_accuracy)
            print(f"Epoch [{epoch+1}/{epochs}] - Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}")
            
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                counter = 0
                best_model_state = model.state_dict().copy()
            else:
                counter += 1
                if counter >= patience:
                    print(f"Early stopping triggered after {epoch+1} epochs")
                    break
        
        model.load_state_dict(best_model_state)
        model.eval()
        test_preds = []
        test_labels = []
        
        with torch.no_grad():
            for batch_X, batch_y in test_loader:
                batch_X, batch_y = batch_X.to(self.device), batch_y.to(self.device)
                outputs = model(batch_X)
                _, preds = torch.max(outputs, 1)
                test_preds.extend(preds.cpu().numpy())
                test_labels.extend(batch_y.cpu().numpy())
        
        test_accuracy = accuracy_score(test_labels, test_preds)
        print(f"\n{parent_tag} Model Test Accuracy: {test_accuracy:.4f}")
        print(f"\n{parent_tag} Classification Report:")
        print(classification_report(test_labels, test_preds, target_names=preprocessors['label_encoder'].classes_, zero_division=0))
        
        return model, (train_losses, val_losses, val_accuracies), test_accuracy
    
    def train_all_models(self, df_path, epochs=100):
        # Load and clean the main dataset
        print("Loading and cleaning data...")
        df = pd.read_csv(df_path)
        df.loc[(df["tag"] == "SPAN") & ((df["type"] == "RECTANGLE") | (df["type"] == "GROUP")), "tag"] = "DIV"
        children_cols = ['child_1_html_tag', 'child_2_html_tag']
        for col in children_cols:
            df[col] = df[col].apply(lambda x: "DIV" if isinstance(x, str) and '-' in x else x)
        for col in ['tag', 'prev_sibling_html_tag', 'child_1_html_tag', 'child_2_html_tag']:
            df[col] = df[col].str.upper()
        
        # Make a folder for models
        os.makedirs('../models/sub_classifiers', exist_ok=True)
        
        # Train a model for each parent tag
        for parent_tag, subtags in self.tag_hierarchy.items():
            print(f"\n{'='*60}")
            print(f"Training {parent_tag} sub-classifier")
            print(f"{'='*60}")
            result = self.prepare_data_for_subtask(df, parent_tag, subtags)
            if result[0] is None:
                print(f"Skipping {parent_tag} due to insufficient data")
                continue
            X, y, preprocessors, cat_cols, cont_cols, label_encoder = result
            model, training_history, test_accuracy = self.train_subtask_model(X, y, preprocessors, parent_tag, epochs)
            self.models[parent_tag] = model
            self.preprocessors[parent_tag] = preprocessors
            self.label_encoders[parent_tag] = label_encoder
            model_path = f'../models/sub_classifiers/{parent_tag.lower()}_classifier.pth'
            torch.save({
                'model_state_dict': model.state_dict(),
                'input_size': X.shape[1],
                'output_size': len(np.unique(y)),
                'preprocessors': preprocessors,
                'test_accuracy': test_accuracy
            }, model_path)
            print(f"Saved {parent_tag} model to {model_path}")
            self.plot_training_history(training_history, parent_tag)
    
    def plot_training_history(self, history, parent_tag):
        # Plot training history (good function naming no need for commenting but here we go)
        train_losses, val_losses, val_accuracies = history
        plt.figure(figsize=(12, 5))
        plt.subplot(1, 2, 1)
        plt.plot(train_losses, label='Training Loss')
        plt.plot(val_losses, label='Validation Loss')
        plt.title(f'{parent_tag} Model: Loss over epochs')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
        plt.subplot(1, 2, 2)
        plt.plot(val_accuracies, label='Validation Accuracy')
        plt.title(f'{parent_tag} Model: Accuracy over epochs')
        plt.xlabel('Epoch')
        plt.ylabel('Accuracy')
        plt.legend()
        plt.tight_layout()
        plt.savefig(f'../models/sub_classifiers/{parent_tag.lower()}_training_history.png')
        plt.close()
    
    def load_models(self, model_dir='../models/sub_classifiers'):
        # Load saved models
        for parent_tag in self.tag_hierarchy.keys():
            model_path = f'{model_dir}/{parent_tag.lower()}_classifier.pth'
            if os.path.exists(model_path):
                print(f"Loading {parent_tag} model from {model_path}")
                checkpoint = torch.load(model_path, map_location=self.device,weights_only=False)
                model = ImprovedTagClassifier(checkpoint['input_size'], checkpoint['output_size']).to(self.device)
                model.load_state_dict(checkpoint['model_state_dict'])
                model.eval()
                self.models[parent_tag] = model
                self.preprocessors[parent_tag] = checkpoint['preprocessors']
                self.label_encoders[parent_tag] = checkpoint['preprocessors']['label_encoder']
                print(f"Loaded {parent_tag} model (Test Accuracy: {checkpoint['test_accuracy']:.4f})")
            else:
                print(f"Model file {model_path} not found!")
    
    def predict_hierarchical(self, sample_data, base_prediction):
        # Predict a tag using the right sub-classifier
        if base_prediction not in self.tag_hierarchy:
            return base_prediction, 1.0
        if base_prediction not in self.models:
            print(f"No sub-classifier found for {base_prediction}")
            return base_prediction, 1.0
        preprocessors = self.preprocessors[base_prediction]
        sample_df = pd.DataFrame([sample_data])
        cat_cols = preprocessors['categorical_cols']
        cont_cols = preprocessors['continuous_cols']
        
        # Add missing columns
        for col in cat_cols + cont_cols:
            if col not in sample_df.columns:
                sample_df[col] = 'unknown' if col in cat_cols else 0
        
        sample_df[cat_cols] = sample_df[cat_cols].astype(str).fillna('unknown')
        X_cat = preprocessors['ohe'].transform(sample_df[cat_cols])
        X_cont = preprocessors['imputer'].transform(sample_df[cont_cols])
        X_cont = preprocessors['scaler'].transform(X_cont)
        X_processed = np.concatenate([X_cat, X_cont], axis=1)
        X_tensor = torch.tensor(X_processed, dtype=torch.float32).to(self.device)
        
        model = self.models[base_prediction]
        with torch.no_grad():
            outputs = model(X_tensor)
            probabilities = torch.softmax(outputs, dim=1)
            _, predicted = torch.max(outputs, 1)
        predicted_label = preprocessors['label_encoder'].inverse_transform([predicted.cpu().numpy()[0]])[0]
        confidence = probabilities.max().item()
        return predicted_label, confidence