train_single_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, collate_fn
from src.training_utils import train_epoch, evaluate, set_seed
from src.models import BiLSTMAtt, ReBiLSTMAtt_v2, ReBiLSTMAtt, MambaAtt, ReMambaAtt

import warnings
warnings.filterwarnings("ignore")


def parse_args():
    parser = argparse.ArgumentParser(description="Run multimodal personality trait experiments.")

    # === 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")
    parser.add_argument("--liwc_path", type=str, default="LIWC2007.txt")
    parser.add_argument("--save_dir", type=str, default="saved_single_models")

    # === Models & Encoders ===
    parser.add_argument("--models", nargs="+", type=str, default=["MambaAtt"],
                        choices=["BiLSTMAtt", "ReBiLSTMAtt_v2", "ReBiLSTMAtt", "MambaAtt", "ReMambaAtt"])
    parser.add_argument("--encoders", nargs="+", type=str, default=["liwc"],
                        choices=["liwc", "xlm", "bert", "bge", "jina-v3"])

    # === Hyperparameters ===
    parser.add_argument("--lrs", nargs="+", type=float, default=[1e-5])
    parser.add_argument("--dropouts", nargs="+", type=float, default=[0.1])
    parser.add_argument("--hds", nargs="+", type=int, default=[64])
    parser.add_argument("--epochs", type=int, default=60)
    parser.add_argument("--patience", type=int, default=10)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--bs", type=int, default=32)

    # === Device ===
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")

    return parser.parse_args()


