def export()

in diffq/ts_export.py [0:0]

• 29 lines of code
• 5 McCabe index (conditional complexity)

def export(quantizer: DiffQuantizer, path: tp.Union[str, Path]):
    """Export the given quantized model to the given path.
    We must save the quantized model ourselves, as we need to recompress
    the zip archive afterwards.
    """
    packed: tp.List[_DiffQPacked] = []
    uniq_name = ''.join([random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(12)])
    with tempfile.TemporaryDirectory() as tmpdir:
        sys.path.insert(0, tmpdir)
        try:
            code = _codegen(quantizer)
            with open(Path(tmpdir) / f'{uniq_name}.py', 'w') as f:
                f.write(code)
            module = importlib.import_module(uniq_name)
            ts_klass = module.DiffQTSModel
            state = quantizer.get_quantized_state(packed=True, torch_pack=True)
            quantized = state["quantized"]
            for qparam in quantizer._qparams:
                if qparam.other is None:
                    levels, scales, bits = quantized.pop(0)
                    size = qparam.param.size()
                    packed.append((levels, scales, bits, list(size)))
                    qparam.param.data.zero_()
            quantizer.detach()
            ts_premodel = ts_klass(quantizer.model, quantizer.group_size,
                                   quantizer.min_bits, packed)
            ts_model = jit.script(ts_premodel)
            if path is not None:
                jit.save(ts_model, path)
                recompress(path)
        finally:
            sys.path.pop(0)

    return ts_model