Upload 8 files

Dockerfile

@@ -0,0 +1,43 @@
+# Use Python 3.9 as base image
+FROM python:3.9-slim
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    curl \
+    software-properties-common \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements file
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Create necessary directories with proper permissions
+RUN mkdir -p /app/uploads \
+    /app/predictions \
+    && chmod -R 777 /app/uploads \
+    /app/predictions
+
+# Copy the application code and utilities
+COPY . /app/
+COPY ../voting.py /app/
+COPY ../config.py /app/
+COPY ../dataset_utils.py /app/
+COPY ../label_encoders.pkl /app/
+
+# Set environment variables
+ENV PYTHONPATH=/app
+ENV PYTHONUNBUFFERED=1
+ENV PORT=7861
+
+# Expose the port the app runs on
+EXPOSE 7861
+
+# Command to run the application
+CMD ["python", "app.py"] 

app.py

@@ -0,0 +1,149 @@
+from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File
+from fastapi.responses import FileResponse
+from pydantic import BaseModel
+from typing import Optional, Dict, Any, List
+import uvicorn
+import torch
+import logging
+import os
+import asyncio
+import pandas as pd
+from datetime import datetime
+import shutil
+from pathlib import Path
+import numpy as np
+import sys
+
+# Add parent directory to Python path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+from voting import perform_voting_ensemble, save_predictions
+from config import LABEL_COLUMNS, PREDICTIONS_SAVE_DIR
+from dataset_utils import load_label_encoders
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="Ensemble Voting API")
+
+# Create necessary directories
+UPLOAD_DIR = Path("uploads")
+PREDICTIONS_DIR = Path(PREDICTIONS_SAVE_DIR)
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+PREDICTIONS_DIR.mkdir(parents=True, exist_ok=True)
+
+class EnsembleConfig(BaseModel):
+    model_names: List[str]
+    weights: Optional[Dict[str, float]] = None
+
+class EnsembleResponse(BaseModel):
+    message: str
+    metrics: Dict[str, Any]
+    predictions: List[Dict[str, Any]]
+
+class PredictionData(BaseModel):
+    model_name: str
+    probabilities: List[List[float]]
+    true_labels: Optional[List[int]] = None
+
+@app.get("/")
+async def root():
+    return {"message": "Ensemble Voting API"}
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+@app.post("/ensemble/vote")
+async def perform_ensemble(
+    config: EnsembleConfig
+):
+    """Perform ensemble voting using specified models"""
+    try:
+        # Perform ensemble voting
+        ensemble_reports, true_labels, ensemble_predictions = perform_voting_ensemble(config.model_names)
+        
+        # Load label encoders for decoding predictions
+        label_encoders = load_label_encoders()
+        
+        # Format predictions with original labels
+        formatted_predictions = []
+        for i, (col, preds) in enumerate(zip(LABEL_COLUMNS, ensemble_predictions)):
+            if true_labels[i] is not None:
+                label_encoder = label_encoders[col]
+                true_labels_orig = label_encoder.inverse_transform(true_labels[i])
+                pred_labels_orig = label_encoder.inverse_transform(preds)
+                
+                for true, pred in zip(true_labels_orig, pred_labels_orig):
+                    formatted_predictions.append({
+                        "field": col,
+                        "true_label": true,
+                        "predicted_label": pred
+                    })
+        
+        return EnsembleResponse(
+            message="Ensemble voting completed successfully",
+            metrics=ensemble_reports,
+            predictions=formatted_predictions
+        )
+        
+    except Exception as e:
+        logger.error(f"Ensemble voting failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Ensemble voting failed: {str(e)}")
+
+@app.post("/ensemble/save-predictions")
+async def save_model_predictions(
+    prediction_data: PredictionData
+):
+    """Save predictions from a model for later ensemble voting"""
+    try:
+        # Convert probabilities to numpy arrays
+        all_probs = [np.array(probs) for probs in prediction_data.probabilities]
+        true_labels = [np.array(prediction_data.true_labels) if prediction_data.true_labels else None]
+        
+        # Save predictions
+        save_predictions(
+            prediction_data.model_name,
+            all_probs,
+            true_labels
+        )
+        
+        return {
+            "message": f"Predictions saved successfully for model {prediction_data.model_name}",
+            "model_name": prediction_data.model_name
+        }
+        
+    except Exception as e:
+        logger.error(f"Failed to save predictions: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Failed to save predictions: {str(e)}")
+
+@app.get("/ensemble/available-models")
+async def get_available_models():
+    """Get list of models with saved predictions"""
+    try:
+        model_dirs = [d for d in os.listdir(PREDICTIONS_DIR) 
+                     if os.path.isdir(os.path.join(PREDICTIONS_DIR, d))]
+        
+        available_models = []
+        for model_name in model_dirs:
+            model_dir = os.path.join(PREDICTIONS_DIR, model_name)
+            has_all_files = all(
+                os.path.exists(os.path.join(model_dir, f"{col}_probs.pkl"))
+                for col in LABEL_COLUMNS
+            )
+            if has_all_files:
+                available_models.append(model_name)
+        
+        return {
+            "available_models": available_models,
+            "total_models": len(available_models)
+        }
+        
+    except Exception as e:
+        logger.error(f"Failed to get available models: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Failed to get available models: {str(e)}")
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7861))
+    uvicorn.run(app, host="0.0.0.0", port=port) 

config.py

@@ -0,0 +1,69 @@
+ # config.py
+
+import torch
+import os
+
+# --- Paths ---
+# Adjust DATA_PATH to your actual data location
+DATA_PATH = './data/synthetic_transactions_samples_5000.csv'
+TOKENIZER_PATH = './tokenizer/'
+LABEL_ENCODERS_PATH = './label_encoders.pkl'
+MODEL_SAVE_DIR = './saved_models/'
+PREDICTIONS_SAVE_DIR = './predictions/' # To save predictions for voting ensemble
+
+# --- Data Columns ---
+TEXT_COLUMN = "Sanction_Context"
+# Define all your target label columns
+LABEL_COLUMNS = [
+    "Red_Flag_Reason",
+    "Maker_Action",
+    "Escalation_Level",
+    "Risk_Category",
+    "Risk_Drivers",
+    "Investigation_Outcome"
+] 
+# Example metadata columns. Add actual numerical/categorical metadata if available in your CSV.
+# For now, it's an empty list. If you add metadata, ensure these columns exist and are numeric or can be encoded.
+METADATA_COLUMNS = [] # e.g., ["Risk_Score", "Transaction_Amount"]
+
+# --- Model Hyperparameters ---
+MAX_LEN = 128       # Maximum sequence length for transformer tokenizers
+BATCH_SIZE = 16     # Batch size for training and evaluation
+LEARNING_RATE = 2e-5 # Learning rate for AdamW optimizer
+NUM_EPOCHS = 3      # Number of training epochs. Adjust based on convergence.
+DROPOUT_RATE = 0.3  # Dropout rate for regularization
+
+# --- Device Configuration ---
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# --- Specific Model Configurations ---
+BERT_MODEL_NAME = 'bert-base-uncased'
+ROBERTA_MODEL_NAME = 'roberta-base'
+DEBERTA_MODEL_NAME = 'microsoft/deberta-base'
+
+# TF-IDF
+TFIDF_MAX_FEATURES = 5000 # Max features for TF-IDF vectorizer
+
+# --- Field-Specific Strategy (Conceptual) ---
+# This dictionary provides conceptual strategies for enhancing specific fields.
+# Actual implementation requires adapting the models (e.g., custom loss functions, metadata integration).
+FIELD_STRATEGIES = {
+    "Maker_Action": {
+        "loss": "focal_loss", # Requires custom Focal Loss implementation
+        "enhancements": ["action_templates", "context_prompt_tuning"] # Advanced NLP concepts
+    },
+    "Risk_Category": {
+        "enhancements": ["numerical_metadata", "transaction_patterns"] # Integrate METADATA_COLUMNS
+    },
+    "Escalation_Level": {
+        "enhancements": ["class_balancing", "policy_keyword_patterns"] # Handled by class weights/metadata
+    },
+    "Investigation_Outcome": {
+        "type": "classification_or_generation" # If generation, T5/BART would be needed.
+    }
+}
+
+# Ensure model save and predictions directories exist
+os.makedirs(MODEL_SAVE_DIR, exist_ok=True)
+os.makedirs(PREDICTIONS_SAVE_DIR, exist_ok=True)
+os.makedirs(TOKENIZER_PATH, exist_ok=True)

dataset_utils.py

