Commit b9ab9a64 authored by Casper's avatar Casper
Browse files

Add wikitext example

parent d832a218
...@@ -10,6 +10,7 @@ quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": ...@@ -10,6 +10,7 @@ quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version":
model = AutoAWQForCausalLM.from_pretrained(model_path) model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
# Define data loading methods
def load_dolly(): def load_dolly():
data = load_dataset('databricks/databricks-dolly-15k', split="train") data = load_dataset('databricks/databricks-dolly-15k', split="train")
...@@ -20,8 +21,12 @@ def load_dolly(): ...@@ -20,8 +21,12 @@ def load_dolly():
concatenated = data.map(concatenate_data) concatenated = data.map(concatenate_data)
return [text for text in concatenated["text"]] return [text for text in concatenated["text"]]
def load_wikitext():
data = load_dataset('wikitext', 'wikitext-2-raw-v1', split="train")
return [text for text in data["text"] if text.strip() != '' and len(text.split(' ')) > 20]
# Quantize # Quantize
model.quantize(tokenizer, quant_config=quant_config, calib_data=load_dolly()) model.quantize(tokenizer, quant_config=quant_config, calib_data=load_wikitext())
# Save quantized model # Save quantized model
model.save_quantized(quant_path) model.save_quantized(quant_path)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment