import os import json import argparse import torch import random import glog from lm_eval import evaluator from eval_utils import LMEvalAdaptor from tokenization_bitnet import BitnetTokenizer from modeling_bitnet import BitnetForCausalLM parser = argparse.ArgumentParser() parser.add_argument('--seed', default=0, type=int) parser.add_argument('--hf_path', default='1bitLLM/bitnet_b1_58-3B', type=str) parser.add_argument('--batch_size', type=int, default=1, help='batch size') parser.add_argument("--tasks", type=str) parser.add_argument("--output_path", default=None, type=str) parser.add_argument('--num_fewshot', type=int, default=0) parser.add_argument('--ctx_size', default=2048, type=int) def main(args): model_str = args.hf_path model = BitnetForCausalLM.from_pretrained( args.hf_path, device_map='auto', low_cpu_mem_usage=True, use_flash_attention_2=True, torch_dtype=torch.float16, ).half() tokenizer = BitnetTokenizer.from_pretrained(args.hf_path, use_fast=False) glog.info('loaded model!') task_names = args.tasks.split(",") lm_eval_model = LMEvalAdaptor(model_str, model, tokenizer, args.batch_size, args.ctx_size) results = evaluator.simple_evaluate( model=lm_eval_model, tasks=task_names, batch_size=args.batch_size, no_cache=True, num_fewshot=args.num_fewshot, ) print(evaluator.make_table(results)) if args.output_path is not None: os.makedirs(os.path.dirname(args.output_path), exist_ok=True) # otherwise cannot save results["config"]["model"] = args.hf_path with open(args.output_path, "w") as f: json.dump(results, f, indent=2) if __name__ == '__main__': torch.set_grad_enabled(False) args = parser.parse_args() random.seed(args.seed) torch.random.manual_seed(args.seed) main(args)