quantization_metric/alphalora/evaluation.py

import os
import re
import json
import argparse
import random
from tqdm import tqdm
from datasets import load_from_disk
import torch
from src.mola_peft_model_hacked import PeftModel
from transformers import GenerationConfig, LlamaTokenizer, AutoConfig
import sys
from typing import Union
from src.mola_modeling_llama_hacked import LlamaForCausalLM_d
from transformers import AutoTokenizer
from src.mola_modeling_mistral_hacked import MistralForCausalLM_d
from src.mistralconfig import MistralConfig
from src.gemma_config import GemmaConfig
from src.mola_modeling_gemma import GemmaForCausalLM_d

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

try:
    if torch.backends.mps.is_available():
        device = "mps"
except:  # noqa: E722
    pass

seed = 10
random.seed(seed)
torch.manual_seed(0)


class Prompter(object):
    __slots__ = ("template", "_verbose")

    def __init__(self, template_name: str = "", verbose: bool = False):
        self._verbose = verbose
        self.template = {
            "description": "Template used by Alpaca-LoRA.",
            "prompt_input": "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
            "prompt_no_input": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n",
            "response_split": "### Response:"
        }

        if self._verbose:
            print(
                f"Using prompt template {template_name}: {self.template['description']}"
            )

    def generate_prompt(
            self,
            instruction: str,
            input: Union[None, str] = None,
            label: Union[None, str] = None,
    ) -> str:
        # returns the full prompt from instruction and optional input
        # if a label (=response, =output) is provided, it's also appended.
        if input:
            res = self.template["prompt_input"].format(
                instruction=instruction, input=input
            )
        else:
            res = self.template["prompt_no_input"].format(
                instruction=instruction
            )
        if label:
            res = f"{res}{label}"
        if self._verbose:
            print(res)
        return res

    def get_response(self, output: str) -> str:
        return output.split(self.template["response_split"])[1].strip()


