# -*- coding:utf-8 -*-

"""
      ┏┛ ┻━━━━━┛ ┻┓
      ┃　　　　　　 ┃
      ┃　　　━　　　┃
      ┃　┳┛　  ┗┳　┃
      ┃　　　　　　 ┃
      ┃　　　┻　　　┃
      ┃　　　　　　 ┃
      ┗━┓　　　┏━━━┛
        ┃　　　┃   神兽保佑
        ┃　　　┃   代码无BUG！
        ┃　　　┗━━━━━━━━━┓
        ┃CREATE BY SNIPER┣┓
        ┃　　　　         ┏┛
        ┗━┓ ┓ ┏━━━┳ ┓ ┏━┛
          ┃ ┫ ┫   ┃ ┫ ┫
          ┗━┻━┛   ┗━┻━┛

"""

import tensorflow as tf
import numpy as np

for gpu in tf.config.experimental.list_physical_devices('GPU'):
    tf.config.experimental.set_memory_growth(gpu, True)

from transformers import BertTokenizer, TFBertForMaskedLM

tokenizer = BertTokenizer.from_pretrained('bert-base-cased')

model = TFBertForMaskedLM.from_pretrained('bert-base-cased', return_dict=True)

inputs = tokenizer("The capital of France is [MASK].", return_tensors="tf")

outputs = model(inputs)
logits = outputs.logits

output = np.argmax(logits[0][6])
o1 = tokenizer.decode(int(output))

inputs = tokenizer("The capital of [MASK] is BeiJing.", return_tensors="tf")

outputs = model(inputs)
logits = outputs.logits

output = np.argmax(logits[0][4])
o2 = tokenizer.decode(int(output))

print()

[((26, 11, 26, 62), 'tensorflow.config.experimental.list_physical_devices', 'tf.config.experimental.list_physical_devices', ({(26, 56, 26, 61): '"""GPU"""'}, {}), "('GPU')", True, 'import tensorflow as tf\n'), ((31, 12, 31, 60), 'transformers.BertTokenizer.from_pretrained', 'BertTokenizer.from_pretrained', ({(31, 42, 31, 59): '"""bert-base-cased"""'}, {}), "('bert-base-cased')", False, 'from transformers import BertTokenizer, TFBertForMaskedLM\n'), ((33, 8, 33, 78), 'transformers.TFBertForMaskedLM.from_pretrained', 'TFBertForMaskedLM.from_pretrained', (), '', False, 'from transformers import BertTokenizer, TFBertForMaskedLM\n'), ((40, 9, 40, 32), 'numpy.argmax', 'np.argmax', ({(40, 19, 40, 31): 'logits[0][6]'}, {}), '(logits[0][6])', True, 'import numpy as np\n'), ((48, 9, 48, 32), 'numpy.argmax', 'np.argmax', ({(48, 19, 48, 31): 'logits[0][4]'}, {}), '(logits[0][4])', True, 'import numpy as np\n'), ((27, 4, 27, 55), 'tensorflow.config.experimental.set_memory_growth', 'tf.config.experimental.set_memory_growth', ({(27, 45, 27, 48): 'gpu', (27, 50, 27, 54): '(True)'}, {}), '(gpu, True)', True, 'import tensorflow as tf\n')]