get_explanation_with_LLM.py

import os
import argparse
import string
from collections import defaultdict
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib import transforms
from tqdm import tqdm
import pandas as pd
import pickle
import numpy as np
import re
import textwrap
import liwc
import torch
import csv
from transformers import AutoModelForCausalLM, AutoTokenizer
from src.training_utils import set_seed

def get_train_data(path):
    with open(path, "rb") as input1:
        dict_importance_files_train = pickle.load(input1)

    dict_traits_sum = {}
    dict_traits_valid_count = {}

    for k, v in tqdm(dict_importance_files_train.items()):
        for trait, values in v.items():
            hc_sum = np.array(values['hc_sum'], dtype=np.float32)
            mask = ~np.isnan(hc_sum) 
            hc_sum_safe = np.where(np.isnan(hc_sum), 0, hc_sum)

            if trait in dict_traits_sum:
                dict_traits_sum[trait] += hc_sum_safe
                dict_traits_valid_count[trait] += mask
            else:
                dict_traits_sum[trait] = hc_sum_safe
                dict_traits_valid_count[trait] = mask.astype(int)
    return dict_traits_sum


def get_test_data(path):
    with open(path, "rb") as f:
        return pickle.load(f)

def normalize_vector(vector):
    min_val = np.min(vector)
    max_val = np.max(vector)
    if max_val == min_val:
        return np.zeros_like(vector)
    normalized = 2 * ((vector - min_val) / (max_val - min_val)) - 1
    return np.round(normalized, 2)

def deep_normalize_attention(weights, all_weights, min_val=-1, max_val=1):
    min_w, max_w = np.min(weights), np.max(weights)
    if max_w == min_w:
        return np.zeros_like(weights)
    return (weights - min_w) / (max_w - min_w) * (max_val - min_val) + min_val

def hc_normalize_attention(vector):
    vector = np.array(vector)
    non_zero = vector[vector != 0]
    if len(non_zero) == 0:
        return np.zeros_like(vector)
    min_val, max_val = np.min(non_zero), np.max(non_zero)
    if min_val == max_val:
        normalized_non_zero = np.zeros_like(non_zero)
    else:
        normalized_non_zero = 2 * (non_zero - min_val) / (max_val - min_val) - 1
    result = vector.copy()
    result[result != 0] = normalized_non_zero
    return result

def get_word_indices(tokens):
    word_indices = defaultdict(list)
    for idx, word in enumerate(tokens):
        word_indices[word].append(idx)
    return word_indices

def get_hc_weights(tokens, weights, idx_sorted, sorted_categories, group_words):
    test = np.zeros(len(tokens))
    test_counter = np.zeros(len(tokens))
    test_weights_sort = [weights[i] for i in idx_sorted]
    word_indices = get_word_indices(tokens)
    word_to_category_ids = defaultdict(list)
    
    for category_id, category in enumerate(sorted_categories):
        words = group_words.get(category, [])
        if not words:
            continue
        for word in words:
            word_to_category_ids[word].append(category_id)
            for idx in word_indices.get(word, []):
                test[idx] += test_weights_sort[category_id]
                test_counter[idx] += 1

    result = np.zeros_like(test)
    mask = test_counter != 0
    result[mask] = test[mask] / test_counter[mask]

    return result, word_to_category_ids

def center_words_with_padding(words):
    max_len = max(len(word) for word in words) if words else 1
    return [word.center(max_len, ' ') for word in words]

def remove_punctuation(word):
    return word.translate(str.maketrans('', '', string.punctuation))

def rainbow_text(x, y, strings, colors, ax=None, **kw):
    if ax is None:
        ax = plt.gca()
    t = ax.transData
    canvas = ax.figure.canvas
    for s, c in zip(strings, colors):
        text = ax.text(x, y, s + " ", color=c, transform=t, fontsize=12, **kw)
        text.draw(canvas.get_renderer())
        ex = text.get_window_extent()
        t = transforms.offset_copy(text._transform, x=ex.width, units='dots')

def get_dict_split_texts(tokens, parse_text_features, category_text_features):
    dict_split_text = defaultdict(list)
    for token in tokens:
        for cat in parse_text_features(token):
            if cat in category_text_features:
                dict_split_text[cat].append(token)
    dict_split_text_new = {k: list(set(v)) for k, v in dict_split_text.items()}
    for cat in category_text_features:
        dict_split_text_new.setdefault(cat, None)
    return dict_split_text_new

def visualize_importance_new(
    scale_weights, deep_weights, hand_weights, hand_weights_2,
    word_to_category_ids, sorted_categories, trait, tokens,
    pred_score, text_traits, group_words, transcription, indent_sizes,
    need_words=3, cof_weights=0.60, draw=True, cmap='coolwarm', output_dir="."
):
    mask1 = hand_weights != 0.0
    masked_hand_weights = np.ma.masked_where(~mask1, hand_weights)
    mask2 = hand_weights_2 != 0.0
    masked_hand_weights2 = np.ma.masked_where(~mask2, hand_weights_2)

    group_words_order = [
        ','.join([str(j) for j in sorted(word_to_category_ids.get(token, []), reverse=True)][:need_words]) or ''
        for token in tokens
    ]

    deep_words_red = [tokens[i] for i in range(len(deep_weights)) if deep_weights[i] >= cof_weights]
    deep_words_blue = [tokens[i] for i in range(len(deep_weights)) if deep_weights[i] <= -cof_weights]
    hand_words_red = [tokens[i] for i in range(len(hand_weights_2)) if hand_weights_2[i] >= cof_weights]
    hand_words_blue = [tokens[i] for i in range(len(hand_weights_2)) if hand_weights_2[i] <= -cof_weights]

    trait_labels = {
        'OPE': ("$\widetilde{FM}^D_O$", "$\widetilde{FM}^H_O$", "$\widetilde{GFM}^H_O$"),
        'CON': ("$\widetilde{FM}^D_C$", "$\widetilde{FM}^H_C$", "$\widetilde{GFM}^H_C$"),
        'EXT': ("$\widetilde{FM}^D_E$", "$\widetilde{FM}^H_E$", "$\widetilde{GFM}^H_E$"),
        'AGR': ("$\widetilde{FM}^D_A$", "$\widetilde{FM}^H_A$", "$\widetilde{GFM}^H_A$"),
        'NNEU': ("$\widetilde{FM}^D_N$", "$\widetilde{FM}^H_N$", "$\widetilde{GFM}^H_N$"),
    }
    ldt, lht, ght = trait_labels.get(trait, ("D", "H", "GH"))

    if draw:
        left_sig_pad = indent_sizes['left_sig_pad']
        explai_pad = indent_sizes['explai_pad']
        id_cat_pad = (len(max(group_words_order, key=len)) + 2) / 100 if group_words_order else 0.02
        longest_word = (len(max(tokens, key=len)) + 2) / 200 if tokens else 0.02
        cize = max(10, round((22 * len(scale_weights)) / 64))
        fig = plt.figure(figsize=(cize, 5))
        gs = gridspec.GridSpec(5, 1, height_ratios=[0.04, longest_word, 0.04, id_cat_pad, explai_pad], hspace=0.3)

        ax1 = fig.add_subplot(gs[0])
        ax_label_1 = fig.add_subplot(gs[1])
        ax2 = fig.add_subplot(gs[2])
        ax_label_2 = fig.add_subplot(gs[3])
        ax3 = fig.add_subplot(gs[4])

        ax_label_1.axis('off')
        ax_label_2.axis('off')
        
        im1 = ax1.imshow(deep_weights[np.newaxis, :], aspect='auto', cmap=cmap, vmin=-1, vmax=1)
        ax1.text(len(tokens)/2-6, -0.7, f"Importance of word for {trait[0]} trait", fontsize=12)
        ax1.text(-left_sig_pad, +0.75, f'Importance scale for\nlocal heatmap of\ndeep features ({ldt})\nPred. score: {pred_score[0]:.2f}', fontsize=10)
        ax1.set_xticks(np.arange(len(tokens)))
        ax1.set_xticklabels(center_words_with_padding(tokens), rotation=90)
        ax1.text(-left_sig_pad, +1.5, "Word:", fontsize=9)
        ax1.set_yticks([])

        im2 = ax2.imshow(masked_hand_weights2[np.newaxis, :], aspect='auto', cmap=cmap, vmin=-1, vmax=1)
        ax2.text(-left_sig_pad, +0.75, f'Importance scale for\nlocal heatmap of\nhand-crafted features ({lht})\nPred. score: {pred_score[1]:.2f}', fontsize=10)
        ax2.set_xticks(np.arange(len(tokens)))
        ax2.set_xticklabels(group_words_order, rotation=90)
        ax2.tick_params(axis='x', which='both', top=True, bottom=True, labeltop=False, labelbottom=True)
        ax2.text(-left_sig_pad, +1.85, "Category IDs:", fontsize=9)
        ax2.set_yticks([])

        im3 = ax3.imshow(scale_weights[np.newaxis, :], aspect='auto', cmap=cmap, vmin=-1, vmax=1)
        ax3.set_xticks(np.arange(len(sorted_categories)))
        ax3.set_xticklabels(sorted_categories, rotation=90)
        ax3.text(0, -1.3, text_traits[0], fontsize=9)
        ax3.text(len(sorted_categories)-1, -1.3, text_traits[1], fontsize=9, horizontalalignment='right')
        ax3.text(-9.1, -0.75, "Category ID:", fontsize=9)
        ax3.text(-9.1, +1.3, "Category:", fontsize=9)
        ax3.text(-9.1, +0.60, f'Importance scale for\nglobal heatmap of\nhand-crafted features ({ght})', fontsize=10)
        ax3.set_yticks([])

        ax3_top = ax3.twiny()
        ax3_top.set_xlim(ax3.get_xlim())
        ax3_top.set_xticks(np.arange(len(sorted_categories)))
        ax3_top.set_xticklabels(list(range(len(sorted_categories))), rotation=0)

        transcriptions = transcription.split('\n')
        for i, trans in enumerate(transcriptions):
            text = f"Input transcription: {trans}" if i == 0 else trans
            colors = [
                'red' if remove_punctuation(j).lower() in set(deep_words_red + hand_words_red) else 'black'
                for j in text.split()
            ]
            y_pos = 3.2 if i == 0 else 3.6
            rainbow_text(0 if i == 0 else 6.7, y_pos, text.split(), colors, ax=ax3)

    # Explanation
    top_positive = [sorted_categories[i] for i in np.argsort(hand_weights)[-3:] if hand_weights[i] > cof_weights]
    top_negative = [sorted_categories[i] for i in np.argsort(hand_weights)[:3] if hand_weights[i] < -cof_weights]
    
    pos_words_str = ", ".join(sorted(set(deep_words_red + hand_words_red))) or "no distinct key words"
    neg_words_str = ", ".join(sorted(set(deep_words_blue + hand_words_blue))) or "no significant negative factors"
    pos_cats_str = ", ".join(sorted(top_positive)) or "no highlighted categories"
    neg_cats_str = ", ".join(sorted(top_negative)) or "no highlighted categories"

    if pred_score[2] > 0.55:
        explanation = (
            f"The fused predicted score ({pred_score[2]:.2f}) of {trait[0]} trait is high. It indicates that the person tends to be {text_traits[1].split(': ')[1].lower()}.\n"
            f"                     - Key words: {pos_words_str}.\n"
            f"                     - Significant categories: {pos_cats_str}.\n"
            f"                     These factors indicate a strong presence of {trait[0]}."
        )
    elif pred_score[2] < 0.45:
        explanation = (
            f"The fused predicted score ({pred_score[2]:.2f}) of {trait[0]} trait is low. It indicates that the person is more likely to be {text_traits[0].split(': ')[1].lower()}.\n"
            f"                     - Words that lowered the score: {neg_words_str}.\n"
            f"                     - Categories with a negative impact: {neg_cats_str}.\n"
            f"                     These factors indicate a weak presence of {trait[0]}."
        )
    else:
        explanation = (
            f"The fused predicted score ({pred_score[2]:.2f}) of {trait[0]} trait is medium. It indicates that the person exhibits characteristics ({text_traits[0].split(': ')[1].lower()} and {text_traits[1].split(': ')[1].lower()}) of both opposing traits.\n"
            f"                     - Positive words: {pos_words_str}.\n"
            f"                     - Negative words: {neg_words_str}.\n"
            f"                     - Important categories: {pos_cats_str}.\n"
            f"                     The score falls in the intermediate range due to a balanced contribution of factors."
        )

    if draw:
        ax3.text(0, 5.9, f"Explanation: {explanation}", fontsize=12, wrap=True)
        cbar = fig.colorbar(im3, ax=[ax1, ax2, ax3], orientation='vertical', fraction=0.04, pad=0.01)
        os.makedirs(output_dir, exist_ok=True)
        plt.savefig(os.path.join(output_dir, f"text_visualisation_{trait}.png"), bbox_inches='tight')
        plt.close(fig)

    explanation_full = (
        f'Input transcription: {transcription}.\n'
        f'      Attention of tokens and categories for {trait[0]} trait:\n'
        f'      Tokens: {tokens}.\n'
        f'      ID word categories for each token: {group_words_order}.\n'
        f'      Deep weights for tokens: {deep_weights}.\n'
        f'      Hand-crafted weights for tokens: {masked_hand_weights2}.\n'
        f'      Categories: {sorted_categories}.\n'
        f'      Category weights: {scale_weights}.\n'
        f"The deep and hand-crafted weights show what each model focuses on for this example. The category weights reflect which word groups globally help the hand-crafted model predict high or low {trait[0]} scores. Each word can correspond to several categories. "
        f'Based on this, the deep model predicts {pred_score[0]:.2f}, the hand-crafted model predicts {pred_score[1]:.2f}, and the fusion model predicts {pred_score[2]:.2f} for {trait[0]} trait. '
        f'I can explain such predictions using only high/low attention: {explanation}. But that is insufficient for accuracy. For {trait[0]} trait, a new score and a deeper explanation are required.'
    )
    
    return explanation, explanation_full

