eval_react.py

from unsloth import FastLanguageModel
# from transformers import TextStreamer
# from unsloth import add_new_tokens
from peft import AutoPeftModelForCausalLM, PeftModel
import os
import json
import re
from tqdm import tqdm
from collections import Counter
import random
import argparse
from enum import Enum

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)
import packaging.version
import torch
import transformers

random.seed(42)

parser = argparse.ArgumentParser(description='Evaluation Reasoner')
parser.add_argument('--beam', type=int, default=1, help='Beam size for evaluation')
parser.add_argument('--ckpt', type=str, default='', help='model ckpt path for evaluation')
parser.add_argument('--data', type=str, default='', help='test data path for evaluation')
parser.add_argument('--log_path', type=str, default='', help='logging file saving path')
parser.add_argument('--load_in_4bit', type=bool, default=False, help='Set true or false for 4bit quantization')
args = parser.parse_args()

beam = args.beam


# os.environ["CUDA_VISIBLE_DEVICES"] = "4"

max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
load_in_4bit = args.load_in_4bit # Use 4bit quantization to reduce memory usage. Can be False.
load_in_8bit = True

alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{}

### Input:
{}

### Response:
{}"""




# for qwen 2.5 14b instruct
chatML_template_test_qwen = """<|im_start|>user
{}<|im_end|>
<|im_start|>assistant
{}"""

# for llama 3.1 8b instruct
chatML_template_test_llama = """<|start_header_id|>user<|end_header_id|>
{}<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
{}"""

chatML_template_test = chatML_template_test_qwen if 'Qwen' in args.ckpt else chatML_template_test_llama


DEFAULT_CHATML_CHAT_TEMPLATE = "{% for message in messages %}\n{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% if loop.last and add_generation_prompt %}{{'<|im_start|>assistant\n' }}{% endif %}{% endfor %}"
DEFAULT_ZEPHYR_CHAT_TEMPLATE = "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"

class ZephyrSpecialTokens(str, Enum):
    user = "<|user|>"
    assistant = "<|assistant|>"
    system = "<|system|>"
    eos_token = "</s>"
    bos_token = "<s>"
    pad_token = "<pad>"

    @classmethod
    def list(cls):
        return [c.value for c in cls]


class ChatmlSpecialTokens(str, Enum):
    user = "<|im_start|>user"
    assistant = "<|im_start|>assistant"
    system = "<|im_start|>system"
    eos_token = "<|im_end|>"
    bos_token = "<s>"
    pad_token = "<pad>"

    @classmethod
    def list(cls):
        return [c.value for c in cls]

# model_name = "meta-llama/Llama-2-7b-hf"
# model_name = "meta-llama/Llama-2-13b-hf"
# model_name = "meta-llama/Llama-3.1-8B"
# model_name = "Qwen/Qwen2.5-14B-Instruct"
# model_path = f"ckpts/{model_name}"

# model_name = "meta-llama/Llama-3.2-3B-Instruct"

# model_path = "/local_data/xywen22/project/train_llama/outputs/checkpoint-10/"
# model_path = "/local_data/xywen22/project/train_llama/ckpts/meta-llama/Llama-2-7b-hf"
# model_path = "/local_data/xywen22/project/train_with_unsloth/ckpts/verification/StrategyQA_20250121/checkpoint-315"
# model_path = "ckpts/meta-llama/Llama-2-7b-hf/stratgeqa/step_type=memory-2-4-efficient=lora+prompt-tuning-lr=0.0002-soft-prompt=True-mcts-dependency-2/checkpoint-1488"
model_path = args.ckpt
# model_path = model_name

bnb_config = BitsAndBytesConfig(
    load_in_4bit=False,
    load_in_8bit=False,
    # bnb_4bit_quant_type=args.bnb_4bit_quant_type,
    # bnb_4bit_compute_dtype=compute_dtype,
    # bnb_4bit_use_double_quant=args.use_nested_quant,
    # bnb_4bit_quant_storage=quant_storage_dtype,
)

# base_model = AutoModelForCausalLM.from_pretrained(model_name)

# base_model.resize_token_embeddings(151688)

# model = PeftModel.from_pretrained(base_model, model_path)

# model = AutoPeftModelForCausalLM.from_pretrained(
#     model_path,
#     quantization_config=bnb_config,
#     trust_remote_code=True,
#     attn_implementation="flash_attention_2",
#     # torch_dtype=torch_dtype,
#     use_cache=False,
# )

chat_template_format = "chatml"

special_tokens = None
chat_template = None
if chat_template_format == "chatml":
    special_tokens = ChatmlSpecialTokens
    chat_template = DEFAULT_CHATML_CHAT_TEMPLATE
elif chat_template_format == "zephyr":
    special_tokens = ZephyrSpecialTokens
    chat_template = DEFAULT_ZEPHYR_CHAT_TEMPLATE

# if special_tokens is not None:
#     tokenizer = AutoTokenizer.from_pretrained(
#         # args.model_name_or_path,
#         model_path,
#         pad_token=special_tokens.pad_token.value,
#         bos_token=special_tokens.bos_token.value,
#         eos_token=special_tokens.eos_token.value,
#         additional_special_tokens=special_tokens.list(),
#         trust_remote_code=True,
#     )
#     tokenizer.chat_template = chat_template

#     # make embedding resizing configurable?
#     # Transformers 4.46.0+ defaults uses mean_resizing by default, which fails with QLoRA + FSDP because the
#     # embedding could be on meta device, therefore, we set mean_resizing=False in that case (i.e. the status quo
#     # ante). See https://github.com/huggingface/accelerate/issues/1620.
#     uses_transformers_4_46 = packaging.version.parse(transformers.__version__) >= packaging.version.parse("4.46.0")
#     uses_fsdp = os.environ.get("ACCELERATE_USE_FSDP").lower() == "true"
#     # if (bnb_config is not None) and uses_fsdp and uses_transformers_4_46:
#     #     model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=8, mean_resizing=False)
#     # else:
#     #     model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=8)
# else:
#     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=True)
#     tokenizer.pad_token = tokenizer.eos_token


# model = AutoPeftModelForCausalLM.from_pretrained(
#     model_path, # YOUR MODEL YOU USED FOR TRAINING
#     load_in_4bit = load_in_4bit,
# ).to("cuda")
# tokenizer = AutoTokenizer.from_pretrained(model_path)


model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = model_path, # YOUR MODEL YOU USED FOR TRAINING
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
    # resize_model_vocab=128264 # for llama3.1 8b
    resize_model_vocab=151672 if 'Qwen' in model_path else 128264
)


# # add_new_tokens(model, tokenizer, new_tokens = ["<GSM8K>", "<TruthfulQA>", "<MATH>", "<BBH>"])

FastLanguageModel.for_inference(model) # Enable native 2x faster inference

# a = tokenizer.encode("<|im_start|>assistant")
# print(model.get_input_embeddings().weight[a])
# print(a)
# b = tokenizer.encode("<|im_start|>user")
# print(model.get_input_embeddings().weight[b])
# print(b)

# c = tokenizer.encode("<knowledge_0>")
# print(model.get_input_embeddings().weight[c])
# print(c)

# d = tokenizer.encode("<reason_7>")
# print(model.get_input_embeddings().weight[d])
# print(d)

# e = tokenizer.encode("reason")
# print(model.get_input_embeddings().weight[e])
# print(e)

# data_path = "dataset/verification/StrategyQA_verification_test.json"
# data_path = "dataset/StrategyQA/StrategyQA_mcts_test.json"
data_path = args.data

def extract_answer(output):
    # 使用正则表达式提取JSON内容
    json_content = re.search(r'```json(.*?)```', output, re.DOTALL)
    if json_content:
        extracted_json = json_content.group(1).strip()
        try:
            extracted_json = json.loads(extracted_json)
        except json.JSONDecodeError as e:
            print(f"JSON decoding error: {e}")
            extracted_json = None
        # print("Extracted JSON content:", extracted_json)
    else:
        print("No JSON content found.")
    
    answer = extracted_json['answer']

    return answer
    


def load_data(data_path):
    data_path = data_path
    with open(data_path, 'r', encoding='utf-8') as file:
        data_list = json.load(file)
    # data_list = random.sample(data_list, 200)

    all_questions = []
    all_answers = []
    random_questions = []
    for i in range(len(data_list)):
        data = data_list[i]
        random_questions.append({
            'question': data['question'],
            'answer': data['answer']
        })

        all_questions.append(data['question'])
        all_answers.append(data['answer'])

    print(len(random_questions))
    print(random_questions[:3])

    return random_questions

random_questions = load_data(data_path)

counter = []
generated_results = []
# save_generate_results_path = f"eval_results/generations_beam{beam}_1488.json"
save_generate_results_path = args.log_path
for i in tqdm(range(len(random_questions)), total=len(random_questions)):
    data = random_questions[i]
    # output = data['output'].replace('\n', '')
    # print(f"input: {data['input']}")

    # print(extracted_json)
    # gt_answer = extract_answer(output)
    gt_answer = data['answer']
    question = data['question']
    # prompt = 'Question:' + ' ' + question
    prompt = question

    inputs = tokenizer(
    [
        chatML_template_test.format(
            # data['instruction'],
            # StrategyQA:
            # 'Answer the following question True or False step by step using the supporting facts in your knowledge.',
            # GPQA:
            # 'Answer the following single-choice question step by step using the supporting facts in your knowledge.',
            # f"Answer the following single-choice question step by step using the supporting facts in your knowledge.[[\nYour input is:\n{prompt}]]",
            # f"Answer the following single-choice question step by step using the supporting facts in your knowledge.\n{prompt}",
            """Solve a multiple choice question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: 
