File size: 1,346 Bytes
2051791 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | import os
import time
from hqq.engine.hf import AutoTokenizer as hggAutoTokenizer
from hqq.engine.hf import HQQModelForCausalLM
def create_mxq_model(model_id, quant_config, config_id, load_quantized, save_dir):
quantized = False
model_file_size = 0
quant_path = f"{save_dir}/{model_id}-{config_id}-mxq"
if load_quantized and os.path.exists(quant_path):
model = HQQModelForCausalLM.from_quantized(quant_path)
tokenizer = hggAutoTokenizer.from_pretrained(model_id)
quantized = True
model_file_size = os.path.getsize(os.path.join(quant_path, "qmodel.pt"))
else:
model = HQQModelForCausalLM.from_pretrained(model_id)
tokenizer = hggAutoTokenizer.from_pretrained(model_id)
return model, tokenizer, quantized, model_file_size
def quantize_mxq_model(model, tokenizer, quant_config, model_id, config_id, save_dir):
model_file_size = 0
t1 = time.time()
model.quantize_model(quant_config=quant_config)
t2 = time.time()
print("Took " + str(t2 - t1) + " seconds to quantize the model with MXQ")
quant_path = f"{save_dir}/{model_id}-{config_id}-mxq"
model.save_quantized(quant_path)
# persistent the quantized model
os.sync()
model_file_size = os.path.getsize(os.path.join(quant_path, "qmodel.pt"))
return model, t2 - t1, model_file_size
|