generate_logits_chat.py

import json
import os
os.environ['TRANSFORMERS_CACHE'] = os.environ['HF_HOME']
os.environ['HF_DATASETS_CACHE'] = os.environ['HF_HOME']
import random
import torch
import argparse
import pickle
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import LlamaForCausalLM, GenerationConfig
from tqdm import tqdm
import prompt as pt


seed = 1
random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False


few_shot_exp_ids = {
    "MMLU": [1, 3, 5, 7, 9],
    "HellaSwag": [1, 3, 5, 7, 9],
    "CosmosQA": [1, 3, 5, 7, 9],
    "Halu-OpenDialKG": [5, 7, 9],
    "Halu-CNN/DailyMail": [9]
}

options = ["Answer: A", "Answer: B", "Answer: C", "Answer: D", "Answer: E", "Answer: F"]
options_alt = ["\nA", "\nB", "\nC", "\nD", "\nE", "\nF"]

def load_data(data_file):
    data = json.load(open(data_file, "r"))
    return data

def load_model(args):
    if "Yi" in args.model:
        tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=False)
    else:
        tokenizer = AutoTokenizer.from_pretrained(args.model)
    tokenizer.truncation_side = "left"
    tokenizer.model_max_length = min(tokenizer.model_max_length, 2048)
    
    if "falcon" in args.model:
        model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, device_map="auto")
    elif "deepseek" in args.model:
        model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, device_map="auto")
        model.generation_config = GenerationConfig.from_pretrained(args.model)
        model.generation_config.pad_token_id = model.generation_config.eos_token_id
    elif "Llama" in args.model:
        model = LlamaForCausalLM.from_pretrained(args.model, torch_dtype=torch.float16, device_map="auto")
    elif "Qwen" in args.model:
        model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, device_map="auto")
    elif "Yi" in args.model:
        model = AutoModelForCausalLM.from_pretrained(args.model, device_map="auto", torch_dtype="auto")
    elif "gemma" in args.model.lower():  # Add this for Gemma support
        model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, device_map="auto")
    else:
        raise NotImplementedError
    
    model.eval()
    return tokenizer, model

def get_fewshot_exps(data):
    src = data[0]["source"]
    fewshot_exps = []
    for idx in few_shot_exp_ids[src]:
        fewshot_exps.append(data[idx])
        assert data[idx]["id"] == idx
    return fewshot_exps

def format_example(example, prompt, with_answer=False, num_options=4):
    # QA
    if example["source"] == "MMLU":
        prompt += "Question: " + example["question"] + "\nChoices:\n"
    # Reading Comprehension
    elif example["source"] == "CosmosQA":
        prompt += "Context: " + example["context"] + "\n" + "Question: " + example["question"] + "\nChoices:\n"
    # Commonsense NLI
    elif example["source"] == "HellaSwag":
        prompt += "Context: " + example["context"] + "\n" + "Question: " + example["question"] + "\nChoices:\n"
    # Dialogue Response
    elif example["source"] == "Halu-OpenDialKG":
        prompt += "Dialogue: " + example["context"] + "\n" + "Question: " + example["question"] + "\nChoices:\n"
    # Document Summarization
    elif example["source"] == "Halu-CNN/DailyMail":
        prompt += "Document: " + example["context"] + "\n" + "Question: " + example["question"] + "\nChoices:\n"
    else:
        raise NotImplementedError("Not supported dataset.")
    # Only include first num_options choices (e.g., A, B, C, D)
    valid_options = ["A", "B", "C", "D", "E", "F"][:num_options]
    for k, v in example["choices"].items():
        if k in valid_options:
            prompt += k + ". " + str(v) + "\n"
    prompt += "Answer:"
    if with_answer:
        prompt += " " + example["answer"] + "\n"
    return prompt

def format_base_prompt(example, args, tokenizer, fewshot_exps=None):
    exp = {}
    exp["id"] = example["id"]
    if args.few_shot == 0 and not args.cot:
        prompt = ""
    elif args.few_shot > 0 and not args.cot:
        prompt = ""
        for fs_exp in fewshot_exps:
            prompt = format_example(fs_exp, prompt, with_answer=True, num_options=args.num_options)
    elif args.few_shot == 0 and args.cot:
        prompt = pt.base_cot_prompt
    else:
        raise NotImplementedError("Not supported method.")
    prompt = format_example(example, prompt, num_options=args.num_options)
    
    # We treat the prompt message by now as the user input
    if "falcon" in args.model:
        prompt = "User: " + prompt + "\n" + "Assistant:"
    else:
        message = [
            {"role": "user", "content": prompt}
        ]
        prompt = tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)

    exp["prompt"] = prompt
    return exp

def format_shared_prompt(example, args, tokenizer, fewshot_exps=None):
    exp = {}
    exp["id"] = example["id"]
    if args.few_shot == 0 and not args.cot:
        prompt = pt.shared_zero_prompt
    elif args.few_shot > 0 and not args.cot:
        prompt = pt.shared_few_prompt
        for fs_exp in fewshot_exps:
            prompt = format_example(fs_exp, prompt, with_answer=True, num_options=args.num_options)
        prompt += "\nNow make your best effort and select the correct answer for the following question. You only need to output the option.\n\n"
    elif args.few_shot == 0 and args.cot:
        prompt = pt.shared_cot_prompt
    else:
        raise NotImplementedError("Not supported method.")
    prompt = format_example(example, prompt, num_options=args.num_options)
    
    # We treat the prompt message by now as the user input
    if "falcon" in args.model:
        prompt = "User: " + prompt + "\n" + "Assistant:"
    else:
        message = [
            {"role": "user", "content": prompt}
        ]
        prompt = tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)

    exp["prompt"] = prompt
    return exp

def format_task_prompt(example, args, tokenizer, fewshot_exps=None):
    exp = {}
    exp["id"] = example["id"]
    if args.few_shot == 0 and not args.cot:
        pt_dict = json.loads(pt.task_zero_prompt, strict=False)
        prompt = pt_dict[example["source"]]
    elif args.few_shot > 0 and not args.cot:
        pt_dict = json.loads(pt.task_few_prompt, strict=False)
        prompt = pt_dict[example["source"]]
        for fs_exp in fewshot_exps:
            prompt = format_example(fs_exp, prompt, with_answer=True, num_options=args.num_options)
        prompt += "\nNow make your best effort and select the correct answer for the following question. You only need to output the option.\n\n"
    elif args.few_shot == 0 and args.cot:
        pt_dict = json.loads(pt.task_cot_prompt, strict=False)
        prompt = pt_dict[example["source"]]
    else:
        raise NotImplementedError("Not supported method.")
    prompt = format_example(example, prompt, num_options=args.num_options)

    # We treat the prompt message by now as the user input
    if "falcon" in args.model:
        prompt = "User: " + prompt + "\n" + "Assistant:"
    else:
        message = [
            {"role": "user", "content": prompt}
        ]
        prompt = tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
    
    exp["prompt"] = prompt
    return exp

def prepare_inputs(tokenizer, exp, device):
    """Prepare inputs and move to the correct device."""
    inputs = tokenizer(exp["prompt"], return_tensors="pt", truncation=True)
    # Move all tensors to the model's device
    inputs = {k: v.to(device) for k, v in inputs.items()}
    return inputs

def log_softmax(logits):
    logits = logits - max(logits)
    return F.log_softmax(logits, dim=0)


def extract_activations(model, inputs, verbose=True):
    """
    Extract MLP activations from the final token using forward hooks.
    """
    collected = []

    def hook_fn(module, inp, out):
        if isinstance(out, torch.Tensor) and out.ndim == 3:
            # Immediately move to CPU and detach to free GPU memory
            collected.append(out[:, -1, :].detach().cpu().clone())

    # Handle Falcon's different architecture
    model_type = type(model).__name__
    if model_type == "FalconForCausalLM":
        layers = model.transformer.h  # Falcon uses transformer.h instead of model.layers
    else:
        layers = model.model.layers  # Works for Llama, Qwen, DeepSeek, Gemma

    handles = [
        layer.mlp.register_forward_hook(hook_fn)
        for layer in layers if hasattr(layer, "mlp")
    ]

    with torch.no_grad():
        outputs = model(**inputs)

    for h in handles:
        h.remove()

    # Stack and immediately free the list
    activations = torch.stack(collected)  # (L, B, H)
    collected.clear()

    return outputs, activations

# def extract_activations(model, inputs, verbose=True):
#     """
#     Extract residual stream hidden states from the final token using forward hooks.
#     """
#     collected = []