@@ -0,0 +1,165 @@
+# dataset_utils.py
+
+import pandas as pd
+import torch
+from torch.utils.data import Dataset, DataLoader
+from sklearn.preprocessing import LabelEncoder
+from transformers import BertTokenizer, RobertaTokenizer, DebertaTokenizer
+import pickle
+import os
+
+from config import TEXT_COLUMN, LABEL_COLUMNS, MAX_LEN, TOKENIZER_PATH, LABEL_ENCODERS_PATH, METADATA_COLUMNS
+
+class ComplianceDataset(Dataset):
+    """
+    Custom Dataset class for handling text and multi-output labels for PyTorch models.
+    """
+    def __init__(self, texts, labels, tokenizer, max_len):
+        self.texts = texts
+        self.labels = labels
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+
+    def __len__(self):
+        """Returns the total number of samples in the dataset."""
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        """
+        Retrieves a sample from the dataset at the given index.
+        Tokenizes the text and converts labels to a PyTorch tensor.
+        """
+        text = str(self.texts[idx])
+        # Tokenize the text, padding to max_length and truncating if longer.
+        # return_tensors="pt" ensures PyTorch tensors are returned.
+        inputs = self.tokenizer(
+            text,
+            padding='max_length',
+            truncation=True,
+            max_length=self.max_len,
+            return_tensors="pt"
+        )
+        # Squeeze removes the batch dimension (which is 1 here because we process one sample at a time)
+        inputs = {key: val.squeeze(0) for key, val in inputs.items()}
+        # Convert labels to a PyTorch long tensor
+        labels = torch.tensor(self.labels[idx], dtype=torch.long)
+        return inputs, labels
+
+class ComplianceDatasetWithMetadata(Dataset):
+    """
+    Custom Dataset class for handling text, additional numerical metadata, and multi-output labels.
+    Used for hybrid models combining text and tabular features.
+    """
+    def __init__(self, texts, metadata, labels, tokenizer, max_len):
+        self.texts = texts
+        self.metadata = metadata # Expects metadata as a NumPy array or list of lists
+        self.labels = labels
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+
+    def __len__(self):
+        """Returns the total number of samples in the dataset."""
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        """
+        Retrieves a sample, its metadata, and labels from the dataset at the given index.
+        Tokenizes text, converts metadata and labels to PyTorch tensors.
+        """
+        text = str(self.texts[idx])
+        inputs = self.tokenizer(
+            text,
+            padding='max_length',
+            truncation=True,
+            max_length=self.max_len,
+            return_tensors="pt"
+        )
+        inputs = {key: val.squeeze(0) for key, val in inputs.items()}
+        # Convert metadata for the current sample to a float tensor
+        metadata = torch.tensor(self.metadata[idx], dtype=torch.float)
+        labels = torch.tensor(self.labels[idx], dtype=torch.long)
+        return inputs, metadata, labels
+
+def load_and_preprocess_data(data_path):
+    """
+    Loads data from a CSV, fills missing values, and encodes categorical labels.
+    Also handles converting specified METADATA_COLUMNS to numeric.
+
+    Args:
+        data_path (str): Path to the CSV data file.
+
+    Returns:
+        tuple: A tuple containing:
+            - data (pd.DataFrame): The preprocessed DataFrame.
+            - label_encoders (dict): A dictionary of LabelEncoder objects for each label column.
+    """
+    data = pd.read_csv(data_path)
+    data.fillna("Unknown", inplace=True) # Fill any missing text values with "Unknown"
+
+    # Convert metadata columns to numeric, coercing errors and filling NaNs with 0
+    # This ensures metadata is suitable for neural networks.
+    for col in METADATA_COLUMNS:
+        if col in data.columns:
+            data[col] = pd.to_numeric(data[col], errors='coerce').fillna(0) # Fill NaN with 0 or a suitable value
+
+    label_encoders = {col: LabelEncoder() for col in LABEL_COLUMNS}
+    for col in LABEL_COLUMNS:
+        # Fit and transform each label column using its respective LabelEncoder
+        data[col] = label_encoders[col].fit_transform(data[col])
+    return data, label_encoders
+
+def get_tokenizer(model_name):
+    """
+    Returns the appropriate Hugging Face tokenizer based on the model name.
+
+    Args:
+        model_name (str): The name of the pre-trained model (e.g., 'bert-base-uncased').
+
+    Returns:
+        transformers.PreTrainedTokenizer: The initialized tokenizer.
+    """
+    if "bert" in model_name.lower():
+        return BertTokenizer.from_pretrained(model_name)
+    elif "roberta" in model_name.lower():
+        return RobertaTokenizer.from_pretrained(model_name)
+    elif "deberta" in model_name.lower():
+        return DebertaTokenizer.from_pretrained(model_name)
+    else:
+        raise ValueError(f"Unsupported tokenizer for model: {model_name}")
+
+def save_label_encoders(label_encoders):
+    """
+    Saves a dictionary of label encoders to a pickle file.
+    This is crucial for decoding predictions back to original labels.
+
+    Args:
+        label_encoders (dict): Dictionary of LabelEncoder objects.
+    """
+    with open(LABEL_ENCODERS_PATH, "wb") as f:
+        pickle.dump(label_encoders, f)
+    print(f"Label encoders saved to {LABEL_ENCODERS_PATH}")
+
+def load_label_encoders():
+    """
+    Loads a dictionary of label encoders from a pickle file.
+
+    Returns:
+        dict: Loaded dictionary of LabelEncoder objects.
+    """
+    with open(LABEL_ENCODERS_PATH, "rb") as f:
+        return pickle.load(f)
+    print(f"Label encoders loaded from {LABEL_ENCODERS_PATH}")
+
+
+def get_num_labels(label_encoders):
+    """
+    Returns a list containing the number of unique classes for each label column.
+    This list is used to define the output dimensions of the model's classification heads.
+
+    Args:
+        label_encoders (dict): Dictionary of LabelEncoder objects.
+
+    Returns:
+        list: A list of integers, where each integer is the number of classes for a label.
+    """
+    return [len(label_encoders[col].classes_) for col in LABEL_COLUMNS]

