Spaces:
Sleeping
Sleeping
| """Create an optional INT8 ONNX variant; evaluate it before selecting for serving.""" | |
| import argparse | |
| import json | |
| import shutil | |
| from pathlib import Path | |
| from onnxruntime.quantization import quantize_dynamic, QuantType | |
| from image_to_latex.artifacts import digest | |
| def main(): | |
| p = argparse.ArgumentParser(description=__doc__) | |
| p.add_argument("--source", type=Path, default=Path("artifacts/latex-ocr")) | |
| p.add_argument("--output", type=Path, default=Path("artifacts/latex-ocr-int8")) | |
| args = p.parse_args() | |
| if args.output.exists(): | |
| p.error("Output must be new; original artifacts will not be overwritten") | |
| args.output.mkdir(parents=True) | |
| for name in ["encoder", "decoder"]: | |
| quantize_dynamic( | |
| str(args.source / f"{name}.onnx"), | |
| str(args.output / f"{name}.int8.onnx"), | |
| weight_type=QuantType.QInt8, | |
| per_channel=True, | |
| op_types_to_quantize=["MatMul", "Gemm"], | |
| extra_options={"MatMulConstBOnly": True}, | |
| ) | |
| for name in ["image_resizer.onnx", "tokenizer.json"]: | |
| shutil.copy2(args.source / name, args.output / name) | |
| metadata = { | |
| p.name: dict(bytes=p.stat().st_size, sha256=digest(p)) | |
| for p in args.output.iterdir() | |
| if p.is_file() | |
| } | |
| (args.output / "quantization.json").write_text( | |
| json.dumps(metadata, indent=2) + "\n" | |
| ) | |
| print(json.dumps(metadata, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |