Improved quantize script
I improved the quantize script by adding error handling and allowing to select many models for quantization at once in the command line. I also converted it to Python for generalization as well as extensibility.
This commit is contained in:
parent
721311070e
commit
f8db3d6cd9
1 changed files with 59 additions and 0 deletions
59
quantize.py
Normal file
59
quantize.py
Normal file
|
@ -0,0 +1,59 @@
|
||||||
|
#!/usr/bin/python3
|
||||||
|
|
||||||
|
"""Script to execute quantization on a given model."""
|
||||||
|
|
||||||
|
import subprocess
|
||||||
|
import argparse
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
"""Parse the command line arguments and execute the script."""
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
prog='Quantization Script',
|
||||||
|
description='This script quantizes a model.'
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument("models", nargs='+', dest='models')
|
||||||
|
parser.add_argument(
|
||||||
|
'-r', '--remove-16', action='store_true', dest='remove_f16',
|
||||||
|
help='Remove the f16 model after quantizing it.'
|
||||||
|
)
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
for model in args.models:
|
||||||
|
|
||||||
|
model_path = os.path.join("models", model, "ggml-model-f16.bin")
|
||||||
|
|
||||||
|
for i in os.listdir(model_path):
|
||||||
|
subprocess.run(
|
||||||
|
["./quantize", i, i.replace("f16", "q4_0"), "2"],
|
||||||
|
shell=True,
|
||||||
|
check=True
|
||||||
|
)
|
||||||
|
|
||||||
|
if args.remove_f16:
|
||||||
|
os.remove(i)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
try:
|
||||||
|
main()
|
||||||
|
|
||||||
|
except subprocess.CalledProcessError:
|
||||||
|
print("An error ocurred while trying to quantize the models.")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
except FileNotFoundError as err:
|
||||||
|
print(
|
||||||
|
f'A FileNotFoundError exception was raised while executing the \
|
||||||
|
script:\n{err}\nMake sure you are located in the root of the \
|
||||||
|
repository and that the models are in the "models" directory.'
|
||||||
|
)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
sys.exit(0)
|
Loading…
Add table
Add a link
Reference in a new issue