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