from auto_transformers import AutoTransformers
import onnx
import warnings

warnings.filterwarnings('ignore')

f = open('onnx.csv', 'a+')
f.write('model_name, run_op, save_onnx, check_onnx\n')

# full_models = AutoTransformers.supported_model_names()
# checked_models = AutoTransformers.supported_model_names(format='onnx')
# models = [x for x in full_models if x not in checked_models]
models = ['funnel-transformer/large', 'funnel-transformer/medium', 'funnel-transformer/small', 'funnel-transformer/xlarge']

for name in models:
    f.write(f'{name},')
    try:
        op = AutoTransformers(model_name=name)
        out1 = op('hello, world.')
        f.write('success,')
    except Exception as e:
        f.write('fail')
        print(f'Fail to load op for {name}: {e}')
        pass
    try:
        op.save_model(format='onnx')
        f.write('success,')
    except Exception as e:
        f.write('fail')
        print(f'Fail to save onnx for {name}: {e}')
        pass
    try:
        saved_name = name.replace('/', '-')
        onnx_model = onnx.load(f'saved/onnx/{saved_name}.onnx', load_external_data=False)
        onnx.checker.check_model(onnx_model)
        f.write('success')
    except Exception as e:
        f.write('fail')
        print(f'Fail to check onnx for {name}: {e}')
        pass
    
    f.write('\n')
print('Finished.')