experiment.py

import torch
import StaICC
import util.load_model_and_data as load_model_and_data
from transformers import AutoTokenizer, AutoModelForCausalLM
import numpy as np
import util.default_config as default_config
import random
from StaICC.prefabricate_inference import model_kernel as model_kernel
from functools import partial
import argparse
from util import model_prepare as model_prepare
from util.head_tuning import head_tuner
import pickle
from datetime import datetime
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training 
import time
from util import baseline
from util.functional import generate_text_from_prefix_set
import os

## Parse arguments
parser = argparse.ArgumentParser(description="ICL Head Tuning Prototype")

parser.add_argument("--ICL_model_name", type=str, help="Name of the ICL model")
parser.add_argument("--ICL_dataset_index", type=int, help="ICL dataset index")
parser.add_argument("--method", type=str, help="Method of inference", default="ABFT")
parser.add_argument("--ICL_k", type=int, default=default_config.ICL_k, help="ICL k value")
parser.add_argument("--quantized", default=default_config.quantized, help="Whether the model is quantized", dest="quantized", action="store_true")
parser.add_argument("--induction_head_attention_score_threshold", type=float, default=default_config.induction_head_attention_score_threshold, help="Induction head attention score threshold")
parser.add_argument("--correct_label_award", type=float, default=default_config.correct_label_award, help="Correct label award")
parser.add_argument("--wrong_label_penalty", type=float, default=default_config.wrong_label_penalty, help="Wrong label penalty")
parser.add_argument("--train_sample_num", type=int, default=default_config.train_sample_num, help="Number of training samples")
parser.add_argument("--lr", type=float, default=default_config.lr, help="Learning rate")
parser.add_argument("--pseudo_batch_size", type=int, default=default_config.pseudo_batch_size, help="Pseudo batch size")
parser.add_argument("--epoch", type=int, default=default_config.epoch, help="Number of epochs")
parser.add_argument("--random_seed", type=int, default=default_config.random_seed, help="Random seed")
parser.add_argument("--no_fore_test", default=default_config.no_fore_test, help="The model will not be tested for accuracy before training for a baseline", dest="no_fore_test", action="store_true")
parser.add_argument("--no_post_test", default=default_config.no_post_test, help="The model will not be tested for accuracy after training for a result report", dest="no_post_test", action="store_true")
parser.add_argument("--test_type", type=str, default=default_config.test_type, help="Normal: normal test for accuracy; TempSensitivity: test for template sensitivity; DemoSensitivity: test for demo sensitivity")
parser.add_argument("--generalization_test", default=default_config.generalization_test, help="The model will be tested on all the 8 datasets in the post test", dest="generalization_test", action="store_true")
parser.add_argument("--dont_save_model", default=default_config.dont_save_model, help="The model will not be saved", dest="dont_save_model", action="store_true")
parser.add_argument("--restricted_gradient", default=default_config.restricted_gradient, help="Whether to restrict the gradient to the induction head", dest="restricted_gradient", action="store_true")
parser.add_argument("--induction_head_method", type=str, default=default_config.induction_head_method, help="Induction head judgement method")
parser.add_argument("--self_adaption_loss", default=default_config.self_adaption_loss, type=str, help="Whether to use the loss factor adaption, and the algorithm. P, PID", dest="self_adaption_loss")
parser.add_argument("--loss_adaption_factor_for_P", default=default_config.loss_adaption_factor, type=float, help="Loss adaption factor", dest="loss_adaption_factor")
parser.add_argument("--PID_param_P", default=default_config.PID_param_P, type=float, help="P factor in PID", dest="PID_param_P")
parser.add_argument("--PID_param_I", default=default_config.PID_param_I, type=float, help="I factor in PID", dest="PID_param_I")
parser.add_argument("--PID_param_D", default=default_config.PID_param_D, type=float, help="D factor in PID", dest="PID_param_D")
parser.add_argument("--lora_r", default=default_config.lora_r, type=int, help="Lora rank", dest="lora_r")
parser.add_argument("--MAUVE_test", default=default_config.MAUVE_test, help="Whether to test with MAUVE. Add with 2 testing phase before and after fine-tuning. VERY SLOW.", dest="MAUVE_test", action="store_true")
parser.add_argument("--MAUVE_prefix", default=default_config.MAUVE_test_prefix, type=str, help="The prefix file for the generation of the MAUVE test")

# TODO: Lora specific arguments


args = parser.parse_args()

ICL_model_name = args.ICL_model_name
ICL_k = args.ICL_k
method = args.method
ICL_dataset_index = args.ICL_dataset_index
quantized = args.quantized
induction_head_attention_score_threshold = args.induction_head_attention_score_threshold
correct_label_award = args.correct_label_award
wrong_label_penalty = args.wrong_label_penalty
train_sample_num = args.train_sample_num
lr = args.lr
pseudo_batch_size = args.pseudo_batch_size
epoch = args.epoch
random_seed = args.random_seed
no_fore_test = args.no_fore_test
no_post_test = args.no_post_test
test_type = args.test_type
generalization_test = args.generalization_test
dont_save_model = args.dont_save_model
restricted_gradient = args.restricted_gradient
induction_head_method = args.induction_head_method
self_adaption_loss = args.self_adaption_loss
loss_adaption_factor = args.loss_adaption_factor
PID_param_P = args.PID_param_P
PID_param_I = args.PID_param_I
PID_param_D = args.PID_param_D
lora_r = args.lora_r
MAUVE_test = args.MAUVE_test
MAUVE_prefix = args.MAUVE_prefix

huggingface_token = default_config.huggingface_token

## Solve model type
if method == "ABFT" or method == "E2E":
    if "llama-3" in ICL_model_name.lower() or "llama3" in ICL_model_name.lower() or "s1" in ICL_model_name.lower():
        model_type = "llama-3"
    elif "gpt2" in ICL_model_name.lower():
        model_type = "gpt2"
    elif "qwen" in ICL_model_name.lower():
        model_type = "qwen"
    elif "falcon3" in ICL_model_name.lower():
        model_type = "falcon3"
    else:
        raise ValueError("Model type not supported. Please write a new model_prepare function for the new model type.")

if method == "ABFT" or method == "E2E":
    if model_type == "llama-3":
        model_set = model_prepare.setup_model_for_Llama3
        locat = model_prepare.attention_parameter_locating_function_for_Llama3
    elif model_type == "gpt2":
        model_set = model_prepare.setup_model_for_GPT2
        locat = model_prepare.attention_parameter_locating_function_for_GPT2
    elif model_type == "qwen":
        model_set = model_prepare.setup_model_for_qwen
        locat = model_prepare.attention_parameter_locating_function_for_qwen
    elif model_type == "falcon3":
        model_set = model_prepare.setup_model_for_falcon3
        locat = model_prepare.attention_parameter_locating_function_for_falcon3
    else:
        raise ValueError("Model type not supported. Please write a new model_prepare function for the new model type.")

## Load model and data
if test_type == "Normal":
    benchmark = StaICC.Normal(ICL_k)
    experimentor = benchmark[ICL_dataset_index]
    experimentor.prompt_former.replace_space_to_label()
elif test_type == "TempSensitivity":
    benchmark = StaICC.Template_sens(ICL_k)
    experimentor = benchmark[ICL_dataset_index]
    experimentor.prompt_former.replace_space_to_label()
elif test_type == "DemoSensitivity":
    benchmark = StaICC.Demo_sens(ICL_k)
    experimentor = benchmark[ICL_dataset_index]
    experimentor.prompt_former.replace_space_to_label()

if method != "metaicl":
    ICL_model, ICL_tknz = load_model_and_data.load_ICL_model(ICL_model_name, huggingface_token = huggingface_token, quantized = quantized)
    for param in ICL_model.parameters():
        param.requires_grad = False
else:
    ICL_tknz = AutoTokenizer.from_pretrained("openai-community/gpt2-large", trust_remote_code=True)
    dict = torch.load("baseline_models/MetaICL/checkpoints/metaicl/hr_to_lr/model.pt")
    ICL_model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2-large", state_dict=dict)