(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage.
(3) Finish[answer], which returns the answer (A, B, C, D, etc.) and finishes the task.""".strip() + '\n' + prompt,
            "", # output - leave this blank for generation!
        )
    ], return_tensors = "pt")

    # Ensure all inputs are on the same device
    inputs = {key: value.to("cuda") for key, value in inputs.items()}

    # text_streamer = TextStreamer(tokenizer)
    # outputs = model.generate(**inputs, streamer = text_streamer, num_beams = 2, max_new_tokens = 512)
    generated_ids = model.generate(**inputs, num_beams = beam, num_return_sequences=beam, max_new_tokens = 2048, do_sample=False)
    # generated_ids = model.generate(**inputs, num_beams = 1, num_return_sequences=1, max_new_tokens = 512)
    # print(outputs)
    generated_texts_beam = []
    for k in range(len(generated_ids)):
        # print(generated_ids[k])
        decoded_text = tokenizer.decode(generated_ids[k], skip_special_tokens=True)
        # print(decoded_text)
        generated_texts_beam.append(decoded_text)

    answer_predicts = []
    generated_texts = []
    for k in range(len(generated_texts_beam)):
        # print("*" * 100)
        # print(generated_texts_beam[i])
        # answer = extract_answer(generated_texts_beam[k].split('The answer is: ')[1].strip())
        generated_texts.append(generated_texts_beam[k])
        if 'The answer is: ' in generated_texts_beam[k]:
            answer = generated_texts_beam[k].split('The answer is: ')[1].strip()
        else:
            print("cannnot extract")
            answer = "None"
        answer_predicts.append(answer)

    # print(answer_predicts)
    # Count the occurrences of each answer in answer_predicts
    answer_counts = Counter(answer_predicts)

    # Find the most common answer
    final_predict = answer_counts.most_common(1)[0][0]

    generated_results.append({
        "generated_text": generated_texts,
        "beam_predict": answer_predicts,
        "predict": final_predict,
        "gt_answer": gt_answer
    })

    with open(save_generate_results_path, 'w') as f:
        json.dump(generated_results, f, indent=4)
    # Print or use the final_predict as needed
    print(f"Final predicted answer: {final_predict}, GT: {gt_answer}")

    counter.append(1 if final_predict == gt_answer else 0)

print(f"reasoner accuracy: {sum(counter)/len(counter)}, ({sum(counter)}/{len(counter)})")