from auto_transformers import AutoTransformers

import torch

models = ['bert-base-cased', 'distilbert-base-cased', 'distilgpt2']

for name in models:
    try:
        op = AutoTransformers(model_name=name)
        out1 = op('hello, world.')
        op.save_model()
        op.model = torch.jit.load(name + '.pt')
        out2 = op('hello, world.')
        assert (out1 == out2).all()
        print(f'[SUCCESS] Saved torchscript for model "{name}"')
    except Exception as e:
        print(f'[ERROR] Fail for model "{name}": {e}.')
        continue