## MAUVE pre-fine-tuning generation (if any)
if MAUVE_test:
    with open(MAUVE_prefix, 'rb') as file:
        MAUVE_prefix_data = pickle.load(file)
    cache_path = f"temp/MAUVE_cache_{ICL_model_name.replace('/', '_')}.pkl"
    if os.path.exists(cache_path):
        with open(cache_path, 'rb') as file:
            generated_fore_training_cache = pickle.load(file)
            if generated_fore_training_cache[0] == MAUVE_prefix_data:
                generated_fore_training = generated_fore_training_cache[1]
            else:
                generated_fore_training = generate_text_from_prefix_set(ICL_model, ICL_tknz, MAUVE_prefix_data)
                with open(cache_path, 'wb') as file:
                    pickle.dump((MAUVE_prefix_data, generated_fore_training), file)
    else:
        generated_fore_training = generate_text_from_prefix_set(ICL_model, ICL_tknz, MAUVE_prefix_data)
        with open(cache_path, 'wb') as file:
            pickle.dump((MAUVE_prefix_data, generated_fore_training), file)

## Fore test
inference = partial(model_kernel.standard_ICL_inference_with_torch_Causal_LM, model = ICL_model, tokenizer = ICL_tknz, label_space = experimentor.get_label_space())
if not no_fore_test:
    print("Fore test:")
    fore_test_res = experimentor(inference)
    print(fore_test_res)


## Training
if method == "ABFT":
    if not quantized:
        model_set(ICL_model) # Enable the gradient of the attention parameters
    else:
        config = LoraConfig(
            r=lora_r,
            lora_alpha=8,
            target_modules=model_prepare.get_target_modules(model_type),
        )
        ICL_model = get_peft_model(ICL_model, config)

    head_tuner_instance = head_tuner(
        ICL_model = ICL_model,
        ICL_tokenizer = ICL_tknz,
        induction_head_attention_score_threshold = induction_head_attention_score_threshold,
        correct_label_award = correct_label_award,
        wrong_label_penalty = wrong_label_penalty,
        StaICC_experimentor = experimentor,
        training_parameters = {
            "train_sample_num" : train_sample_num,
            "lr" : lr,
            "optimizer" : torch.optim.Adam,
            "pseudo_batch_size" : pseudo_batch_size,
            "epoch" : epoch,
        },
        random_seed = random_seed,
        restricted_gradient = restricted_gradient,
        attention_parameter_locating_function = locat,
        induction_head_judge = induction_head_method,
        self_adaption_loss = self_adaption_loss,
        self_adaption_loss_decay = loss_adaption_factor,
        PID_param_P = PID_param_P,
        PID_param_I = PID_param_I,
        PID_param_D = PID_param_D
    )
    start_time = time.time()
    train_loss, avg_heads = head_tuner_instance.fine_tune()
    end_time = time.time()
elif method == "E2E":
    if not quantized:
        model_set(ICL_model, method = "E2E") # Enable the gradient of the attention parameters
    else:
        prepare_model_for_kbit_training(ICL_model)
        config = LoraConfig(
            r=lora_r,
            lora_alpha=8,
            target_modules=model_prepare.get_target_modules(model_type),
        )
        ICL_model = get_peft_model(ICL_model, config)
    e2e_tuner_instance = baseline.end_to_end_tuner(
        ICL_model = ICL_model,
        ICL_tokenizer = ICL_tknz,
        StaICC_experimentor = experimentor,
        training_parameters = {
            "train_sample_num" : train_sample_num,
            "lr" : lr,
            "optimizer" : torch.optim.Adam,
            "pseudo_batch_size" : pseudo_batch_size,
            "epoch" : epoch,
        },
        random_seed = 42
    )
    start_time = time.time()
    train_loss = e2e_tuner_instance.fine_tune()
    end_time = time.time()
elif method == "hidden calibration":
    start_time = time.time()
    inference = baseline.hidden_calibration(ICL_model, ICL_tknz, experimentor)
    end_time = time.time()
elif method == "contextual calibration":
    start_time = time.time()
    inference = baseline.contextual_calibration(ICL_model, ICL_tknz, experimentor)
    end_time = time.time()
elif method == "none" or method == "metaicl":
    no_post_test = False if no_fore_test else True
    start_time = time.time()
    end_time = time.time()
else:
    raise ValueError("Method not supported")

time_used = end_time - start_time
print(f"Training completed in {time_used} seconds")


## Post test
for param in ICL_model.parameters():
    param.requires_grad = False

