| import argparse |
| import json |
| import math |
| import os |
| import re |
| from dataclasses import dataclass, field |
| from typing import Any, Optional |
|
|
| import torch |
| import transformers |
| import yaml |
| from accelerate.utils import set_seed |
| from datasets import load_dataset |
| from tqdm import tqdm |
|
|
|
|
| DEFAULT_DATA_PATHS = { |
| "gsm-hard": "/data/oss_bucket_0/liqi/dataset/gsm-hard/gsm-hard.json", |
| "multi-arith": "/data/oss_bucket_0/liqi/dataset/MultiArith/test.json", |
| "svamp": "/data/oss_bucket_0/liqi/dataset/SVAMP/test.json", |
| "gsm8k": "/data/oss_bucket_0/liqi/gsk8k-aug/gsm_test.json", |
| } |
| ANSWER_MARKER = "The answer is:" |
| NUMBER_PATTERN = re.compile(r"[-+]?(?:(?:\d{1,3}(?:,\d{3})+)|\d+|\.\d+)(?:\.\d+)?") |
|
|
|
|
| @dataclass |
| class ModelArguments: |
| model_name_or_path: str = field(metadata={"help": "Base causal language model path or Hugging Face ID."}) |
| token: Optional[str] = field(default=None) |
| torch_dtype: str = field(default="bfloat16") |
| device: str = field(default="auto") |
|
|
|
|
| @dataclass |
| class DataArguments: |
| data_name: str = field(default="gsm8k") |
| data_path: Optional[str] = field(default=None) |
| batch_size: int = field(default=16) |
|
|
|
|
| @dataclass |
| class GenerationArguments: |
| output_dir: str = field(default="./output") |
| run_name: str = field(default="zero_shot_cot") |
| max_input_tokens: int = field(default=512) |
| max_new_tokens: int = field(default=256) |
| do_sample: bool = field(default=False) |
| temperature: float = field(default=0.7) |
| top_p: float = field(default=0.95) |
| top_k: int = field(default=40) |
| seed: int = field(default=42) |
| use_chat_template: bool = field(default=False) |
| print_outputs: bool = field(default=True) |
|
|
|
|
| @dataclass |
| class EvalExample: |
| question: str |
| raw_answer: Any |
| answer_kind: str |
| normalized_answer: Any |
|
|
|
|
| def read_json(file_path: str) -> Any: |
| with open(file_path, "r", encoding="utf-8") as file: |
| return json.load(file) |
|
|
|
|
| def write_json(data: Any, file_path: str) -> None: |
| with open(file_path, "w", encoding="utf-8") as file: |
| json.dump(data, file, ensure_ascii=False, indent=4) |
|
|
|
|
| def normalize_gold_answer(answer: Any) -> Any: |
| if isinstance(answer, bool): |
| return answer |
|
|
| text = str(answer).strip() |
| if text.lower() == "true": |
| return True |
| if text.lower() == "false": |
| return False |
| if re.fullmatch(r"[A-E]", text, flags=re.IGNORECASE): |
| return text.upper() |
| if "####" in text: |
| text = text.rsplit("####", 1)[-1].strip() |
|
|
| try: |
| return float(text.replace(",", "")) |
| except ValueError: |
| return None |
|
|
|
|
| def infer_answer_kind(answer: Any) -> str: |
| if isinstance(answer, bool): |
| return "boolean" |
| if isinstance(answer, str) and re.fullmatch(r"[A-E]", answer): |
| return "choice" |
| return "number" |
|
|
|
|
| def load_examples(data_args: DataArguments) -> list[EvalExample]: |
| data_name = data_args.data_name.lower() |
| if data_name not in {*DEFAULT_DATA_PATHS, "commonsense"}: |
| raise ValueError(f"Unsupported data_name: {data_args.data_name}") |
|
|
| if data_name == "commonsense" and data_args.data_path is None: |
| raw_data = load_dataset("zen-E/CommonsenseQA-GPT4omini")["validation"] |
| else: |
| data_path = data_args.data_path or DEFAULT_DATA_PATHS[data_name] |
| raw_data = read_json(data_path) |
|
|
| examples = [] |
| for item in raw_data: |
| if data_name == "gsm-hard": |
| question, raw_answer = item["input"], item["target"] |
| elif data_name == "multi-arith": |
| question, raw_answer = item["question"], item["final_ans"] |
| elif data_name == "svamp": |
| question = f"{item['Body']} {item['Question']}" |
| raw_answer = item["Answer"] |
| else: |
| question, raw_answer = item["question"], item["answer"] |
|
|
| normalized_answer = normalize_gold_answer(raw_answer) |
| examples.append( |
| EvalExample( |
| question=question.strip().replace(" ", " "), |
| raw_answer=raw_answer, |
| answer_kind=infer_answer_kind(normalized_answer), |
| normalized_answer=normalized_answer, |
| ) |
| ) |
|
|
| if not examples: |
| raise ValueError("The evaluation dataset is empty.") |
| return examples |
|
|
|
|
| def build_prompt(question: str) -> str: |
| return ( |
| f"Question: {question}\n" |
| "Answer the question by reasoning step by step. Keep the entire response concise " |
| "and within 200 tokens. For multiple-choice questions, give the option letter; " |
| "for true/false questions, give True or False. " |
| f"End with exactly `{ANSWER_MARKER} <final answer>`.\n" |
| "Answer: Let's think step by step." |
| ) |
|
|
|
|
| def format_prompt(question: str, tokenizer: transformers.PreTrainedTokenizerBase, use_chat_template: bool) -> str: |
| prompt = build_prompt(question) |
| if not use_chat_template: |
| return prompt |
| if tokenizer.chat_template is None: |
| raise ValueError("use_chat_template=True, but the tokenizer has no chat template.") |
| return tokenizer.apply_chat_template( |
| [{"role": "user", "content": prompt}], |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
|
|
|
|
| def answer_segment(text: str) -> tuple[str, bool]: |
| matches = list(re.finditer(r"The answer is\s*:\s*", text, flags=re.IGNORECASE)) |
| if not matches: |
| return text, False |
| return text[matches[-1].end():], True |
|
|
|
|
| def extract_prediction(text: str, answer_kind: str) -> Any: |
| segment, has_marker = answer_segment(text) |
|
|
| if answer_kind == "choice": |
| matches = re.findall(r"\b(?:option\s*)?([A-E])\b", segment, flags=re.IGNORECASE) |
| return matches[-1].upper() if matches else None |
|
|
| if answer_kind == "boolean": |
| matches = re.findall(r"\b(True|False)\b", segment, flags=re.IGNORECASE) |
| if not matches: |
| return None |
| return matches[-1].lower() == "true" |
|
|
| matches = NUMBER_PATTERN.findall(segment.replace("$", "")) |
| if not matches: |
| return None |
| number = matches[0] if has_marker else matches[-1] |
| try: |
| return float(number.replace(",", "")) |
| except ValueError: |
| return None |
|
|
|
|
| def compute_accuracy(gold: list[Any], predictions: list[Any]) -> float: |
| return sum(prediction == answer for answer, prediction in zip(gold, predictions)) / len(gold) |
|
|
|
|
| def resolve_dtype(dtype_name: str) -> torch.dtype: |
| dtype_map = { |
| "auto": None, |
| "bfloat16": torch.bfloat16, |
| "float16": torch.float16, |
| "float32": torch.float32, |
| } |
| if dtype_name not in dtype_map: |
| raise ValueError(f"Unsupported torch_dtype: {dtype_name}") |
| return dtype_map[dtype_name] |
|
|
|
|
| def evaluate( |
| model_args: ModelArguments, |
| data_args: DataArguments, |
| generation_args: GenerationArguments, |
| ) -> float: |
| examples = load_examples(data_args) |
| device = model_args.device |
| if device == "auto": |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| tokenizer = transformers.AutoTokenizer.from_pretrained( |
| model_args.model_name_or_path, |
| token=model_args.token, |
| padding_side="left", |
| truncation_side="left", |
| use_fast=False, |
| ) |
| if tokenizer.pad_token_id is None: |
| if tokenizer.eos_token_id is None: |
| raise ValueError("The tokenizer must define either a pad token or an EOS token.") |
| tokenizer.pad_token = tokenizer.eos_token |
|
|
| torch_dtype = resolve_dtype(model_args.torch_dtype) |
| model_load_kwargs = {"token": model_args.token} |
| if torch_dtype is not None: |
| model_load_kwargs["torch_dtype"] = torch_dtype |
| model = transformers.AutoModelForCausalLM.from_pretrained( |
| model_args.model_name_or_path, |
| **model_load_kwargs, |
| ).to(device) |
| model.eval() |
| set_seed(generation_args.seed) |
|
|
| prompts = [ |
| format_prompt(example.question, tokenizer, generation_args.use_chat_template) |
| for example in examples |
| ] |
| predictions = [] |
| full_outputs = [] |
| output_lengths = [] |
| eval_steps = math.ceil(len(examples) / data_args.batch_size) |
|
|
| for step in tqdm(range(eval_steps), desc="Evaluating", unit="batch"): |
| start = step * data_args.batch_size |
| end = min(start + data_args.batch_size, len(examples)) |
| batch_prompts = prompts[start:end] |
| encoded = tokenizer( |
| batch_prompts, |
| return_tensors="pt", |
| padding=True, |
| truncation=True, |
| max_length=generation_args.max_input_tokens, |
| ).to(device) |
|
|
| generate_kwargs = { |
| "max_new_tokens": generation_args.max_new_tokens, |
| "do_sample": generation_args.do_sample, |
| "pad_token_id": tokenizer.pad_token_id, |
| "eos_token_id": tokenizer.eos_token_id, |
| "use_cache": True, |
| } |
| if generation_args.do_sample: |
| generate_kwargs.update( |
| temperature=generation_args.temperature, |
| top_p=generation_args.top_p, |
| top_k=generation_args.top_k, |
| ) |
|
|
| with torch.inference_mode(): |
| generated = model.generate(**encoded, **generate_kwargs) |
|
|
| generated_tokens = generated[:, encoded["input_ids"].shape[1]:] |
| for batch_index, token_ids in enumerate(generated_tokens): |
| token_list = token_ids.tolist() |
| if tokenizer.eos_token_id in token_list: |
| token_list = token_list[: token_list.index(tokenizer.eos_token_id) + 1] |
| decoded = tokenizer.decode(token_list, skip_special_tokens=True).strip() |
| example = examples[start + batch_index] |
| prediction = extract_prediction(decoded, example.answer_kind) |
|
|
| full_outputs.append(decoded) |
| output_lengths.append(len(token_list)) |
| predictions.append(prediction) |
|
|
| accuracy = compute_accuracy( |
| [example.normalized_answer for example in examples], |
| predictions, |
| ) |
| save_list = [ |
| { |
| "question": example.question, |
| "ground truth": example.raw_answer, |
| "full output": output, |
| "prediction": prediction, |
| } |
| for example, output, prediction in zip(examples, full_outputs, predictions) |
| ] |
| save_dir = os.path.join(generation_args.output_dir, generation_args.run_name) |
| os.makedirs(save_dir, exist_ok=True) |
| write_json( |
| { |
| "accuracy": 100 * accuracy, |
| "average length of COT": sum(output_lengths) / len(output_lengths), |
| "output": save_list, |
| }, |
| os.path.join(save_dir, "answer.json"), |
| ) |
| print(f"Zero-shot CoT accuracy: {100 * accuracy:.2f}%") |
| print(f"Average length of CoT: {sum(output_lengths) / len(output_lengths):.2f}") |
| return 100 * accuracy |
|
|
|
|
| def parse_arguments() -> tuple[ModelArguments, DataArguments, GenerationArguments]: |
| config_parser = argparse.ArgumentParser(add_help=False) |
| config_parser.add_argument("--config_file", type=str) |
| config_args, remaining_args = config_parser.parse_known_args() |
| parser = transformers.HfArgumentParser((ModelArguments, DataArguments, GenerationArguments)) |
|
|
| if config_args.config_file is None: |
| return parser.parse_args_into_dataclasses(args=remaining_args) |
|
|
| with open(config_args.config_file, "r", encoding="utf-8") as file: |
| config = yaml.safe_load(file) or {} |
| config_cli_args = [] |
| for key, value in config.items(): |
| if value is not None: |
| config_cli_args.extend([f"--{key}", str(value)]) |
| return parser.parse_args_into_dataclasses(args=config_cli_args + remaining_args) |
|
|
|
|
| if __name__ == "__main__": |
| parsed_model_args, parsed_data_args, parsed_generation_args = parse_arguments() |
| evaluate(parsed_model_args, parsed_data_args, parsed_generation_args) |
|
|
| |
