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 argparse
|
||||||
import torch
|
import torch
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
from datasets import concatenate_datasets, load_dataset
|
||||||
|
from ipex_llm.utils.common import invalidInputError
|
||||||
|
|
||||||
|
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
parser.add_argument("--model_path", required=True, type=str)
|
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("--chunk_size", type=int, default=512)
|
||||||
parser.add_argument("--stride", type=int, default=0)
|
parser.add_argument("--stride", type=int, default=0)
|
||||||
parser.add_argument("--device", type=str, default="xpu")
|
parser.add_argument("--device", type=str, default="xpu")
|
||||||
parser.add_argument("--precision", type=str, default="sym_int4")
|
parser.add_argument("--precision", type=str, default="sym_int4")
|
||||||
parser.add_argument("--use-cache", action="store_true")
|
parser.add_argument("--use-cache", action="store_true")
|
||||||
|
parser.add_argument("--max_length", type=int, default=None)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
if args.precision == "fp16": # ipex fp16
|
if args.precision == "fp16": # ipex fp16
|
||||||
|
|
@ -46,14 +50,29 @@ else: # ipex-llm
|
||||||
model = model.to(args.device)
|
model = model.to(args.device)
|
||||||
model = model.eval()
|
model = model.eval()
|
||||||
|
|
||||||
with open(args.data_path, "rb") as f:
|
|
||||||
data = f.read()
|
|
||||||
|
|
||||||
from transformers import AutoTokenizer
|
from transformers import AutoTokenizer
|
||||||
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
|
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
|
||||||
encodings = tokenizer(data.decode("utf-8").strip("\n"), return_tensors="pt")
|
|
||||||
|
|
||||||
max_length = model.config.max_position_embeddings
|
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
|
stride = args.chunk_size if args.stride <= 0 else args.stride
|
||||||
seq_len = encodings.input_ids.size(1)
|
seq_len = encodings.input_ids.size(1)
|
||||||
num_chunks = seq_len // stride
|
num_chunks = seq_len // stride
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue