train_fusion_models.py

import os
import argparse
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import pandas as pd
import torch
import torch.nn as nn
import liwc
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, AutoModel
from datetime import datetime
from src.text_preprocessing import get_embeddings_sequence, extract_liwc_sequence, load_and_preprocess_data
from src.datasets import PersonalityDataset, EnsembleDataset, collate_fn
from src.training_utils import evaluate_fusion_model, set_seed, get_predictions, load_model_from_checkpoint
from src.models import BiLSTMAtt, ReBiLSTMAtt_v2, ReBiLSTMAtt, MambaAtt, ReMambaAtt, fusion_model

import warnings
warnings.filterwarnings("ignore")


def parse_args():
    parser = argparse.ArgumentParser(description="Train ensemble fusion model from two base models.")

    # === Checkpoint paths ===
    parser.add_argument("--nn_model_path", type=str, required=True, help="Path to deep-based model base model checkpoint (.pt)")
    parser.add_argument("--hc_model_path", type=str, required=True, help="Path to hand-crafted-based base model checkpoint (.pt)")
    parser.add_argument("--deep_model_architecture", type=str, default="MambaAtt",
                        choices=["BiLSTMAtt", "ReBiLSTMAtt_v2", "ReBiLSTMAtt", "MambaAtt", "ReMambaAtt"])
    parser.add_argument("--hc_model_architecture", type=str, default="MambaAtt",
                        choices=["BiLSTMAtt", "ReBiLSTMAtt_v2", "ReBiLSTMAtt", "MambaAtt", "ReMambaAtt"])
    parser.add_argument("--deep_encoder", type=str, default="xlm",
                        choices=["xlm", "bert", "bge", "jina-v3"])

    # === Data paths ===
    parser.add_argument("--train_csv", type=str, default="src/prepered_dataframes/train_full_with_ASR.csv")
    parser.add_argument("--dev_csv", type=str, default="src/prepered_dataframes/dev_full_with_ASR.csv")
    parser.add_argument("--test_csv", type=str, default="src/prepered_dataframes/test_full_with_ASR.csv")

    # === Text column ===
    parser.add_argument("--text_column", type=str, default="text_ASR", help="Name of the text column in CSV")

    # === LIWC ===
    parser.add_argument("--liwc_path", type=str, default="LIWC2007.txt", help="Path to LIWC dictionary file")

    # === Training hyperparameters ===
    parser.add_argument("--lr", type=float, default=1e-2, help="Learning rate for fusion model")
    parser.add_argument("--bs", type=int, default=128, help="Batch size for fusion model training")
    parser.add_argument("--epochs", type=int, default=500, help="Max number of epochs")
    parser.add_argument("--patience", type=int, default=100, help="Patience for early stopping")

    # === Misc ===
    parser.add_argument("--save_dir", type=str, default="saved_fusion_models", help="Directory to save fusion model and results")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda/cpu)")

    return parser.parse_args()

def get_encoder_config(encoder_name):
    """Map encoder names to config dicts."""
    config = {
        "xlm": {"name": "xlm", "path": "FacebookAI/xlm-roberta-base"},
        "bert": {"name": "bert", "path": "google-bert/bert-base-multilingual-cased"},
        "bge": {"name": "bge", "path": "BAAI/bge-large-en"},
        "jina-v3": {"name": "jina-v3", "path": "jinaai/jina-embeddings-v3"},
    }
    return config[encoder_name]["path"]

def get_model_classes(model_name):
    """Map model names to classes."""
    mapping = {
        "BiLSTMAtt": BiLSTMAtt,
        "ReBiLSTMAtt_v2": ReBiLSTMAtt_v2,
        "ReBiLSTMAtt": ReBiLSTMAtt,
        "MambaAtt": MambaAtt,
        "ReMambaAtt": ReMambaAtt,
    }
    return mapping[model_name]

