Upload 9 files

Dockerfile

@@ -0,0 +1,49 @@
+# 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/saved_models \
+    /app/predictions \
+    /app/tokenizer \
+    /app/cache \
+    && chmod -R 777 /app/uploads \
+    /app/saved_models \
+    /app/predictions \
+    /app/tokenizer \
+    /app/cache
+
+# Copy the application code and utilities
+COPY . /app/
+COPY ../dataset_utils.py /app/
+COPY ../train_utils.py /app/
+COPY ../config.py /app/
+COPY ../label_encoders.pkl /app/
+
+# Set environment variables
+ENV PYTHONPATH=/app
+ENV PYTHONUNBUFFERED=1
+ENV PORT=7860
+
+# Expose the port the app runs on
+EXPOSE 7860
+
+# Command to run the application
+CMD ["python", "app.py"] 

app.py

@@ -0,0 +1,283 @@
+from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File, Form
+from fastapi.responses import FileResponse
+from pydantic import BaseModel
+from typing import Optional, Dict, Any, List
+import uvicorn
+import logging
+import os
+import pandas as pd
+from datetime import datetime
+import shutil
+from pathlib import Path
+import numpy as np
+import sys
+import json
+import joblib
+
+# Import existing utilities
+from dataset_utils import (
+    load_and_preprocess_data,
+    save_label_encoders,
+    load_label_encoders
+)
+from config import (
+    TEXT_COLUMN,
+    LABEL_COLUMNS,
+    BATCH_SIZE,
+    MODEL_SAVE_DIR
+)
+from tfidf_based_models.tfidf_lgbm import TfidfLightGBM
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="LGBM Compliance Predictor API")
+
+UPLOAD_DIR = Path("uploads")
+MODEL_SAVE_DIR = Path("saved_models")
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+MODEL_SAVE_DIR.mkdir(parents=True, exist_ok=True)
+
+training_status = {
+    "is_training": False,
+    "current_epoch": 0,
+    "total_epochs": 0,
+    "current_loss": 0.0,
+    "start_time": None,
+    "end_time": None,
+    "status": "idle",
+    "metrics": None
+}
+
+class TrainingConfig(BaseModel):
+    batch_size: int = 32
+    num_epochs: int = 1  # Not used for LGBM, but kept for API compatibility
+    random_state: int = 42
+
+class TrainingResponse(BaseModel):
+    message: str
+    training_id: str
+    status: str
+    download_url: Optional[str] = None
+
+class ValidationResponse(BaseModel):
+    message: str
+    metrics: Dict[str, Any]
+    predictions: List[Dict[str, Any]]
+
+class TransactionData(BaseModel):
+    Transaction_Id: str
+    Message: str
+    # ... (other fields as needed) ...
+
+class PredictionRequest(BaseModel):
+    transaction_data: TransactionData
+    model_name: str = "lgbm_models"  # Default to tfidf_lgbm if not specified
+
+class BatchPredictionResponse(BaseModel):
+    message: str
+    predictions: List[Dict[str, Any]]
+    metrics: Optional[Dict[str, Any]] = None
+
+@app.get("/")
+async def root():
+    return {"message": "LGBM Compliance Predictor API"}
+
+@app.get("/v1/lgbm/health")
+async def health_check():
+    return {"status": "healthy"}
+
+@app.get("/v1/lgbm/training-status")
+async def get_training_status():
+    return training_status
+
+@app.post("/v1/lgbm/train", response_model=TrainingResponse)
+async def start_training(
+    config: str = Form(...),
+    background_tasks: BackgroundTasks = None,
+    file: UploadFile = File(...)
+):
+    if training_status["is_training"]:
+        raise HTTPException(status_code=400, detail="Training is already in progress")
+    if not file.filename.endswith('.csv'):
+        raise HTTPException(status_code=400, detail="Only CSV files are allowed")
+    try:
+        config_dict = json.loads(config)
+        training_config = TrainingConfig(**config_dict)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid config parameters: {str(e)}")
+    file_path = UPLOAD_DIR / file.filename
+    with file_path.open("wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+    training_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+    training_status.update({
+        "is_training": True,
+        "current_epoch": 0,
+        "total_epochs": 1,
+        "start_time": datetime.now().isoformat(),
+        "status": "starting"
+    })
+    background_tasks.add_task(train_model_task, training_config, str(file_path), training_id)
+    download_url = f"/v1/lgbm/download-model/{training_id}"
+    return TrainingResponse(
+        message="Training started successfully",
+        training_id=training_id,
+        status="started",
+        download_url=download_url
+    )
+
+@app.post("/v1/lgbm/validate")
+async def validate_model(
+    file: UploadFile = File(...),
+    model_name: str = "lgbm_models"
+):
+    if not file.filename.endswith('.csv'):
+        raise HTTPException(status_code=400, detail="Only CSV files are allowed")
+    try:
+        file_path = UPLOAD_DIR / file.filename
+        with file_path.open("wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+        data_df, label_encoders = load_and_preprocess_data(str(file_path))
+        model_path = MODEL_SAVE_DIR / f"{model_name}.pkl"
+        if not model_path.exists():
+            raise HTTPException(status_code=404, detail="LGBM model file not found")
+        model = TfidfLightGBM(label_encoders)
+        model.load_model(model_name)
+        X = data_df[TEXT_COLUMN]
+        y = data_df[LABEL_COLUMNS]
+        reports, y_true_list, y_pred_list = model.evaluate(X, y)
+        all_probs = model.predict_proba(X)
+        predictions = []
+        for i, col in enumerate(LABEL_COLUMNS):
+            label_encoder = label_encoders[col]
+            true_labels_orig = label_encoder.inverse_transform(y_true_list[i])
+            pred_labels_orig = label_encoder.inverse_transform(y_pred_list[i])
+            for true, pred, probs in zip(true_labels_orig, pred_labels_orig, all_probs[i]):
+                class_probs = {label: float(prob) for label, prob in zip(label_encoder.classes_, probs)}
+                predictions.append({
+                    "field": col,
+                    "true_label": true,
+                    "predicted_label": pred,
+                    "probabilities": class_probs
+                })
+        return ValidationResponse(
+            message="Validation completed successfully",
+            metrics=reports,
+            predictions=predictions
+        )
+    except Exception as e:
+        logger.error(f"Validation failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Validation failed: {str(e)}")
+    finally:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+
+@app.post("/v1/lgbm/predict")
+async def predict(
+    request: Optional[PredictionRequest] = None,
+    file: UploadFile = File(None),
+    model_name: str = "lgbm_models"
+):
+    try:
+        model_path = MODEL_SAVE_DIR / f"{model_name}.pkl"
+        if not model_path.exists():
+            raise HTTPException(status_code=404, detail=f"Model {model_name} not found")
+        label_encoders = load_label_encoders()
+        model = TfidfLightGBM(label_encoders)
+        model.load_model(model_name)
+        # Batch prediction
+        if file and file.filename:
+            if not file.filename.endswith('.csv'):
+                raise HTTPException(status_code=400, detail="Only CSV files are allowed")
+            file_path = UPLOAD_DIR / file.filename
+            with file_path.open("wb") as buffer:
+                shutil.copyfileobj(file.file, buffer)
+            try:
+                data_df, _ = load_and_preprocess_data(str(file_path))
+                X = data_df[TEXT_COLUMN]
+                all_probabilities = model.predict_proba(X)
+                predictions = []
+                for i, row in data_df.iterrows():
+                    transaction_pred = {}
+                    for j, col in enumerate(LABEL_COLUMNS):
+                        probs = all_probabilities[j][i]
+                        pred = np.argmax(probs)
+                        decoded_pred = label_encoders[col].inverse_transform([pred])[0]
+                        class_probs = {label: float(probs[j]) for j, label in enumerate(label_encoders[col].classes_)}
+                        transaction_pred[col] = {
+                            "prediction": decoded_pred,
+                            "probabilities": class_probs
+                        }
+                    predictions.append({
+                        "transaction_id": row.get('Transaction_Id', f"transaction_{i}"),
+                        "predictions": transaction_pred
+                    })
+                return BatchPredictionResponse(
+                    message="Batch prediction completed successfully",
+                    predictions=predictions
+                )
+            finally:
+                if os.path.exists(file_path):
+                    os.remove(file_path)
+        # Single prediction
+        elif request and request.transaction_data:
+            input_data = pd.DataFrame([request.transaction_data.dict()])
+            X = input_data[TEXT_COLUMN]
+            all_probabilities = model.predict_proba(X)
+            response = {}
+            for i, col in enumerate(LABEL_COLUMNS):
+                probs = all_probabilities[i][0]
+                pred = np.argmax(probs)
+                decoded_pred = label_encoders[col].inverse_transform([pred])[0]
+                class_probs = {label: float(probs[j]) for j, label in enumerate(label_encoders[col].classes_)}
+                response[col] = {
+                    "prediction": decoded_pred,
+                    "probabilities": class_probs
+                }
+            return response
+        else:
+            raise HTTPException(
+                status_code=400,
+                detail="Either provide a transaction in the request body or upload a CSV file"
+            )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/v1/lgbm/download-model/{model_id}")
+async def download_model(model_id: str):
+    model_path = MODEL_SAVE_DIR / f"{model_id}.pkl"
+    if not model_path.exists():
+        raise HTTPException(status_code=404, detail="Model not found")
+    return FileResponse(
+        path=model_path,
+        filename=f"lgbm_model_{model_id}.pkl",
+        media_type="application/octet-stream"
+    )
+
+async def train_model_task(config: TrainingConfig, file_path: str, training_id: str):
+    try:
+        data_df_original, label_encoders = load_and_preprocess_data(file_path)
+        save_label_encoders(label_encoders)
+        X = data_df_original[TEXT_COLUMN]
+        y = data_df_original[LABEL_COLUMNS]
+        model = TfidfLightGBM(label_encoders)
+        model.train(X, y)
+        model.save_model(training_id)
+        training_status.update({
+            "is_training": False,
+            "end_time": datetime.now().isoformat(),
+            "status": "completed"
+        })
+    except Exception as e:
+        logger.error(f"Training failed: {str(e)}")
+        training_status.update({
+            "is_training": False,
+            "end_time": datetime.now().isoformat(),
+            "status": "failed",
+            "error": str(e)
+        })
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    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

models/tfidf_lgbm.py

@@ -0,0 +1,163 @@
+# tfidf_based_models/tfidf_lgbm.py
+
+from sklearn.feature_extraction.text import TfidfVectorizer
+import lightgbm as lgb
+from sklearn.pipeline import Pipeline
+from sklearn.metrics import classification_report
+from sklearn.preprocessing import LabelEncoder
+from sklearn.utils.class_weight import compute_class_weight
+import numpy as np
+import pandas as pd
+import joblib
+import os
+
+from config import TEXT_COLUMN, LABEL_COLUMNS, TFIDF_MAX_FEATURES, MODEL_SAVE_DIR
+
+class TfidfLightGBM:
+    """
+    TF-IDF based LightGBM model for multi-output classification.
+    It trains a separate LightGBM classifier for each target label
+    after converting text data into TF-IDF features.
+    """
+    def __init__(self, label_encoders):
+        """
+        Initializes the TfidfLightGBM model.
+
+        Args:
+            label_encoders (dict): A dictionary of LabelEncoder objects.
+        """
+        self.label_encoders = label_encoders
+        self.models = {} # Stores the trained Pipeline for each label
+
+    def train(self, X_train_text, y_train_df):
+        """
+        Trains a TF-IDF + LightGBM pipeline for each label.
+
+        Args:
+            X_train_text (pd.Series): Training text data.
+            y_train_df (pd.DataFrame): DataFrame of training labels (encoded).
+        """
+        print("Training TF-IDF + LightGBM models...")
+        for i, col in enumerate(LABEL_COLUMNS):
+            print(f"  Training for {col}...")
+            num_classes = len(self.label_encoders[col].classes_)
+            # Determine LightGBM objective based on number of classes
+            objective = 'multiclass' if num_classes > 2 else 'binary'
+            # `num_class` parameter is required for 'multiclass' objective
+            num_class_param = {'num_class': num_classes} if num_classes > 2 else {}
+
+            pipeline = Pipeline([
+                ('tfidf', TfidfVectorizer(max_features=TFIDF_MAX_FEATURES)),
+                ('lgbm', lgb.LGBMClassifier(
+                    objective=objective,
+                    **num_class_param, # Unpack num_class_param if it's not empty
+                    random_state=42,
+                    n_estimators=100
+                ))
+            ])
+            # Fit the pipeline on the training data.
+            # LightGBM handles class imbalance with `is_unbalance=True` or `scale_pos_weight`
+            # for binary classification, or implicitly for multiclass with default settings.
+            pipeline.fit(X_train_text, y_train_df[col])
+            self.models[col] = pipeline
+        print("TF-IDF + LightGBM training complete.")
+
+    def predict(self, X_test_text):
+        """
+        Makes class predictions for all labels.
+
+        Args:
+            X_test_text (pd.Series): Test text data.
+
+        Returns:
+            dict: A dictionary where keys are label names and values are NumPy arrays
+                  of predicted class indices.
+        """
+        predictions = {}
+        for col, model_pipeline in self.models.items():
+            predictions[col] = model_pipeline.predict(X_test_text)
+        return predictions
+
+    def predict_proba(self, X_test_text):
+        """
+        Returns prediction probabilities for each class for all labels.
+
+        Args:
+            X_test_text (pd.Series): Test text data.
+
+        Returns:
+            list: A list of NumPy arrays. Each array corresponds to a label column
+                  and contains the probability distribution over classes for each sample.
+        """
+        probabilities = []
+        for col in LABEL_COLUMNS:
+            if col in self.models:
+                probabilities.append(self.models[col].predict_proba(X_test_text))
+            else:
+                print(f"Warning: Model for {col} not found, cannot predict probabilities.")
+                probabilities.append(np.array([]))
+        return probabilities
+
+    def evaluate(self, X_test_text, y_test_df):
+        """
+        Evaluates the models and returns classification reports.
+
+        Args:
+            X_test_text (pd.Series): Test text data.
+            y_test_df (pd.DataFrame): DataFrame of true test labels (encoded).
+
+        Returns:
+            tuple: A tuple containing:
+                - reports (dict): Classification reports for each label column.
+                - truths (list): List of true label arrays.
+                - preds (list): List of predicted label arrays.
+        """
+        reports = {}
+        truths = [[] for _ in range(len(LABEL_COLUMNS))]
+        preds = [[] for _ in range(len(LABEL_COLUMNS))]
+
+        for i, col in enumerate(LABEL_COLUMNS):
+            if col in self.models:
+                y_pred = self.models[col].predict(X_test_text)
+                y_true = y_test_df[col].values
+                try:
+                    report = classification_report(y_true, y_pred, output_dict=True, zero_division=0)
+                    reports[col] = report
+                except ValueError:
+                    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}}
+                truths[i].extend(y_true)
+                preds[i].extend(y_pred)
+            else:
+                print(f"Warning: Model for {col} not found for evaluation.")
+
+        return reports, truths, preds
+
+    def save_model(self, model_name="tfidf_lgbm", save_format='pickle'):
+        """
+        Saves the trained TF-IDF LightGBM models.
+
+        Args:
+            model_name (str): The base name for the saved model file.
+            save_format (str): Format to save the model in (default: 'pickle').
+        """
+        if save_format != 'pickle':
+            raise ValueError("TF-IDF models only support 'pickle' format")
+        save_path = os.path.join(MODEL_SAVE_DIR, f"{model_name}.pkl")
+        joblib.dump(self.models, save_path)
+        print(f"TF-IDF LightGBM models saved to {save_path}")
+
+    def load_model(self, model_name="tfidf_lgbm"):
+        """
+        Loads trained TF-IDF LightGBM models from a file.
+
+        Args:
+            model_name (str): The base name of the model file to load.
+        """
+        load_path = os.path.join(MODEL_SAVE_DIR, f"{model_name}.pkl")
+        if os.path.exists(load_path):
+            self.models = joblib.load(load_path)
+            print(f"TF-IDF LightGBM models loaded from {load_path}")
+        else:
+            print(f"Error: Model file not found at {load_path}. Initialize models as empty.")
+            self.models = {}

saved_models/lgbm_models.pkl

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

saved_models/tfidf_vectorizer.pkl

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

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