set_seed(42)

def main():
    parser = argparse.ArgumentParser(description="Generate explanations for all 5 traits and refine with LLM.")

    # --- Paths ---
    parser.add_argument("--test_csv", type=str, required=True, help="Path to test CSV with labels")
    parser.add_argument("--train_weights", type=str, required=True, help="Path to train attention weights pickle")
    parser.add_argument("--test_weights", type=str, required=True, help="Path to test attention weights pickle")
    parser.add_argument("--liwc_path", type=str, default="LIWC2007.txt", help="Path to LIWC dictionary")
    parser.add_argument("--output_dir", type=str, default="visualizations", help="Directory to save plots and outputs")

    # --- Target ---
    parser.add_argument("--video_name", type=str, required=True, help="Video name to process")

    # --- LLM ---
    parser.add_argument("--llm_model_path", type=str, default="tiiuae/Falcon-H1-7B-Instruct", help="Path to LLM")
    parser.add_argument("--run_llm", action="store_true", help="Run LLM refinement after generating explanations")

    # --- Visualization params ---
    parser.add_argument("--need_words", type=int, default=4, help="Max category IDs per token to show")
    parser.add_argument("--cof_weights", type=float, default=0.6, help="Threshold for significant words")
    parser.add_argument("--left_sig_pad", type=float, default=6.45, help="Padding for left text")
    parser.add_argument("--explai_pad", type=float, default=0.05, help="Padding for explanation")

    args = parser.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    # === 1. Load data ===
    test_df = pd.read_csv(args.test_csv)
    if args.video_name not in test_df['video_name'].values:
        raise ValueError(f"Video {args.video_name} not found in test CSV")

    row = test_df[test_df['video_name'] == args.video_name].iloc[0]
    text_ASR = row['text_ASR']
    true_scores = [
        row['openness'],
        row['conscientiousness'],
        row['extraversion'],
        row['agreeableness'],
        row['non-neuroticism']
    ]

    # Load weights
    dict_train = get_train_data(args.train_weights)
    dict_test = get_test_data(args.test_weights)
    if args.video_name not in dict_test:
        raise ValueError(f"Video {args.video_name} not found in test weights")

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

    # === 2. Trait configs ===
    TRAIT_KEYS = ["OPE", "CON", "EXT", "AGR", "NNEU"]
    TRAIT_NAMES = ["openness", "conscientiousness", "extraversion", "agreeableness", "non-neuroticism"]
    TRAIT_TEXT = {
        'OPE': ['Low score: Traditional, Conventional', 'High score: Philosophical, Reflective'],
        'CON': ['Low score: Spontaneous, Impulsive', 'High score: Organized, Systematic'],
        'EXT': ['Low score: Reserved, Quiet', 'High score: Assertive, Enthusiastic'],
        'AGR': ['Low score: Argumentative, Demanding', 'High score: Polite, Empathetic'],
        'NNEU': ['Low score: Sensitive, Tense', 'High score: Relaxed, Resilient'],
    }

    # === 3. Generate explanations for all traits ===
    all_explanations = {}
    indent_sizes = {"left_sig_pad": args.left_sig_pad, "explai_pad": args.explai_pad}

    for trait_key, trait_name in zip(TRAIT_KEYS, TRAIT_NAMES):
        print(f"\nProcessing trait: {trait_key} ({trait_name})")
        curr_test = dict_test[args.video_name]
        text = curr_test[trait_key]['text']
        text = re.sub(r'\.\s*\.', '.', text)
        text = textwrap.fill(text, width=150)
        tokens = [t.lower() for t in curr_test[trait_key]['tokens']]
        if not tokens:
            print(f"No tokens for {trait_key}")
            continue

        group_words = get_dict_split_texts(tokens, parse_text_features, category_text_features)

        # Global weights
        idx_sorted = np.argsort(dict_train[trait_key])
        sorted_categories = [category_text_features[i] for i in idx_sorted]
        sorted_values = normalize_vector([dict_train[trait_key][i] for i in idx_sorted])

        # Local weights
        nn_full_weights = []
        for t in TRAIT_KEYS:
            nn_full_weights.extend(curr_test[t]['nn'])
        nn_weights = deep_normalize_attention(np.array(curr_test[trait_key]['nn'][:-4]), np.array(nn_full_weights[:-4]))

        hc_weights_by_category = hc_normalize_attention(np.array(curr_test[trait_key]['hc']))
        hc_weights_by_category = [hc_weights_by_category[i] for i in idx_sorted]

        hc_weights_by_word, word_to_category_ids = get_hc_weights(
            tokens, curr_test[trait_key]['hc'], idx_sorted, sorted_categories, group_words
        )
        hc_weights_by_word = hc_normalize_attention(hc_weights_by_word)[:-4]

        # Visualize and get explanation
        explanation_short, explanation_full = visualize_importance_new(
            sorted_values,
            nn_weights,
            hc_weights_by_category,
            hc_weights_by_word,
            word_to_category_ids,
            sorted_categories,
            trait_key,
            tokens[:-4],
            [
                curr_test[trait_key]['nn_score'],
                curr_test[trait_key]['hc_score'],
                curr_test[trait_key]['final_score']
            ],
            TRAIT_TEXT[trait_key],
            group_words,
            text,
            indent_sizes,
            need_words=args.need_words,
            cof_weights=args.cof_weights,
            draw=True,
            output_dir=os.path.join(args.output_dir, args.video_name)
        )

        all_explanations[trait_key] = {
            "explanation": explanation_short,
            "final_score": float(curr_test[trait_key]['final_score'])
        }

        print(f"  → Short explanation: {explanation_short[:100]}...")

    # === 4. Save explanations for LLM ===
    llm_input = {args.video_name: all_explanations}
    llm_input_path = os.path.join(args.output_dir, f"{args.video_name}_explanations.pickle")
    with open(llm_input_path, "wb") as f:
        pickle.dump(llm_input, f)
    print(f"\nAll explanations saved to: {llm_input_path}")

    # === 5. Run LLM refinement (optional) ===
    if args.run_llm:
        print("\nRunning LLM refinement...")
        refined_scores, response_clean = run_llm_refinement(
            video_name=args.video_name,
            text_ASR=text_ASR,
            explanations=all_explanations,
            llm_path=args.llm_model_path,
            output_dir=args.output_dir
        )

        # Save LLM result
        csv_path = os.path.join(args.output_dir, f"{args.video_name}_llm_refined.csv")
        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f)
            writer.writerow([
                "name_video", "text_ASR",
                "openness", "conscientiousness", "extraversion", "agreeableness", "non_neuroticism",
                "response_clean"
            ])
            writer.writerow([args.video_name, text_ASR, *refined_scores, response_clean])
        print(f"LLM result saved to: {csv_path}")

    print("\nDone!")


def run_llm_refinement(video_name, text_ASR, explanations, llm_path, output_dir):
    """Runs LLM on a single example using your prompt."""
    TRAIT_KEYS = ["OPE", "CON", "EXT", "AGR", "NNEU"]
    scores = [explanations[trait]["final_score"] for trait in TRAIT_KEYS]
    explanations_list = [explanations[trait]["explanation"] for trait in TRAIT_KEYS]

    model = AutoModelForCausalLM.from_pretrained(
        llm_path,
        torch_dtype=torch.bfloat16,
        # device_map="auto"
    )

    model = torch.compile(model)
    tokenizer = AutoTokenizer.from_pretrained(llm_path)

    prompt = f"""
You calibrate preliminary Big Five score. Do not predict personality from scratch.
Always transform the preliminary scores. All five final scores must differ from
the corresponding preliminary scores. 

Use this priority:
1) preliminary score = anchor;
2) trait explanation = main signal for calibration magnitude;
3) transcript = weak limiter only: it may reduce the magnitude by one step, but must not eliminate the correction and must not change its sign.

Use a centered two-slope calibration pattern with center c = 0.5.

For each trait independently:
- let x be the preliminary score;
- if x < 0.5, the correction must be negative;
- if x >= 0.5, the correction must be positive.

Allowed corrections:
- if x < 0.5, d must be one of {-0.018, -0.012, -0.006};
- if x >= 0.5, d must be one of {+0.006, +0.012}.

Magnitude rule:
- the explanation selects the weaker or stronger allowed correction;
- the transcript may only reduce the magnitude by one step;
- the transcript must not change the sign;
- d = 0 is forbidden.

Update rule:
- set y = clip(x + d, 0, 1).

Hard rules:
- all five traits must be adjusted;
- every final value must differ from its preliminary value by at least 0.006 unless clipping prevents it;
- direction is determined only by whether x is below or at/above 0.5;
- explanations control magnitude only, not direction;
- the transcript may weaken magnitude only, not direction;
- do not leave any trait unchanged;
- do not copy any preliminary value exactly unless clipping makes this unavoidable;
- keep all final values inside [0,1].

Output exactly one line with five comma-separated values in this order:\nOpenness, Conscientiousness, Extraversion, Agreeableness, Non-neuroticism

Write each value with exactly five digits after the decimal point. Do not output labels, explanations, reasoning, or any extra text.

Transcript:\n{text_ASR}\nPreliminary scores:\n{scores}\nTrait explanations:\n{explanations_list}
"""

    messages = [{"role": "user", "content": prompt}]
    text_prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
    )

    inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)
    generated_ids = model.generate(**inputs, max_new_tokens=512)
    response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)

    if "assistant" in response:
        response_clean = response.split("assistant")[-1].strip()
    else:
        response_clean = response.strip()

    numbers = re.findall(r"\d\.\d+", response_clean)
    if len(numbers) == 5:
        refined_scores = [float(n) for n in numbers]
    else:
        refined_scores = [0.5] * 5  # fallback

    return refined_scores, response_clean


if __name__ == "__main__":
    main()