#     def hook_fn(module, inp, out):
#         # Transformer blocks return (hidden_state, past_key_value, ...) tuples
#         hidden = out[0] if isinstance(out, tuple) else out
#         if isinstance(hidden, torch.Tensor) and hidden.ndim == 3:
#             collected.append(hidden[:, -1, :].detach().cpu().clone())

#     model_type = type(model).__name__
#     if model_type == "FalconForCausalLM":
#         layers = model.transformer.h
#     else:
#         layers = model.model.layers

#     handles = [
#         layer.register_forward_hook(hook_fn)  # block, not layer.mlp
#         for layer in layers
#     ]

#     with torch.no_grad():
#         outputs = model(**inputs)

#     for h in handles:
#         h.remove()

#     activations = torch.stack(collected).squeeze(1)  # (L, H)
#     collected.clear()

#     return outputs, activations


def get_model_outputs(model, tokenizer, data, args):
    all_outputs = []
    device = next(model.parameters()).device
    batch_size = args.batch_size

    if "Yi" in args.model:
        option_ids = [tokenizer.encode(opt)[-1] for opt in options_alt[:args.num_options]]
    else:
        option_ids = [tokenizer.encode(opt)[-1] for opt in options[:args.num_options]]

    # Left-padding so last token is always at position -1 for all samples
    original_padding_side = tokenizer.padding_side
    tokenizer.padding_side = "left"
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    for batch_start in tqdm(range(0, len(data), batch_size)):
        batch = data[batch_start:batch_start + batch_size]
        prompts = [exp["prompt"] for exp in batch]

        inputs = tokenizer(prompts, return_tensors="pt", truncation=True, padding=True)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        if args.extract_activations:
            verbose = (batch_start == 0)
            outputs, activations = extract_activations(model, inputs, verbose=verbose)
        else:
            with torch.no_grad():
                outputs = model(**inputs)
            activations = None

        # Logits at last position (left-padded, so -1 is the real last token)
        batch_logits = outputs.logits[:, -1, :].detach().cpu()  # (B, vocab)

        for i, exp in enumerate(batch):
            logits_options = batch_logits[i][option_ids]

            out = {}
            out["id"] = exp["id"]
            out["logits_options"] = logits_options.float().numpy()

            if args.extract_activations and activations is not None:
                out["activations"] = activations[:, i, :]  # (L, H)

            all_outputs.append(out)

        del outputs, batch_logits, inputs
        if activations is not None:
            del activations

    tokenizer.padding_side = original_padding_side

    return all_outputs


def validate_batching(model, tokenizer, data, args, n=10):
    """Compare batch_size=1 vs batch_size=N on first n samples to verify correctness."""
    subset = data[:min(n, len(data))]
    orig_batch = args.batch_size

    print(f"\n  Validating batching (batch_size={orig_batch} vs 1) on {len(subset)} samples...")

    args.batch_size = 1
    outputs_single = get_model_outputs(model, tokenizer, subset, args)

    args.batch_size = orig_batch
    outputs_batched = get_model_outputs(model, tokenizer, subset, args)

    max_diff = 0
    for s, b in zip(outputs_single, outputs_batched):
        diff = abs(s["logits_options"] - b["logits_options"]).max()
        max_diff = max(max_diff, float(diff))

    print(f"  Max logit difference: {max_diff:.2e}")
    if max_diff < 1e-3:
        print(f"  PASS — batching is safe\n")
    else:
        print(f"  WARNING — differences too large, falling back to batch_size=1\n")
        args.batch_size = 1


