from transformers import BertJapaneseTokenizer
from transformers import BertConfig
from transformers import BertForMaskedLM
from transformers import pipeline

inputs = ['[MASK]もそう思います', '[MASK]なんというかその', 'これは[MASK]私が子供の頃の話なんですけど']

model_name_list = ['models/1-6_layer-wise', 'models/tapt512_60K', 'models/dapt128-tapt512']


for input_, model_name in zip(inputs, model_name_list):

    tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)
    config = BertConfig.from_pretrained(model_name)
    model = BertForMaskedLM.from_pretrained(model_name)

    print('model name:',model_name)
    print('input:',input_)
    fill_mask = pipeline('fill-mask', model=model, tokenizer=tokenizer, config=config)
    print('output:',fill_mask(input_))
    print()