from neural_compressor.experimental import Quantization, common

class Dataset(object):
    def __init__(self, *args):
        # TODO:initialize dataset related info here
        pass

    def __getitem__(self, index):
        # TODO:get item magic method
        # return a tuple containing 1 image and 1 label
        # for example, return img, label
        pass

    def __len__(self):
        # TODO:get total length of dataset, such as how many images in the dataset
        # if the total length is not able to know, pls implement __iter__() magic method
        # rather than above two methods.
        pass

quantizer = Quantization('{{config_path}}')
quantizer.model = '{{model_path}}'
dataset = Dataset()
quantizer.calib_dataloader = common.DataLoader(dataset, batch_size=1)
quantizer.eval_dataloader = common.DataLoader(dataset, batch_size=1)
quantized_model = quantizer.fit()
quantized_model.save('{{model_output_path}}')