def main():
    parser = argparse.ArgumentParser(description='Evaluation')
    # Defining arguments
    parser.add_argument('--test_dataset', type=str, default="./scienceqa/scienceq_test.json",
                        help='test_dataset')
    parser.add_argument('--base_model', type=str, default="NousResearch/Llama-2-7b-hf", help='base_model')
    parser.add_argument('--mola_weights', type=str, default="./scienceqa_mola",
                        help='mola_model')
    parser.add_argument('--number_experts', type=str,
                        default="2,2,2,2,2,2,2,2,4,4,4,4,4,4,4,4,6,6,6,6,6,6,6,6,8,8,8,8,8,8,8,8",
                        help='experts number')
    parser.add_argument('--top_k', type=str,
                        default="2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2",
                        help='lora_model')
    parser.add_argument('--save_path', type=str,
                        default="./results/mola_test_sciqa_seed_10.json",
                        help='lora_model')
    parser.add_argument('--lora_target_modules', type=str,
                        default="q_proj,v_proj,k_proj,o_proj,gate_proj,down_proj,up_proj", help='lora_target_modules')
    parser.add_argument('--batch_size', type=int, default=8, help='base_model')

    # Parsing arguments
    args = parser.parse_args()
    if args.test_dataset.endswith(".json") or args.test_dataset.endswith(".jsonl"):
        data_a = json.load(open(args.test_dataset))
    else:
        data_aa = load_from_disk(args.test_dataset)["test"]
        data_a = [d for d in data_aa]

    base_model = args.base_model
    mola_weights = args.mola_weights
    max_batch_size = args.batch_size

    lora_target_modules = args.lora_target_modules.split(",")
    lora_target_modules = [str(lr) for lr in lora_target_modules]
    print(lora_target_modules)
    number_experts = args.number_experts.split(",")
    number_experts = [int(lr) for lr in number_experts]
    print(number_experts)
    top_k = args.top_k.split(",")
    top_k = [int(lr) for lr in top_k]
    print(top_k)

    print(args.test_dataset)
    print(args.base_model)
    print(args.mola_weights)

    load_8bit = False

    if "Llama" in base_model or "llama" in base_model:
        tokenizer = LlamaTokenizer.from_pretrained(base_model, padding_side='left')
        config = AutoConfig.from_pretrained(base_model)

        config.lora_target_modules = lora_target_modules
        if device == "cuda":
            model = LlamaForCausalLM_d.from_pretrained(
            base_model,
            config=config,
            load_in_8bit=load_8bit,
            torch_dtype=torch.float16,
            device_map="auto",
            )
            model = PeftModel.from_pretrained(
            model,
            mola_weights,
            torch_dtype=torch.float16,
            number_experts=number_experts,
            top_k=top_k,
            )
        else:
            model = LlamaForCausalLM_d.from_pretrained(
                base_model, config=config, device_map={"": device}, low_cpu_mem_usage=True
            )
            model = PeftModel.from_pretrained(
                model,
                mola_weights,
                device_map={"": device},
            )
    elif "Mistral" in base_model:
        tokenizer = AutoTokenizer.from_pretrained(base_model,padding_side='left')
        config = MistralConfig.from_pretrained(base_model)
        config.lora_target_modules = lora_target_modules
        model = MistralForCausalLM_d.from_pretrained(
            base_model,
            config=config,
            load_in_8bit=False,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        model = PeftModel.from_pretrained(
            model,
            mola_weights,
            torch_dtype=torch.float16,
            number_experts=number_experts,
            top_k=top_k,
        )
    else:
        print("Not support")


    obalance = False
    model.get_new_parameters(number_experts, top_k, obalance)

    print(model.config.pad_token_id, tokenizer.pad_token_id)
    print(model.config.bos_token_id, tokenizer.bos_token_id)
    print(model.config.eos_token_id, tokenizer.eos_token_id)
    # unwind broken decapoda-research config
    model.config.pad_token_id = tokenizer.pad_token_id = 0  # unk
    model.config.bos_token_id = 1
    model.config.eos_token_id = 2

    if not load_8bit:
        model.half()  # seems to fix bugs for some users.

    model.eval()
    if torch.__version__ >= "2" and sys.platform != "win32":
        model = torch.compile(model)

    prompter = Prompter(template_name="alpaca")

    max_new_tokens = 128
    save_every = 200

    correct = 0
    results = []
    outputs = []
    gt = []

    for start_idx in tqdm(range(0, len(data_a), max_batch_size)):
        end_idx = min(start_idx + max_batch_size, len(data_a))
        batch = data_a[start_idx:end_idx]

        answers = [str(example["answer"]) for example in batch]

        # generate prompt
        prompts = [prompter.generate_prompt(example['instruction'], example['input']) for example in batch]
        inputs = tokenizer(prompts, padding=True, return_tensors="pt")
        input_ids = inputs["input_ids"].to(device)

        with torch.no_grad():
            generation_output = model.generate(
                input_ids=input_ids,
                return_dict_in_generate=True,
                output_scores=True,
                max_new_tokens=max_new_tokens,
            )
        s = generation_output.sequences
        output = tokenizer.batch_decode(s)
        output = [prompter.get_response(otp) for otp in output]
        # extract the answer
        print(output)

        # TODO: Here we use different pattern to extract the answer for different datasets
        # pattern = re.compile(r'The anwser to the question is (\d+):*')
        if "cola" in args.test_dataset or "mrpc" in args.test_dataset or "rte" in args.test_dataset:
            pattern = re.compile(r'Answer: ([\w\s]+)')
            res = [pattern.search(otp).group(1) for otp in output]

        else:
            pattern = re.compile(r'The answer is ([A-Z]).')
            res = [pattern.findall(otp) for otp in output]
        #-------------------------------------------------------------------------------------------------------
        print(res)
        pred = []
        for r_i in range(len(res)):

            #TODO: for diiferent datasets, we need to use different pattern to match the answer
            if "rte" in args.test_dataset:
                if res[r_i] == "entailment" or res[r_i] == "not entailment":
                    answer = res[r_i]  # 'A', 'B', ...
                else:
                    answer = "FAILED"
                    print(res[r_i])
            elif "mrpc" in args.test_dataset:
                if res[r_i] == "equivalent" or res[r_i] == "not equivalent":
                    answer = res[r_i]  # 'A', 'B', ...
                else:
                    answer = "FAILED"
                    print(res[r_i])
            elif "cola" in args.test_dataset:
                if res[r_i] == "acceptable" or res[r_i] == "unacceptable":
                    answer = res[r_i]  # 'A', 'B', ...
                else:
                    answer = "FAILED"
                    print(res[r_i])
            else:
                if len(res[r_i]) == 1:
                    answer = res[r_i][0]
        #-------------------------------------------------------------------------------------------------------
            pred.append(answer)
            results.append(res[r_i])
            outputs.append(output[r_i])
            gt.append(answers[r_i])

            if str(answer) == str(answers[r_i]):
                correct += 1
                print('correct:', str(answer), str(answers[r_i]))
            else:
                print('gt-ans:', str(answer), str(answers[r_i]))

        acc = correct / len(results) * 100

        if end_idx % save_every == 0 or end_idx == len(data_a):
            result_file = args.save_path
            os.makedirs(result_file, exist_ok=True)
            print(f"{len(results)}/{len(data_a)}, correct: {correct}, acc: {round(acc, 2)}%, saving to {result_file}")
            data = {}
            data['acc'] = acc
            data['correct'] = correct
            data['len'] = len(results)
            data['results'] = results
            data['outputs'] = outputs
            with open(result_file, 'w') as f:
                json.dump(data, f, indent=2, separators=(',', ': '))


if __name__ == "__main__":
    main()