def train_ensemble_model(args):
    set_seed(args.seed)
    device = torch.device(args.device)
    os.makedirs(args.save_dir, exist_ok=True)

    deep_model_architecture = get_model_classes(args.deep_model)
    deep_encoder = get_encoder_config(args.deep_encoder)

    hc_model_architecture = get_model_classes(args.hc_model)

    # Load LIWC
    parse_text_features, category_text_features = liwc.load_token_parser(args.liwc_path)
    category_text_features = sorted(category_text_features)

    # Load data
    train_texts, train_labels, dev_texts, dev_labels, test_texts, test_labels = load_and_preprocess_data(
        args.train_csv, args.dev_csv, args.test_csv, args.text_column
    )

    # Load base models
    print("Loading base models...")
    nn_model, nn_config = load_model_from_checkpoint(
        checkpoint_path=args.nn_model_path,
        model_class=deep_model_architecture,
        device=device
    )
    hc_model, hc_config = load_model_from_checkpoint(
        checkpoint_path=args.hc_model_path,
        model_class=hc_model_architecture,
        device=device
    )

    # Extract embeddings
    print("Extracting embeddings...")
    tokenizer_nn = AutoTokenizer.from_pretrained(deep_encoder)
    encoder_nn = AutoModel.from_pretrained(deep_encoder).to(device).eval()

    X_train_nn = get_embeddings_sequence(train_texts, tokenizer_nn, encoder_nn, device)
    X_dev_nn = get_embeddings_sequence(dev_texts, tokenizer_nn, encoder_nn, device)
    X_test_nn = get_embeddings_sequence(test_texts, tokenizer_nn, encoder_nn, device)

    X_train_hc = extract_liwc_sequence(train_texts, category_text_features, parse_text_features)
    X_dev_hc = extract_liwc_sequence(dev_texts, category_text_features, parse_text_features)
    X_test_hc = extract_liwc_sequence(test_texts, category_text_features, parse_text_features)

    train_loader_nn = DataLoader(PersonalityDataset(X_train_nn, train_labels), batch_size=nn_config['batch_size'], shuffle=False, collate_fn=collate_fn)
    dev_loader_nn = DataLoader(PersonalityDataset(X_dev_nn, dev_labels), batch_size=nn_config['batch_size'], collate_fn=collate_fn)
    test_loader_nn = DataLoader(PersonalityDataset(X_test_nn, test_labels), batch_size=nn_config['batch_size'], collate_fn=collate_fn)

    train_loader_hc = DataLoader(PersonalityDataset(X_train_hc, train_labels), batch_size=hc_config['batch_size'], shuffle=False, collate_fn=collate_fn)
    dev_loader_hc = DataLoader(PersonalityDataset(X_dev_hc, dev_labels), batch_size=hc_config['batch_size'], collate_fn=collate_fn)
    test_loader_hc = DataLoader(PersonalityDataset(X_test_hc, test_labels), batch_size=hc_config['batch_size'], collate_fn=collate_fn)

    # Get predictions
    print("Generating base predictions...")
    pred_train_nn = get_predictions(nn_model, train_loader_nn, device)
    pred_dev_nn = get_predictions(nn_model, dev_loader_nn, device)
    pred_test_nn = get_predictions(nn_model, test_loader_nn, device)

    pred_train_hc = get_predictions(hc_model, train_loader_hc, device)
    pred_dev_hc = get_predictions(hc_model, dev_loader_hc, device)
    pred_test_hc = get_predictions(hc_model, test_loader_hc, device)

    # Ensemble datasets
    train_ens_dataset = EnsembleDataset(pred_train_nn, pred_train_hc, torch.tensor(train_labels, dtype=torch.float32))
    dev_ens_dataset = EnsembleDataset(pred_dev_nn, pred_dev_hc, torch.tensor(dev_labels, dtype=torch.float32))
    test_ens_dataset = EnsembleDataset(pred_test_nn, pred_test_hc, torch.tensor(test_labels, dtype=torch.float32))

    train_ens_loader = DataLoader(train_ens_dataset, batch_size=args.bs, shuffle=True)
    dev_ens_loader = DataLoader(dev_ens_dataset, batch_size=args.bs, shuffle=False)
    test_ens_loader = DataLoader(test_ens_dataset, batch_size=args.bs, shuffle=False)

    # Train fusion model
    print("Training fusion model...")
    meta_model = fusion_model().to(device)
    optimizer = torch.optim.Adam(meta_model.parameters(), lr=args.lr)
    criterion = nn.L1Loss()

    best_test_acc = -1
    best_epoch = 0
    no_improve = 0
    best_weights = None
    best_metrics = {}

    for epoch in range(args.epochs):
        # Train
        meta_model.train()
        for p1, p2, y in train_ens_loader:
            p1, p2, y = p1.to(device), p2.to(device), y.to(device)
            optimizer.zero_grad()
            pred = meta_model(p1, p2)
            loss = criterion(pred, y)
            loss.backward()
            optimizer.step()

        # Evaluate
        meta_model.eval()
        dev_loss, dev_acc, dev_ccc = evaluate_fusion_model(meta_model, dev_ens_loader, criterion, device)
        test_loss, test_acc, test_ccc = evaluate_fusion_model(meta_model, test_ens_loader, criterion, device)

        print(f"Epoch {epoch+1:3d}/{args.epochs} | "
              f"Dev ACC: {dev_acc:.4f}, CCC: {dev_ccc:.4f} | "
              f"Test ACC: {test_acc:.4f}, CCC: {test_ccc:.4f}")

        if test_acc > best_test_acc:
            best_test_acc = test_acc
            best_epoch = epoch
            no_improve = 0
            best_weights = {k: v.cpu().clone() for k, v in meta_model.state_dict().items()}
            best_metrics = {
                "val_mae": dev_loss, "val_acc": dev_acc, "val_ccc": dev_ccc,
                "test_mae": test_loss, "test_acc": test_acc, "test_ccc": test_ccc
            }
        else:
            no_improve += 1

        if no_improve >= args.patience:
            print(f"Early stopping at epoch {epoch+1}")
            break

    # Save
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    if best_weights:
        save_path = os.path.join(args.save_dir, f"ensemble_meta_b5_acc{best_test_acc:.4f}_{timestamp}.pt")
        torch.save(best_weights, save_path)
        print(f"Best ensemble model saved: {save_path}")

    # Save results
    result_row = {
        "encoder": "xlm+liwc",
        "model": "ensemble_b5",
        "lr": args.lr,
        "batch_size": args.bs,
        "patience": args.patience,
        "best_epoch": best_epoch + 1,
        "val_mae": best_metrics["val_mae"],
        "val_acc": best_metrics["val_acc"],
        "val_ccc": best_metrics["val_ccc"],
        "test_mae": best_metrics["test_mae"],
        "test_acc": best_metrics["test_acc"],
        "test_ccc": best_metrics["test_ccc"],
    }

    results_df = pd.DataFrame([result_row])
    csv_path = os.path.join(args.save_dir, f"ensemble_results_{timestamp}.csv")
    results_df.to_csv(csv_path, index=False)
    print(f"\nFinal results saved to: {csv_path}")
    print(results_df.T)

    return results_df


if __name__ == "__main__":
    args = parse_args()
    train_ensemble_model(args)