get_attention_weights.py

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

import pandas as pd
import torch
import torch.nn as nn
import numpy as np
import liwc
from tqdm import tqdm
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, AutoModel
import pickle
from src.text_preprocessing import get_embeddings_sequence_with_tokens, extract_liwc_sequence, load_and_preprocess_one_subset, get_need_cut, get_tokens_weights
from src.datasets import DeepPersonalityDataset, HandPersonalityDataset, collate_fn, deep_collate_fn
from src.training_utils import set_seed, 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 checkpoint (.pt)")
    parser.add_argument("--hc_model_path", type=str, required=True, help="Path to hand-crafted-based model checkpoint (.pt)")
    parser.add_argument("--fusion_model_path", type=str, required=True, help="Path to fusion model checkpoint (.pt)")
    parser.add_argument("--deep_model_architecture", type=str, default="ReBiLSTMAtt",
                        choices=["BiLSTMAtt", "ReBiLSTMAtt_v2", "ReBiLSTMAtt", "MambaAtt", "ReMambaAtt"])
    parser.add_argument("--hc_model_architecture", type=str, default="ReBiLSTMAtt",
                        choices=["BiLSTMAtt", "ReBiLSTMAtt_v2", "ReBiLSTMAtt", "MambaAtt", "ReMambaAtt"])
    parser.add_argument("--deep_encoder", type=str, default="xlm",
                        choices=["xlm", "bert", "bge", "jina-v3"])
    parser.add_argument("--dataset_path", type=str, default="src/prepered_dataframes/train_full_with_ASR.csv")
    parser.add_argument("--subset", type=str, default="train")

    # === 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")

    # === Misc ===
    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 get_attention_weights(args):
    set_seed(args.seed)
    device = torch.device(args.device)

    deep_model_architecture = get_model_classes(args.deep_model_architecture)
    deep_encoder = get_encoder_config(args.deep_encoder)
    hc_model_architecture = get_model_classes(args.hc_model_architecture)
    dataset_path = f"{args.dataset_path}"
    df = pd.read_csv(dataset_path)
    name_text_column = "text_ASR"
    name_class = ['OPE', 'CON', 'EXT', 'AGR', 'NNEU']
    text_ASR = df[name_text_column].tolist()

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

    # Load data
    texts, labels, names = load_and_preprocess_one_subset(df, name_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
    )
    print("Loading fusion models...")
    meta_model = fusion_model().to(device)
    meta_model.load_state_dict(torch.load(args.fusion_model_path, map_location=device))

    nn_model.train()
    nn_model.zero_grad()
    hc_model.train()
    hc_model.zero_grad()
    meta_model.train()
    meta_model.zero_grad()

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

    x_nn, idx_tokens, tokens = get_embeddings_sequence_with_tokens(texts, tokenizer_nn, encoder_nn, device)
    x_hc = extract_liwc_sequence(texts, category_text_features, parse_text_features)

    dataset_nn = DeepPersonalityDataset(x_nn, labels, idx_tokens, tokens, names)
    loader_nn = DataLoader(dataset_nn, batch_size=nn_config['batch_size'], collate_fn=deep_collate_fn)
    dataset_hc = HandPersonalityDataset(x_hc, labels)
    loader_hc = DataLoader(dataset_hc, batch_size=hc_config['batch_size'], collate_fn=collate_fn)

    dict_importanse_traits_all = {}

    need_cut = True

    for idx_text, (nn_data, hc_data) in tqdm(enumerate(zip(loader_nn,loader_hc))):
        batch_nn_features, batch_nn_labels, batch_nn_mask, batch_idx_tokens, batch_tokens, batch_names  = nn_data
        batch_hc_features, batch_hc_labels, batch_hc_mask  = hc_data

        for idx_pred in range(len(batch_nn_features)):
            tokens = batch_tokens[idx_pred]
            
            if need_cut:
                nn_need_cut = get_need_cut(batch_nn_features[idx_pred])
                hc_need_cut = get_need_cut(batch_hc_features[idx_pred])
            else:
                nn_need_cut = 0
                hc_need_cut = 0
                
            nn_features = batch_nn_features[idx_pred].to(device).unsqueeze(0).requires_grad_(True)

            hc_max_groups = torch.max(batch_hc_features[idx_pred], dim = 0).values.numpy()
            hc_features = batch_hc_features[idx_pred].to(device).unsqueeze(0).requires_grad_(True)
            
            nn_output, _ = nn_model(nn_features, mask=batch_nn_mask[idx_pred].unsqueeze(0).to(nn_features.device))
            hc_output, _ = hc_model(hc_features, mask=batch_hc_mask[idx_pred].unsqueeze(0).to(hc_features.device))
            final_output = meta_model(nn_output, hc_output)

            dict_importanse_traits = {}

            for target_class in range(5):
                ## work with nn
                nn_model.zero_grad()
                nn_output[:, target_class].backward(retain_graph=True)
                nn_grad = nn_features.grad.clone()
                nn_features.grad.zero_()
                
                if nn_need_cut != 0:
                    nn_pooled_grads = torch.mean(nn_grad[:, :-nn_need_cut, :], dim=1)
                    nn_weights = nn_features[:, :-nn_need_cut, :] * nn_pooled_grads.unsqueeze(1)
                else:
                    nn_pooled_grads = torch.mean(nn_grad, dim=1)
                    nn_weights = nn_features * nn_pooled_grads.unsqueeze(1)
                    
                nn_weights = nn_weights.detach().cpu().numpy()
                t, nn_weights = get_tokens_weights(tokens[1:-1], nn_weights[0][1:-1])
                if t:
                    nn_heatmap = torch.mean(torch.from_numpy(nn_weights), dim=1)
        
                ## work with hc
                hc_model.zero_grad()
                hc_output[:, target_class].backward(retain_graph=True)
                hc_grad = hc_features.grad.clone()
                hc_features.grad.zero_()
                hc_weights = hc_grad.detach().cpu().numpy()
                if hc_need_cut != 0:
                    hc_weights = hc_weights[:, :-hc_need_cut, :]
                hc_weights_sum = np.sum(hc_weights[0], axis=0)
                hc_weights_zeros = np.where(hc_weights_sum*hc_max_groups==0, 0, hc_weights_sum)
            
                dict_importanse_traits[name_class[target_class]] = {
                                            'nn': nn_heatmap.detach().cpu().numpy(),
                                            'hc': hc_weights_zeros,
                                            'hc_sum': hc_weights_sum,
                                            'tokens': t,
                                            'text': text_ASR[names.index(batch_names[idx_pred])],
                                            'normalized_text': texts[names.index(batch_names[idx_pred])],
                                            'nn_score': nn_output[:, target_class].detach().cpu().numpy()[0], 
                                            'hc_score': hc_output[:, target_class].detach().cpu().numpy()[0],
                                            'final_score': final_output[:, target_class].detach().cpu().numpy()[0],
                                            'true_score': batch_nn_labels[idx_pred][target_class].detach().cpu().numpy()
                                            }
        
            dict_importanse_traits_all[batch_names[idx_pred]] = dict_importanse_traits

    with open(f'{args.subset}_attention_weights.pickle', 'wb') as handle:
        pickle.dump(dict_importanse_traits_all, handle, protocol=pickle.HIGHEST_PROTOCOL)
    
    print("Attention weights were successfully saved")

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