label_encoders.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c336fd07858af76d40c7200de1a769099abeec25d4f48b999351318680d4e4d6
+size 2047

requirements.txt

@@ -0,0 +1,11 @@
+fastapi==0.104.1
+uvicorn==0.24.0
+pydantic==2.4.2
+numpy==1.24.3
+pandas==2.1.2
+scikit-learn==1.3.2
+python-multipart==0.0.6
+python-jose==3.3.0
+passlib==1.7.4
+bcrypt==4.0.1
+python-dotenv==1.0.0 

train_utils.py

@@ -0,0 +1,310 @@
+# train_utils.py
+
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from sklearn.metrics import classification_report
+from sklearn.utils.class_weight import compute_class_weight
+import numpy as np
+from tqdm import tqdm
+import pandas as pd
+import os
+import joblib
+
+from config import DEVICE, LABEL_COLUMNS, NUM_EPOCHS, LEARNING_RATE, MODEL_SAVE_DIR
+
+def get_class_weights(data_df, field, label_encoder):
+    """
+    Computes balanced class weights for a given target field.
+    These weights can be used in the loss function to mitigate class imbalance.
+
+    Args:
+        data_df (pd.DataFrame): The DataFrame containing the original (unencoded) label data.
+        field (str): The name of the label column for which to compute weights.
+        label_encoder (sklearn.preprocessing.LabelEncoder): The label encoder fitted for this field.
+
+    Returns:
+        torch.Tensor: A tensor of class weights for the specified field.
+    """
+    # Get the original labels for the specified field
+    y = data_df[field].values
+    # Use label_encoder.transform directly - it will handle unseen labels
+    try:
+        y_encoded = label_encoder.transform(y)
+    except ValueError as e:
+        print(f"Warning: {e}")
+        print(f"Using only seen labels for class weights calculation")
+        # Filter out unseen labels
+        seen_labels = set(label_encoder.classes_)
+        y_filtered = [label for label in y if label in seen_labels]
+        y_encoded = label_encoder.transform(y_filtered)
+    
+    # Ensure y_encoded is integer type
+    y_encoded = y_encoded.astype(int)
+    
+    # Initialize counts for all possible classes
+    n_classes = len(label_encoder.classes_)
+    class_counts = np.zeros(n_classes, dtype=int)
+    
+    # Count occurrences of each class
+    for i in range(n_classes):
+        class_counts[i] = np.sum(y_encoded == i)
+    
+    # Calculate weights for all classes
+    total_samples = len(y_encoded)
+    class_weights = np.ones(n_classes)  # Default weight of 1 for unseen classes
+    seen_classes = class_counts > 0
+    if np.any(seen_classes):
+        class_weights[seen_classes] = total_samples / (np.sum(seen_classes) * class_counts[seen_classes])
+    
+    return torch.tensor(class_weights, dtype=torch.float)
+
+def initialize_criterions(data_df, label_encoders):
+    """
+    Initializes CrossEntropyLoss criteria for each label column, applying class weights.
+
+    Args:
+        data_df (pd.DataFrame): The original (unencoded) DataFrame. Used to compute class weights.
+        label_encoders (dict): Dictionary of LabelEncoder objects.
+
+    Returns:
+        dict: A dictionary where keys are label column names and values are
+              initialized `torch.nn.CrossEntropyLoss` objects.
+    """
+    field_criterions = {}
+    for field in LABEL_COLUMNS:
+        # Get class weights for the current field
+        weights = get_class_weights(data_df, field, label_encoders[field])
+        # Initialize CrossEntropyLoss with the computed weights and move to the device
+        field_criterions[field] = torch.nn.CrossEntropyLoss(weight=weights.to(DEVICE))
+    return field_criterions
+
+def train_model(model, loader, optimizer, field_criterions, epoch):
+    """
+    Trains the given PyTorch model for one epoch.
+
+    Args:
+        model (torch.nn.Module): The model to train.
+        loader (torch.utils.data.DataLoader): DataLoader for training data.
+        optimizer (torch.optim.Optimizer): Optimizer for model parameters.
+        field_criterions (dict): Dictionary of loss functions for each label.
+        epoch (int): Current epoch number (for progress bar description).
+
+    Returns:
+        float: Average training loss for the epoch.
+    """
+    model.train() # Set the model to training mode
+    total_loss = 0
+    # Use tqdm for a progress bar during training
+    tqdm_loader = tqdm(loader, desc=f"Epoch {epoch + 1} Training")
+
+    for batch in tqdm_loader:
+        # Unpack batch based on whether it contains metadata
+        if len(batch) == 2: # Text-only models (inputs, labels)
+            inputs, labels = batch
+            input_ids = inputs['input_ids'].to(DEVICE)
+            attention_mask = inputs['attention_mask'].to(DEVICE)
+            labels = labels.to(DEVICE)
+            # Forward pass through the model
+            outputs = model(input_ids, attention_mask)
+        elif len(batch) == 3: # Text + Metadata models (inputs, metadata, labels)
+            inputs, metadata, labels = batch
+            input_ids = inputs['input_ids'].to(DEVICE)
+            attention_mask = inputs['attention_mask'].to(DEVICE)
+            metadata = metadata.to(DEVICE)
+            labels = labels.to(DEVICE)
+            # Forward pass through the hybrid model
+            outputs = model(input_ids, attention_mask, metadata)
+        else:
+            raise ValueError("Unsupported batch format. Expected 2 or 3 items in batch.")
+
+        loss = 0
+        # Calculate total loss by summing loss for each label column
+        # `outputs` is a list of logits, one for each label column
+        for i, output_logits in enumerate(outputs):
+            # `labels[:, i]` gets the true labels for the i-th label column
+            # `field_criterions[LABEL_COLUMNS[i]]` selects the appropriate loss function
+            loss += field_criterions[LABEL_COLUMNS[i]](output_logits, labels[:, i])
+
+        optimizer.zero_grad() # Clear previous gradients
+        loss.backward()       # Backpropagation
+        optimizer.step()      # Update model parameters
+        total_loss += loss.item() # Accumulate loss
+        tqdm_loader.set_postfix(loss=loss.item()) # Update progress bar with current batch loss
+
+    return total_loss / len(loader) # Return average loss for the epoch
+
+def evaluate_model(model, loader):
+    """
+    Evaluates the given PyTorch model on a validation/test set.
+
+    Args:
+        model (torch.nn.Module): The model to evaluate.
+        loader (torch.utils.data.DataLoader): DataLoader for evaluation data.
+
+    Returns:
+        tuple: A tuple containing:
+            - reports (dict): Classification reports (dict format) for each label column.
+            - truths (list): List of true label arrays for each label column.
+            - predictions (list): List of predicted label arrays for each label column.
+    """
+    model.eval() # Set the model to evaluation mode (disables dropout, batch norm updates, etc.)
+    # Initialize lists to store predictions and true labels for each output head
+    predictions = [[] for _ in range(len(LABEL_COLUMNS))]
+    truths = [[] for _ in range(len(LABEL_COLUMNS))]
+
+    with torch.no_grad(): # Disable gradient calculations during evaluation for efficiency
+        for batch in tqdm(loader, desc="Evaluation"):
+            if len(batch) == 2:
+                inputs, labels = batch
+                input_ids = inputs['input_ids'].to(DEVICE)
+                attention_mask = inputs['attention_mask'].to(DEVICE)
+                labels = labels.to(DEVICE)
+                outputs = model(input_ids, attention_mask)
+            elif len(batch) == 3:
+                inputs, metadata, labels = batch
+                input_ids = inputs['input_ids'].to(DEVICE)
+                attention_mask = inputs['attention_mask'].to(DEVICE)
+                metadata = metadata.to(DEVICE)
+                labels = labels.to(DEVICE)
+                outputs = model(input_ids, attention_mask, metadata)
+            else:
+                raise ValueError("Unsupported batch format.")
+
+            for i, output_logits in enumerate(outputs):
+                # Get the predicted class by taking the argmax of the logits
+                preds = torch.argmax(output_logits, dim=1).cpu().numpy()
+                predictions[i].extend(preds)
+                # Get the true labels for the current output head
+                truths[i].extend(labels[:, i].cpu().numpy())
+
+    reports = {}
+    # Generate classification report for each label column
+    for i, col in enumerate(LABEL_COLUMNS):
+        try:
+            # `zero_division=0` handles cases where a class might have no true or predicted samples
+            reports[col] = classification_report(truths[i], predictions[i], output_dict=True, zero_division=0)
+        except ValueError:
+            # Handle cases where a label might not appear in the validation set,
+            # which could cause classification_report to fail.
+            print(f"Warning: Could not generate classification report for {col}. Skipping.")
+            reports[col] = {'accuracy': 0, 'weighted avg': {'precision': 0, 'recall': 0, 'f1-score': 0, 'support': 0}}
+    return reports, truths, predictions
+
+def summarize_metrics(metrics):
+    """
+    Summarizes classification reports into a readable Pandas DataFrame.
+
+    Args:
+        metrics (dict): Dictionary of classification reports, as returned by `evaluate_model`.
+
+    Returns:
+        pd.DataFrame: A DataFrame summarizing precision, recall, f1-score, accuracy, and support for each field.
+    """
+    summary = []
+    for field, report in metrics.items():
+        # Safely get metrics, defaulting to 0 if not present (e.g., for empty reports)
+        precision = report['weighted avg']['precision'] if 'weighted avg' in report else 0
+        recall = report['weighted avg']['recall'] if 'weighted avg' in report else 0
+        f1 = report['weighted avg']['f1-score'] if 'weighted avg' in report else 0
+        support = report['weighted avg']['support'] if 'weighted avg' in report else 0
+        accuracy = report['accuracy'] if 'accuracy' in report else 0 # Accuracy is usually top-level
+        summary.append({
+            "Field": field,
+            "Precision": precision,
+            "Recall": recall,
+            "F1-Score": f1,
+            "Accuracy": accuracy,
+            "Support": support
+        })
+    return pd.DataFrame(summary)
+
+def save_model(model, model_name, save_format='pth'):
+    """
+    Saves the state dictionary of a PyTorch model.
+
+    Args:
+        model (torch.nn.Module): The trained PyTorch model.
+        model_name (str): A descriptive name for the model (used for filename).
+        save_format (str): Format to save the model in ('pth' for PyTorch models, 'pickle' for traditional ML models).
+    """
+    # Construct the save path dynamically relative to the project root
+    if save_format == 'pth':
+        model_path = os.path.join(MODEL_SAVE_DIR, f"{model_name}_model.pth")
+        torch.save(model.state_dict(), model_path)
+    elif save_format == 'pickle':
+        model_path = os.path.join(MODEL_SAVE_DIR, f"{model_name}.pkl")
+        joblib.dump(model, model_path)
+    else:
+        raise ValueError(f"Unsupported save format: {save_format}")
+    
+    print(f"Model saved to {model_path}")
+
+def load_model_state(model, model_name, model_class, num_labels, metadata_dim=0):
+    """
+    Loads the state dictionary into a PyTorch model.
+
+    Args:
+        model (torch.nn.Module): An initialized model instance (architecture).
+        model_name (str): The name of the model to load.
+        model_class (class): The class of the model (e.g., BertMultiOutputModel).
+        num_labels (list): List of number of classes for each label.
+        metadata_dim (int): Dimensionality of metadata features, if applicable (default 0 for text-only).
+
+    Returns:
+        torch.nn.Module: The model with loaded state_dict, moved to the correct device, and set to eval mode.
+    """
+    model_path = os.path.join(MODEL_SAVE_DIR, f"{model_name}_model.pth")
+    if not os.path.exists(model_path):
+        print(f"Warning: Model file not found at {model_path}. Returning a newly initialized model instance.")
+        # Re-initialize the model if not found, to ensure it has the correct architecture
+        if metadata_dim > 0:
+            return model_class(num_labels, metadata_dim=metadata_dim).to(DEVICE)
+        else:
+            return model_class(num_labels).to(DEVICE)
+
+    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+    model.to(DEVICE)
+    model.eval() # Set to evaluation mode after loading
+    print(f"Model loaded from {model_path}")
+    return model
+
+def predict_probabilities(model, loader):
+    """
+    Generates prediction probabilities for each label for a given model.
+    This is used for confidence scoring and feeding into a voting ensemble.
+
+    Args:
+        model (torch.nn.Module): The trained PyTorch model.
+        loader (torch.utils.data.DataLoader): DataLoader for the data to predict on.
+
+    Returns:
+        list: A list of lists of numpy arrays. Each inner list corresponds to a label column,
+              containing the softmax probabilities for each sample for that label.
+    """
+    model.eval() # Set to evaluation mode
+    # List to store probabilities for each output head
+    all_probabilities = [[] for _ in range(len(LABEL_COLUMNS))]
+
+    with torch.no_grad():
+        for batch in tqdm(loader, desc="Predicting Probabilities"):
+            # Unpack batch, ignoring labels as we only need inputs
+            if len(batch) == 2:
+                inputs, _ = batch
+                input_ids = inputs['input_ids'].to(DEVICE)
+                attention_mask = inputs['attention_mask'].to(DEVICE)
+                outputs = model(input_ids, attention_mask)
+            elif len(batch) == 3:
+                inputs, metadata, _ = batch
+                input_ids = inputs['input_ids'].to(DEVICE)
+                attention_mask = inputs['attention_mask'].to(DEVICE)
+                metadata = metadata.to(DEVICE)
+                outputs = model(input_ids, attention_mask, metadata)
+            else:
+                raise ValueError("Unsupported batch format.")
+
+            for i, out_logits in enumerate(outputs):
+                # Apply softmax to logits to get probabilities
+                probs = torch.softmax(out_logits, dim=1).cpu().numpy()
+                all_probabilities[i].extend(probs)
+    return all_probabilities