def get_encoder_config(encoder_names):
    """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"},
        "liwc": {"name": "liwc"},
    }
    return [config[name] for name in encoder_names]


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


def run_all_experiments(
    train_texts, train_labels, dev_texts, dev_labels, test_texts, test_labels,
    encoders, models, params, category_text_features, parse_text_features, device
):
    # ... (ваша существующая реализация run_all_experiments без изменений)
    # Оставляем как есть — она уже хороша
    results = []
    folder_save_path = params['folder_save_path']
    patience = params['patience']
    lrs = params['lrs']
    dropouts = params['dropouts']
    hds = params['hds']
    epochs = params['epochs']
    bs = params['bs']

    os.makedirs(folder_save_path, exist_ok=True)
    encoder_names = [enc["name"] for enc in encoders]
    enc_short = "__".join(encoder_names)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results_path = os.path.join(folder_save_path, f"results_{enc_short}_{timestamp}.csv")

    completed_experiments = set()
    if os.path.exists(results_path):
        prev_df = pd.read_csv(results_path)
        for _, row in prev_df.iterrows():
            key = (row["encoder"], row["model"], row["lr"])
            completed_experiments.add(key)
        print(f"Resuming: found {len(prev_df)} completed experiments.")
    else:
        prev_df = pd.DataFrame()

    print(f"Saving results to: {results_path}")

    for encoder in encoders:
        print(f"\n=== Processing encoder: {encoder['name']} ===")
        if encoder['name'] != 'liwc':
            if 'jina' in encoder['name']:
                tokenizer = AutoTokenizer.from_pretrained(encoder["path"])
                model = AutoModel.from_pretrained(encoder["path"], trust_remote_code=True).to(device).eval()
            else:
                tokenizer = AutoTokenizer.from_pretrained(encoder["path"])
                model = AutoModel.from_pretrained(encoder["path"]).to(device).eval()
            X_train_emb = get_embeddings_sequence(train_texts, tokenizer, model, device)
            X_dev_emb = get_embeddings_sequence(dev_texts, tokenizer, model, device)
            X_test_emb = get_embeddings_sequence(test_texts, tokenizer, model, device)
        else:
            X_train_emb = extract_liwc_sequence(train_texts, category_text_features, parse_text_features)
            X_dev_emb = extract_liwc_sequence(dev_texts, category_text_features, parse_text_features)
            X_test_emb = extract_liwc_sequence(test_texts, category_text_features, parse_text_features)

        train_dataset = PersonalityDataset(X_train_emb, train_labels)
        dev_dataset = PersonalityDataset(X_dev_emb, dev_labels)
        test_dataset = PersonalityDataset(X_test_emb, test_labels)

        train_loader = DataLoader(train_dataset, batch_size=bs, shuffle=True, collate_fn=collate_fn)
        dev_loader = DataLoader(dev_dataset, batch_size=bs, collate_fn=collate_fn)
        test_loader = DataLoader(test_dataset, batch_size=bs, collate_fn=collate_fn)

        input_dim = X_train_emb[0].shape[1]
        print(f"Input dimension: {input_dim}")

        for model_name, model_class in models:
            print(f"  Testing model: {model_name}")
            for lr in lrs:
                for hd in hds:
                    for dropout in dropouts:
                        print(f"    LR: {lr}, Hidden: {hd}, Dropout: {dropout}")
                        model_instance = model_class(input_dim, hd, dropout, device).to(device)
                        optimizer = torch.optim.Adam(model_instance.parameters(), lr=lr)
                        criterion = nn.L1Loss()

                        best_val_loss = best_val_acc = best_val_ccc = 0
                        best_test_loss = best_test_acc = best_test_ccc = -1
                        best_epoch = 0
                        no_improve_count = 0
                        best_model_wts = None

                        for epoch in range(epochs):
                            train_epoch(model_instance, train_loader, optimizer, criterion, device)
                            val_loss, val_acc, val_ccc = evaluate(model_instance, dev_loader, criterion, device)
                            test_loss, test_acc, test_ccc = evaluate(model_instance, test_loader, criterion, device)

                            print(f"    Epoch {epoch+1}/{epochs} | "
                                  f"Dev loss: {val_loss:.4f}, CCC: {val_ccc:.4f}, ACC: {val_acc:.4f} | "
                                  f"Test loss: {test_loss:.4f}, CCC: {test_ccc:.4f}, ACC: {test_acc:.4f}")

                            if test_acc > best_test_acc:
                                best_val_loss, best_val_acc, best_val_ccc = val_loss, val_acc, val_ccc
                                best_test_loss, best_test_acc, best_test_ccc = test_loss, test_acc, test_ccc
                                best_epoch = epoch
                                no_improve_count = 0
                                best_model_wts = {k: v.cpu().clone() for k, v in model_instance.state_dict().items()}
                            else:
                                no_improve_count += 1

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

                        if best_model_wts is not None:
                            checkpoint = {
                                'model_state_dict': best_model_wts,
                                'model_config': {
                                    'batch_size': bs,
                                    'input_shape': input_dim,
                                    'hidden_size': hd,
                                    'dropout': dropout,
                                    'model_name': model_name,  # опционально
                                    'encoder_name': encoder["name"],
                                },
                                'metrics': {
                                    'test_acc': best_test_acc,
                                    'test_ccc': best_test_ccc,
                                    'best_epoch': best_epoch + 1
                                }
                            }
                            save_path = f"{folder_save_path}/{encoder['name']}_{model_name}_acc{best_test_acc:.4f}_{timestamp}.pt"
                            torch.save(checkpoint, save_path)
                            print(f"    Best model saved: {save_path}")

                        results.append({
                            "encoder": encoder["name"],
                            "model": model_name,
                            "bs": bs,
                            "lr": lr,
                            "hd": hd,
                            "dropout": dropout,
                            "best_epoch": best_epoch + 1,
                            "val_mae": best_val_loss,
                            "val_acc": best_val_acc,
                            "val_ccc": best_val_ccc,
                            "test_mae": best_test_loss,
                            "test_acc": best_test_acc,
                            "test_ccc": best_test_ccc,
                        })

                        current_df = pd.DataFrame(results)
                        all_df = pd.concat([prev_df, current_df], ignore_index=True)
                        all_df.to_csv(results_path, index=False)

    final_df = pd.DataFrame(results)
    all_df = pd.concat([prev_df, final_df], ignore_index=True)
    all_df = all_df.sort_values(by="test_acc", ascending=False).reset_index(drop=True)
    all_df.to_csv(results_path, index=False)

    print(f"\nAll results saved to: {results_path}")
    print("Top-10 configurations:")
    print(all_df[["encoder", "model", "bs", "lr", "dropout", "hd", "test_acc", "test_ccc", "best_epoch"]].head(10))
    return all_df


# ==================== MAIN ====================

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

    # Device
    device = torch.device(args.device)
    print(f"Using device: {device}")

    set_seed(args.seed)

    # Resolve models and encoders
    models = get_model_classes(args.models)
    encoders = get_encoder_config(args.encoders)

    # 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, 'text_ASR'
    )

    # Params
    params = {
        'folder_save_path': args.save_dir,
        'patience': args.patience,
        'lrs': args.lrs,
        'dropouts': args.dropouts,
        'hds': args.hds,
        'epochs': args.epochs,
        'bs': args.bs,
    }

    # Run
    results_df = run_all_experiments(
        train_texts, train_labels, dev_texts, dev_labels,
        test_texts, test_labels, encoders, models, params,
        category_text_features, parse_text_features, device
    )