from transformers import AutoTokenizer, AutoModelForMaskedLM
from transformers import RobertaTokenizer, RobertaTokenizerFast, RobertaForMaskedLM, pipeline
import torch

def evaluate(framework):
    text = "På biblioteket kan du [MASK] en bok."
    if framework == "flax":
        model = AutoModelForMaskedLM.from_pretrained("./", from_flax=True)
    elif framework == "tensorflow":
        model = AutoModelForMaskedLM.from_pretrained("./", from_tf=True)
    else:
        model = AutoModelForMaskedLM.from_pretrained("./")

    print("Testing with AutoTokenizer")
    tokenizer = AutoTokenizer.from_pretrained("./")
    my_unmasker_pipeline = pipeline('fill-mask', model=model, tokenizer=tokenizer)
    output = my_unmasker_pipeline(text)
    print(output)

    #print("\n\nTesting with RobertaTokenizer")
    #tokenizer = RobertaTokenizer.from_pretrained("./")
    #my_unmasker_pipeline = pipeline('fill-mask', model=model, tokenizer=tokenizer)
    #output = my_unmasker_pipeline(text)
    #print(output)

    #print("\n\nTesting with RobertaTokenizerFast")
    #tokenizer = RobertaTokenizerFast.from_pretrained("./")
    #my_unmasker_pipeline = pipeline('fill-mask', model=model, tokenizer=tokenizer)
    #output = my_unmasker_pipeline(text)
    #print(output)


print("Evaluating PyTorch Model")
evaluate("pytorch")


#print("Evaluating Flax Model")
#evaluate("flax")