def main(args):
    if args.file != "xxx.json":
        all_data_files = [args.file]
    else:
        all_data_files = ['mmlu_10k.json', 'cosmosqa_10k.json', 'hellaswag_10k.json',
                          'halu_dialogue.json', 'halu_summarization.json']
    print(all_data_files)

    # Pre-check which files need processing before loading the model
    files_to_run = []
    for file in all_data_files:
        save_file = args.model.split("/")[-1] + "_" + file.split(".json")[0] + "_" + args.prompt_method
        save_file += "_icl" + str(args.few_shot)
        if args.cot:
            save_file += "_cot"
        save_file = os.path.join(args.output_dir, save_file)
        os.makedirs(args.output_dir, exist_ok=True)

        pkl_filename = save_file + (f"_{args.pkl_suffix}.pkl" if args.pkl_suffix else ".pkl")
        act_filename = save_file + f"_{args.activation_suffix}.pt"

        if not args.force and os.path.exists(pkl_filename):
            skip = True
            try:
                with open(pkl_filename, "rb") as f:
                    existing = pickle.load(f)
                if len(existing) > 0 and len(existing[0]["logits_options"]) != args.num_options:
                    print(f"  Re-running {file} — existing has {len(existing[0]['logits_options'])} logits, need {args.num_options}")
                    skip = False
            except Exception:
                skip = False
            if skip and (not args.extract_activations or os.path.exists(act_filename)):
                print(f"  Skipping {file} — output already exists with {args.num_options} logits.")
                continue
        files_to_run.append(file)

    if not files_to_run:
        print("Nothing to do — all outputs exist.")
        return

    tokenizer, model = load_model(args)

    validated = False
    for file in files_to_run:
        # Compute output paths
        save_file = args.model.split("/")[-1] + "_" + file.split(".json")[0] + "_" + args.prompt_method
        save_file += "_icl" + str(args.few_shot)
        if args.cot:
            save_file += "_cot"
        save_file = os.path.join(args.output_dir, save_file)

        pkl_filename = save_file + (f"_{args.pkl_suffix}.pkl" if args.pkl_suffix else ".pkl")
        act_filename = save_file + f"_{args.activation_suffix}.pt"

        data = load_data(os.path.join(args.data_path, file))
        # get few-shot examples
        if args.few_shot > 0:
            fewshot_exps = get_fewshot_exps(data)
        else:
            fewshot_exps = None
        prompt_data = []
        for datum in data:
            if args.prompt_method == "base":
                prompt_data.append(format_base_prompt(datum, args, tokenizer, fewshot_exps=fewshot_exps))
            elif args.prompt_method == "shared":
                prompt_data.append(format_shared_prompt(datum, args, tokenizer, fewshot_exps=fewshot_exps))
            elif args.prompt_method == "task":
                prompt_data.append(format_task_prompt(datum, args, tokenizer, fewshot_exps=fewshot_exps))

        # Validate batching correctness on first non-skipped dataset
        if args.validate and args.batch_size > 1 and not validated:
            validate_batching(model, tokenizer, prompt_data, args)
            validated = True

        model_outputs = get_model_outputs(model, tokenizer, prompt_data, args)

        logits_only = [{"id": out["id"], "logits_options": out["logits_options"]}
                       for out in model_outputs]
        with open(pkl_filename, "wb") as f:
            pickle.dump(logits_only, f)

        if args.extract_activations:
            save_dict = {
                "outputs": model_outputs,
                "model_name": args.model,
                "file": file,
                "prompt_method": args.prompt_method,
                "few_shot": args.few_shot,
                "cot": args.cot
            }
            torch.save(save_dict, save_file + f"_{args.activation_suffix}.pt")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, required=True)
    parser.add_argument('--data_path', type=str, default="../data")
    parser.add_argument('--file', type=str, default="xxx.json", help="Specify which dataset to use")
    parser.add_argument('--prompt_method', type=str, default="base", help="Select from 'base', 'shared', 'task'")
    parser.add_argument('--output_dir', type=str, default='outputs')
    parser.add_argument('--few_shot', type=int, default=0)
    parser.add_argument('--cot', action="store_true", default=False)
    parser.add_argument('--extract_activations', action="store_true", default=False,
                        help="Extract MLP activations from final token")
    parser.add_argument('--activation_suffix', type=str, default='activations',
                        help="Suffix for activation file (default: 'activations' -> _activations.pt)")
    parser.add_argument('--pkl_suffix', type=str, default='',
                        help="Suffix for pkl file (default: '' -> .pkl, 'residual' -> _residual.pkl)")
    parser.add_argument('--num_options', type=int, default=4,
                        help="Number of answer options to use (default: 4 for A,B,C,D)")
    parser.add_argument('--batch_size', type=int, default=1,
                        help="Batch size for inference (uses left-padding)")
    parser.add_argument('--validate', action="store_true", default=False,
                        help="Validate batching by comparing batch=1 vs batch=N on first 10 samples")
    parser.add_argument('--force', action="store_true", default=False,
                        help="Force re-run even if output files already exist")
    args = parser.parse_args()
    
    main(args)