entity_description.py

# Modified from: https://github.com/edenbiran/HoppingTooLate
# Original Authors: Eden Biran, Daniela Gottesman, Sohee Yang, Mor Geva, Amir Globerson

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import os
import re
from datasets import load_dataset
from patched_generation import get_hidden_states, generate_with_patching_all_layers
from utils import get_layer_names, load_model, load_tokenizer


def generate_entity_description_with_patching(entries, model, tokenizer, source, do_sample):
    
    if source == "e2":
        entities = [subanswer_1[0] for subanswer_1 in entries["subanswer_1"]]
    elif source == "link":
        entities = [subanswer_1[0] + " as" for subanswer_1 in entries["subanswer_1"]]
    elif source == "e3":
        entities = [re.search(f"(.+?) is to {re.escape(subanswer_1[0])} as (.+?) is to", main_query).group(2) for subanswer_1, main_query in zip(entries["subanswer_1"], entries["main_query"])]
    elif source == "last":
        entities = [prompt.strip().split()[-1] for prompt in entries["main_query"]]
    else:
        raise ValueError(f"Source {source} not supported")

    tokenizer.padding_side = "right"
    hidden_states = get_hidden_states(model, tokenizer, entities, entries["main_query"])
    tokenizer.padding_side = "left"

    if args.target_prompt_type == "relational_description_e2_full" and source == "e2":
        entries["target_prompt"] = [target_prompt.format(re.search(f"(.+?) is to {re.escape(subanswer_1[0])} as.+", main_query).group(1)) for target_prompt, subanswer_1, main_query in zip(entries["target_prompt"], entries["subanswer_1"], entries["main_query"])]
    elif args.target_prompt_type == "relational_description_e3_full" and source == "e3":
        entries["target_prompt"] = [target_prompt.format(subanswer_2[0]) for target_prompt, subanswer_2 in zip(entries["target_prompt"], entries["subanswer_2"])]
    elif args.target_prompt_type == "relational_description_last_inserte1" and source == "last":
        entries["target_prompt"] = [target_prompt.format(re.search(f"(.+?) is to {re.escape(subanswer_1[0])} as (.+?) is to", main_query).group(1), subanswer_2[0]) for target_prompt, subanswer_1, subanswer_2, main_query in zip(entries["target_prompt"], entries["subanswer_1"], entries["subanswer_2"], entries["main_query"])]
    elif args.target_prompt_type == "relational_description_last_inserte3" and source == "last":
        entries["target_prompt"] = [target_prompt.format(re.search(f"(.+?) is to {re.escape(subanswer_1[0])} as (.+?) is to", main_query).group(2), subanswer_2[0]) for target_prompt, subanswer_1, subanswer_2, main_query in zip(entries["target_prompt"], entries["subanswer_1"], entries["subanswer_2"], entries["main_query"])]

    generations = generate_with_patching_all_layers(model, tokenizer, hidden_states, entries["target_prompt"], "x", do_sample)

    entry_count = len(entries["main_query"])
    layer_count = len(get_layer_names(model))
    new_entries = {}
    for k in entries.keys():
        new_entries[k] = []
    new_entries["source_layer"] = []
    new_entries["target_layer"] = []
    new_entries["generation"] = []
    for source_layer in range(layer_count):
        for target_layer in range(layer_count):
            for i in range(entry_count):
                for k, v in entries.items():
                    new_entries[k].append(v[i])
                new_entries["source_layer"].append(source_layer)
                new_entries["target_layer"].append(target_layer)
                new_entries["generation"].append(generations[source_layer, target_layer][i])

    return new_entries

def filter_samples(sample):
    if args.mode == "correct":
        return sample["subquery_1_iscorrect"] == 1 and sample["subquery_2_iscorrect"] == 1 and sample["main_query_iscorrect"] == 1 and sample["exclude_e1_e2_iscorrect"] == 0 and sample["exclude_e2_iscorrect"] == 0
    elif args.mode == "incorrect":
        return sample["subquery_1_iscorrect"] == 1 and sample["subquery_2_iscorrect"] == 1 and sample["main_query_iscorrect"] == 0 and sample["exclude_e1_e2_iscorrect"] == 0 and sample["exclude_e2_iscorrect"] == 0
    else:
        raise ValueError(f"Mode {args.mode} not supported")
		
def add_target_prompt(sample):
    if args.target_prompt_type == "default_description": # for section 4
        target_prompt = (
            "Syria: Syria is a country in the Middle East, " +
            "Leonardo DiCaprio: Leonardo DiCaprio is an American actor, " +
            "Samsung: Samsung is a South Korean multinational corporation, " +
            "x"
        )
    elif args.target_prompt_type == "relational_description_e2_full": # for section 5
        target_prompt = (
            "Japan is to Tokyo: capital of, " +
            "Theory of Evolution is to Charles Darwin: founder of, " +
            "Peace is to olive branch: symbol of, " +
            "{} is to x"
        )
    elif args.target_prompt_type == "relational_description_e3_full": # for section 5
        target_prompt = (
            "Japan is to Tokyo: capital of, " +
            "Theory of Evolution is to Charles Darwin: founder of, " +
            "Peace is to olive branch: symbol of, " +
            "x is to {}"
        )
    elif args.target_prompt_type == "relational_description_last_inserte1": # for section 5
        target_prompt = (
            "Japan is to Tokyo: capital of, " +
            "Theory of Evolution is to Charles Darwin: founder of, " +
            "Peace is to olive branch: symbol of, " +
            "{} is x"
        )
    elif args.target_prompt_type == "relational_description_last_inserte3": # for section 5 (same as above, but placeholder replaced with e3)
        target_prompt = (
            "Japan is to Tokyo: capital of, " +
            "Theory of Evolution is to Charles Darwin: founder of, " +
            "Peace is to olive branch: symbol of, " +
            "{} is x"
        )

    sample["target_prompt"] = target_prompt
    return sample

###############

def main(args):
    print(args)
		
    input_path = os.path.join(os.path.dirname(__file__), "results", "evaluate_knowledge", args.model_name, "evaluate_knowledge.jsonl")
    output_path = os.path.join(os.path.dirname(__file__), "results", "patchscopes", args.mode, args.model_name, f"entity_description_{args.target_prompt_type}_{args.source}_max20.csv")

    if not os.path.exists(input_path):
        raise FileNotFoundError
    
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
	
    dataset = load_dataset("json", data_files=input_path)
    dataset = dataset.filter(filter_samples)["train"]
    dataset = dataset.map(add_target_prompt)
	
    if len(dataset) > 500:
        dataset = dataset.shuffle(seed=42).select(range(500))
		
    model = load_model(args.model_name)
    model.eval()
    tokenizer = load_tokenizer(args.model_name)
    generations = dataset.map(
        generate_entity_description_with_patching,
        fn_kwargs={"model": model, "tokenizer": tokenizer, "source": args.source,
                   "do_sample": args.do_sample},
        batched=True,
        batch_size=args.batch_size,
        remove_columns=dataset.column_names
    )

    generations.to_csv(output_path, escapechar='\\')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", type=str, choices=[
        "meta-llama/Llama-2-13b-hf", "google/gemma-7b", "Qwen/Qwen2.5-14B"])
    parser.add_argument("--mode", choices=["correct", "incorrect"])
    parser.add_argument("--source", choices=["e2", "link", "e3", "last"])
    parser.add_argument("--target_prompt_type", choices=["default_description", "relational_description_e2_full", "relational_description_e3_full", "relational_description_last_inserte1","relational_description_last_inserte3"])	
    parser.add_argument("--batch_size", type=int, default=32)
    args = parser.parse_args()
    main(args)