import numpy as np
from transformers import pipeline
from timeit import default_timer as timer

SENTENCES = [
    'EÚ môže vynechať Slovensko a Maďarsko z ##mask## na ruskú ropu.',
    'Manželka amerického prezidenta Jill Bidenová navštívi Slovensko, kde sa počas Dňa ##mask## stretne s ukrajinskými matkami a deťmi, ktoré utiekli z vlasti po začiatku ruskej invázie.',
    'Namiesto zanedbaného územia s ležoviskami divých svíň môžu mať ##mask## v horizonte niekoľkých rokov novú štvrť pre 5-tisíc obyvateľov.'
]

MODELS_TO_TEST = [
    'fav-kky/FERNET-CC_sk',
    # 'gerulata/slovakbert',
    # '/data/ivan/experiment-1',
    # '/data/ivan/experiment-2',
    # '/data/ivan/experiment-3',
    # '/data/ivan/experiment-4',
    # '/data/ivan/experiment-5',
    # '/data/ivan/experiment-6'
]


def get_examples(mask_token='<mask>'):
    examples = []
    for s in SENTENCES:
        examples.append(s.replace('##mask##', mask_token))
    return examples


def main():
    for model_name in MODELS_TO_TEST:
        unmasker = pipeline(model=model_name, task='fill-mask', device=0)

        average_300_examples_time = []
        examples = get_examples('[MASK]' if 'FERNET' in model_name else '<mask>')
        for _ in range(20):
            start = timer()

            for _ in range(100):
                for s in examples:
                    _ = unmasker(s)

            end = timer()
            average_300_examples_time.append(end - start)

        print(f'{model_name} 300 examples time: mean: {np.mean(average_300_examples_time)}, median: {np.median(average_300_examples_time)}')


if __name__ == '__main__':
    main()
