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")
 | 
			
		||||
 | 
			
		||||
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
 | 
			
		||||
seq_len = encodings.input_ids.size(1)
 | 
			
		||||
num_chunks = seq_len // stride
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
		Loading…
	
		Reference in a new issue