quick_eval.py

import argparse
import json
import pickle
import numpy as np
import torch
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
from utils import probing_eval_prompt


def load_model_pickle(clf_path="model.pkl"):
    """
    Load back the model and metadata from pickle.
    """
    with open(clf_path, "rb") as f:
        package = pickle.load(f)
    clf = package["model"]
    metadata = package["metadata"]
    # print("Loaded metadata:", metadata)
    return clf, metadata


def extract_batch_reps(prompt_list, pool, tokenizer, model):
    inputs = tokenizer(
        prompt_list,
        return_tensors="pt",
        padding=True,
        truncation=True,
    )

    # Send inputs to the same device as the model embeddings
    device = next(model.parameters()).device
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        out = model(
            **inputs, 
            output_hidden_states=True, 
            output_attentions=True
        )
    
    hidden = out.hidden_states
    attentions = out.attentions
    reps = {}
    attn_entropy = {}
    
    for layer_idx, h in enumerate(hidden[1:]):
        if pool == "mean":
            res_pooled = h.mean(dim=1).cpu().detach().numpy()
        elif pool == "last":
            seq_lens = (inputs['attention_mask'].sum(dim=1) - 1)
            res_pooled = torch.stack([h[i, seq_lens[i], :] for i in range(h.size(0))]).cpu().detach().numpy()
        elif pool == "min":
            res_pooled = h.min(dim=1)[0].cpu().detach().numpy()
        elif pool == "max":
            res_pooled = h.max(dim=1)[0].cpu().detach().numpy()
        reps[layer_idx] = {pool: res_pooled}

        # Attention entropy
        A = attentions[layer_idx]
        mask = A > 0
        log_A = torch.where(mask, torch.log(A), torch.zeros_like(A))
        ent = - (A * log_A).sum(dim=-1).mean(dim=-1)
        attn_entropy[layer_idx] = ent.cpu().detach().numpy()
    
    return reps, attn_entropy


def extract_features(reps, attn_entropy, layers, pool, use_attn=False):
    features = []
    for l in layers:
        h = reps[l][pool]  # (B, H)
        if use_attn:
            attn = attn_entropy[l]  # (B, n_heads)
            layer_feat = np.hstack([h, attn])  # (B, target_dim_per_layer + n_heads)
        else:
            layer_feat = h
        features.append(layer_feat)
    features_final = np.hstack(features)  # (B, ...)
    return features_final


def evaluate_samples(samples, clf_root='gsm8k_multi_clfs', batch_size=16, 
                     model_name="Qwen/Qwen3-1.7B", top_percent=1.0):
    """
    Evaluate a list of samples and return evaluated results.
    
    Args:
        samples: List of dictionaries, each must have 'question' and 'prediction' keys
        clf_root: Path to classifier pickle root directory
        batch_size: Batch size for processing
        model_name: Huggingface model name for feature extraction
        top_percent: Fraction of top data to keep (1.0 = keep all)
    
    Returns:
        List of dictionaries with original fields plus evaluation scores:
        - semantic_consistency
        - logicality
        - informativeness
        - fluency
        - factuality
        - total_score
    """
    # Load model/tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        output_hidden_states=True,
        output_attentions=True,
        dtype=torch.float16 if torch.cuda.is_available() or torch.backends.mps.is_available() else torch.float32,
        device_map="auto",
        attn_implementation="eager",
        # cache_dir='local_models/',
    )
    model.eval()

    # Score each sample
    res = []
    classification = None  # Will be set from metadata
    for i in tqdm(range(0, len(samples), batch_size)):
        batch = samples[i:i+batch_size]
        tmp_res = batch.copy()
        for dim in ['semantic_consistency', 'logicality', 'informativeness', 'fluency', 'factuality']:
            # Load classifier
            clf_path = clf_root+f'/{dim}.pkl'       # 'gsm8k_multi_clfs/semantic_consistency.pkl'
            clf, metadata = load_model_pickle(clf_path)
            pool = metadata["pool"]
            use_attn = metadata["use_attn"]
            layers = metadata["layers"]
            if classification is None:
                classification = metadata["classification"]

            prompts = [probing_eval_prompt(sample['question'], sample['prediction'], dim) for sample in batch]
            reps, attn_entropy = extract_batch_reps(prompts, pool, tokenizer, model)
            features = extract_features(reps, attn_entropy, layers, pool, use_attn)
            del reps, attn_entropy
            torch.cuda.empty_cache()
            batch_scores = clf.predict(features)
            for j, sample in enumerate(batch):
                score = batch_scores[j]
                if hasattr(score, "item"):
                    tmp_res[j][dim] = score.item()
                else:
                    tmp_res[j][dim] = int(score)

        for s in tmp_res:   # total scores
            s['total_score'] = sum([s[dim] for dim in ['semantic_consistency', 'logicality', 'informativeness', 'fluency', 'factuality']])
        res.extend(tmp_res)
    
    # Sort and keep top X%
    if top_percent != 1.0:
        res.sort(reverse=True, key=lambda x: x['total_score'])
        n_keep = int(len(res) * top_percent)
        res = res[:n_keep]
    
    return res


def main():
    parser = argparse.ArgumentParser(description="Evaluate samples using trained classifiers.")
    parser.add_argument('--clf_root', type=str, default='gsm8k_multi_clfs', help='Path to classifier pickle root')
    parser.add_argument('--batch_size', '-bs', default=16, type=int, help="batch size")
    parser.add_argument('--file_path', type=str, default=None, help='Path to results JSON file (optional)')
    parser.add_argument('--top_percent', type=float, default=1.0, help='Fraction of top data to keep (e.g., 0.2 for 20%)')
    parser.add_argument('--model_name', type=str, default="Qwen/Qwen3-1.7B", help='Huggingface model name for feature extraction')
    parser.add_argument('--output_path', type=str, default=None, help='Output file path (optional, auto-generated if not provided)')
    args = parser.parse_args()
    
    # Load results from file if provided
    if args.file_path:
        with open(args.file_path, 'r') as f:
            samples = json.load(f)
    else:
        print("Error: --file_path is required when running as a script.")
        print("For programmatic use, import evaluate_samples() function and call it with a list.")
        return

    # Evaluate samples
    res = evaluate_samples(
        samples=samples,
        clf_root=args.clf_root,
        batch_size=args.batch_size,
        model_name=args.model_name,
        top_percent=args.top_percent
    )

    # Save
    if args.output_path:
        output_path = args.output_path
    else:
        model_id = args.model_name.split('/')[-1]
        # Get classification from first classifier metadata, multi/binary
        clf_path = args.clf_root + '/semantic_consistency.pkl'
        _, metadata = load_model_pickle(clf_path)
        classification = metadata["classification"]
        output_path = args.file_path.replace('.json', f'_{model_id}_{classification}_evaled.json')
    
    with open(output_path, 'w') as f:
        json.dump(res, f, indent=4)
    print(f"Saved {len(res)} samples to {output_path}")


if __name__ == "__main__":
    main()