Combine two versions of run_wikitext.py (#11597)
* Combine two versions of run_wikitext.py * Update run_wikitext.py * Update run_wikitext.py * aligned the format * update error display * simplified argument parser --------- Co-authored-by: jenniew <jenniewang123@gmail.com>
This commit is contained in:
parent
6f999e6e90
commit
1da1f1dd0e
1 changed files with 25 additions and 6 deletions
|
|
@ -20,16 +20,20 @@
|
|||
import argparse
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from datasets import concatenate_datasets, load_dataset
|
||||
from ipex_llm.utils.common import invalidInputError
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model_path", required=True, type=str)
|
||||
parser.add_argument("--data_path", type=str, default='wikitext-2-raw-v1/wikitext-2-raw/wiki.test.raw')
|
||||
parser.add_argument("--dataset", type=str, default=None)
|
||||
parser.add_argument("--data_path", type=str, default=None)
|
||||
parser.add_argument("--chunk_size", type=int, default=512)
|
||||
parser.add_argument("--stride", type=int, default=0)
|
||||
parser.add_argument("--device", type=str, default="xpu")
|
||||
parser.add_argument("--precision", type=str, default="sym_int4")
|
||||
parser.add_argument("--use-cache", action="store_true")
|
||||
parser.add_argument("--max_length", type=int, default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.precision == "fp16": # ipex fp16
|
||||
|
|
@ -46,14 +50,29 @@ else: # ipex-llm
|
|||
model = model.to(args.device)
|
||||
model = model.eval()
|
||||
|
||||
with open(args.data_path, "rb") as f:
|
||||
data = f.read()
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
|
||||
encodings = tokenizer(data.decode("utf-8").strip("\n"), return_tensors="pt")
|
||||
|
||||
if args.dataset:
|
||||
def parse_kwargs(kwstr):
|
||||
kvpair = [item.split('=') for item in kwstr.split(',') if item != ""]
|
||||
return {k:v for k, v in kvpair}
|
||||
test = load_dataset(**parse_kwargs(args.dataset), split="test")["text"]
|
||||
encodings = tokenizer("\n\n".join(test), return_tensors="pt")
|
||||
elif args.data_path:
|
||||
with open(args.data_path, "rb") as f:
|
||||
data = f.read()
|
||||
encodings = tokenizer(data.decode("utf-8").strip("\n"), return_tensors="pt")
|
||||
else:
|
||||
raise invalidInputError(False, "Must specify either dataset or datapath.")
|
||||
|
||||
if not args.max_length:
|
||||
try:
|
||||
max_length = model.config.max_position_embeddings
|
||||
except:
|
||||
max_length = model.config.seq_length # max_length in config of chatglm is 'seq_length'
|
||||
else:
|
||||
max_length = args.max_length
|
||||
stride = args.chunk_size if args.stride <= 0 else args.stride
|
||||
seq_len = encodings.input_ids.size(1)
|
||||
num_chunks = seq_len // stride
|
||||
|
|
|
|||
Loading…
Reference in a new issue