voting.py

@@ -0,0 +1,152 @@
+# voting.py
+
+import numpy as np
+import pandas as pd
+from collections import defaultdict
+from sklearn.metrics import classification_report
+import os
+import pickle
+
+from config import LABEL_COLUMNS, PREDICTIONS_SAVE_DIR
+
+def save_predictions(model_name, all_probs, true_labels):
+    """
+    Saves the prediction probabilities and true labels for each target field
+    from a specific model. This data is then used by the voting ensemble.
+
+    Args:
+        model_name (str): Unique identifier for the model (e.g., "BERT", "TF_IDF_LR").
+        all_probs (list): A list where each element is a NumPy array of probabilities
+                          for a corresponding label column (shape: num_samples, num_classes).
+        true_labels (list): A list where each element is a NumPy array of true labels
+                            for a corresponding label column (shape: num_samples,).
+    """
+    model_preds_dir = os.path.join(PREDICTIONS_SAVE_DIR, model_name)
+    os.makedirs(model_preds_dir, exist_ok=True) # Ensure the model-specific directory exists
+
+    for i, col in enumerate(LABEL_COLUMNS):
+        # Define file paths for probabilities and true labels for the current field
+        prob_file = os.path.join(model_preds_dir, f"{col}_probs.pkl")
+        true_file = os.path.join(model_preds_dir, f"{col}_true.pkl")
+
+        # Save probabilities (list of arrays) and true labels (list of arrays)
+        with open(prob_file, 'wb') as f:
+            pickle.dump(all_probs[i], f)
+        with open(true_file, 'wb') as f:
+            pickle.dump(true_labels[i], f)
+    print(f"Predictions for {model_name} saved to {model_preds_dir}")
+
+def load_predictions(model_name):
+    """
+    Loads saved prediction probabilities and true labels for a given model.
+
+    Args:
+        model_name (str): Unique identifier for the model.
+
+    Returns:
+        tuple: A tuple containing:
+            - all_probs (list): List of NumPy arrays of probabilities for each label column.
+            - true_labels (list): List of NumPy arrays of true labels for each label column.
+            Returns (None, None) if files are not found.
+    """
+    model_preds_dir = os.path.join(PREDICTIONS_SAVE_DIR, model_name)
+    all_probs = [[] for _ in range(len(LABEL_COLUMNS))]
+    true_labels = [[] for _ in range(len(LABEL_COLUMNS))]
+
+    found_all_files = True
+    for i, col in enumerate(LABEL_COLUMNS):
+        prob_file = os.path.join(model_preds_dir, f"{col}_probs.pkl")
+        true_file = os.path.join(model_preds_dir, f"{col}_true.pkl")
+        if os.path.exists(prob_file) and os.path.exists(true_file):
+            with open(prob_file, 'rb') as f:
+                all_probs[i] = pickle.load(f)
+            with open(true_file, 'rb') as f:
+                true_labels[i] = pickle.load(f)
+        else:
+            print(f"Warning: Prediction files not found for {model_name} - {col}. This model might be excluded for this label in ensemble.")
+            found_all_files = False # Mark that not all files were found
+
+    if not found_all_files:
+        return None, None # Indicate that this model's predictions couldn't be fully loaded
+
+    # Convert list of lists to list of numpy arrays if they were loaded as lists
+    # This ensures consistency for stacking later.
+    all_probs = [np.array(p) for p in all_probs]
+    true_labels = [np.array(t) for t in true_labels]
+
+    return all_probs, true_labels
+
+def perform_voting_ensemble(model_names_to_ensemble):
+    """
+    Performs a soft voting ensemble (averaging probabilities) for each label
+    across a list of specified models.
+
+    Args:
+        model_names_to_ensemble (list): A list of string names of the models
+                                        whose predictions should be ensembled.
+                                        These names should match the directory names
+                                        under `PREDICTIONS_SAVE_DIR`.
+
+    Returns:
+        tuple: A tuple containing:
+            - ensemble_reports (dict): Classification reports for the ensemble predictions.
+            - all_true_labels_for_ensemble (list): List of true labels used for evaluation.
+            - ensemble_predictions (list): List of predicted class indices from the ensemble.
+    """
+    print("\n--- Performing Voting Ensemble ---")
+    # defaultdict stores a list for each key, helpful when appending to potentially new keys
+    all_models_probs = defaultdict(list) # Stores list of probability arrays per label for all models
+    # Initialize with empty lists; true labels for evaluation (should be consistent across models)
+    all_true_labels_for_ensemble = [None for _ in range(len(LABEL_COLUMNS))]
+
+    # Load probabilities from all specified models
+    for model_name in model_names_to_ensemble:
+        print(f"Loading predictions for {model_name}...")
+        probs_per_label, true_labels_per_label = load_predictions(model_name)
+
+        if probs_per_label is None: # Skip this model if loading failed
+            continue
+
+        for i, col in enumerate(LABEL_COLUMNS):
+            if len(probs_per_label[i]) > 0: # Ensure probabilities were actually loaded for this label
+                all_models_probs[col].append(probs_per_label[i])
+                if all_true_labels_for_ensemble[i] is None: # Store true labels only once (they should be identical)
+                    all_true_labels_for_ensemble[i] = true_labels_per_label[i]
+
+    ensemble_predictions = [[] for _ in range(len(LABEL_COLUMNS))]
+    ensemble_reports = {}
+
+    for i, col in enumerate(LABEL_COLUMNS):
+        if not all_models_probs[col]: # If no models provided predictions for this label
+            print(f"No valid predictions available for {col} to ensemble. Skipping.")
+            ensemble_reports[col] = {'accuracy': 0, 'weighted avg': {'precision': 0, 'recall': 0, 'f1-score': 0, 'support': 0}}
+            continue
+
+        # Stack probabilities for the current label from all models that had them.
+        # `stacked_probs` will have shape: (num_contributing_models, num_samples, num_classes)
+        stacked_probs = np.stack(all_models_probs[col], axis=0)
+
+        # Perform soft voting by summing probabilities across models.
+        # `summed_probs` will have shape: (num_samples, num_classes)
+        summed_probs = np.sum(stacked_probs, axis=0)
+
+        # Get the final predicted class by taking the argmax of the summed probabilities.
+        final_preds = np.argmax(summed_probs, axis=1) # (num_samples,)
+
+        ensemble_predictions[i] = final_preds.tolist()
+
+        # Evaluate ensemble predictions
+        y_true_ensemble = all_true_labels_for_ensemble[i]
+        if y_true_ensemble is not None: # Ensure true labels are available
+            try:
+                report = classification_report(y_true_ensemble, final_preds, output_dict=True, zero_division=0)
+                ensemble_reports[col] = report
+            except ValueError:
+                print(f"Warning: Could not generate ensemble classification report for {col}. Skipping.")
+                ensemble_reports[col] = {'accuracy': 0, 'weighted avg': {'precision': 0, 'recall': 0, 'f1-score': 0, 'support': 0}}
+        else:
+            print(f"Warning: True labels not found for {col}, cannot evaluate ensemble.")
+            ensemble_reports[col] = {'accuracy': 0, 'weighted avg': {'precision': 0, 'recall': 0, 'f1-score': 0, 'support': 0}}
+
+
+    return ensemble_reports, all_true_labels_for_ensemble, ensemble_predictions