if not no_post_test:
    print("Post test:")
    if test_type == "TempSensitivity":
        benchmark = StaICC.Template_sens(ICL_k)
        experimentor = benchmark[ICL_dataset_index]
        experimentor.prompt_former.replace_space_to_label()
        post_test_res = experimentor(inference)
    elif test_type == "DemoSensitivity":
        benchmark = StaICC.Demo_sens(ICL_k)
        experimentor = benchmark[ICL_dataset_index]
        experimentor.prompt_former.replace_space_to_label()
        post_test_res = experimentor(inference)
    else:
        if generalization_test:
            post_test_res = []
            for i in range(9):
                if i == 5:
                    continue
                temp_experimentor = benchmark[i]
                post_test_res.append(temp_experimentor(inference))
        else:
            post_test_res = experimentor(inference)
    print(post_test_res)

## MAUVE post-fine-tuning generation (if any)
if MAUVE_test:
    generated_post_training = generate_text_from_prefix_set(ICL_model, ICL_tknz, MAUVE_prefix_data)
    from evaluate import load
    mauve = load('mauve')
    mauve_results = mauve.compute(predictions=generated_post_training, references=generated_fore_training).mauve

## Build the save data
model_cpu = ICL_model.to('cpu')
if quantized:
    model_cpu = model_cpu.merge_and_unload()
    model_cpu = model_cpu.to('cpu')

experiment_params = {
    "ICL_model_name": ICL_model_name,
    "ICL_k": ICL_k,
    "method": method,
    "ICL_dataset_index": ICL_dataset_index,
    "quantized": quantized,
    "huggingface_token": huggingface_token,
    "induction_head_attention_score_threshold": induction_head_attention_score_threshold,
    "correct_label_award": correct_label_award,
    "wrong_label_penalty": wrong_label_penalty,
    "train_sample_num": train_sample_num,
    "lr": lr,
    "pseudo_batch_size": pseudo_batch_size,
    "epoch": epoch,
    "random_seed": random_seed,
    "no_fore_test": no_fore_test,
    "no_post_test": no_post_test,
    "test_type": test_type,
    "generalization_test": generalization_test,
    "MAUVE_test": MAUVE_test,
    "restricted_gradient": restricted_gradient,
    "induction_head_method": induction_head_method,
    "self_adaption_loss": self_adaption_loss,
    "loss_adaption_factor": loss_adaption_factor,
    "PID_param_P": PID_param_P,
    "PID_param_I": PID_param_I,
    "PID_param_D": PID_param_D
}

save_data = {
    "model": (model_cpu if not quantized else model_cpu.state_dict()) if not dont_save_model else None,
    "ICL_model_name": ICL_model_name,
    "post_test_res": post_test_res if 'post_test_res' in locals() and not no_post_test else None,
    "fore_test_res": fore_test_res if 'fore_test_res' in locals() and not no_fore_test else None,
    "experiment_params": experiment_params,
    "train_loss": train_loss if 'train_loss' in locals() else None,
    "avg_heads": avg_heads if 'avg_heads' in locals() else None,
    "correct_label_award_log": head_tuner_instance.correct_label_award_log if 'head_tuner_instance' in locals() else None,
    "wrong_label_penalty_log": head_tuner_instance.wrong_label_penalty_log if 'head_tuner_instance' in locals() else None,
    "time_used": time_used if 'time_used' in locals() else None,
    "MAUVE_results": mauve_results if 'mauve_results' in locals() else None
}

## Save
current_time = datetime.now().strftime("%Y%m%d_%H%M%S_%f")

if not os.path.exists('logs'):
    os.makedirs('logs')

with open(f'logs/trained_model_and_results_{current_time}.pkl', 'wb') as f:
    pickle.dump(save_data, f)

# Save experiment parameters to a txt file
params_filename = f"logs/experiment_params_{current_time}.txt"

with open(params_filename, 'w') as f:
    for key, value in experiment_params.items():
        f.write(f"{key}: {value}\n")

# Save test results to a txt file if they exist
if not no_fore_test:
    fore_test_filename = f"logs/fore_test_res_{current_time}.txt"
    with open(fore_test_filename, 'w') as f:
        f.write(str(fore_test_res))

if (not no_post_test) or (MAUVE_test):
    post_test_filename = f"logs/post_test_res_{current_time}.txt"
    with open(post_test_filename, 'w') as f:
        f.write(str(post_test_res) if 'post_test_res' in locals() else "No post test")
        f.write("\n")
        f.write(str(time_used))
        f.write("\n")
        f.write(str(mauve_results) if 'mauve_results' in locals() else "No MAUVE test")