bert_pipeline_unmasker.py

# !pip install --upgrade transformers plotly pandas

import plotly.graph_objects as go
import pandas as pd
pd.options.display.max_columns = None
from transformers import pipeline
unmasker = pipeline('fill-mask', model='bert-base-uncased')

results = []
cars = ['mercedes', 'audi', 'bmw', 'volkswagen', 'ford', 'toyota',
        'peugeot', 'citroen', 'volvo', 'land rover', 'nissan', 'renault', 'fiat',
        'kia', 'lexus', 'jaguar', 'hyundai', 'honda', 'porsche']

for car in cars:
    result = unmasker(f'{car.title()} is a [MASK] car.')
    results.append(result)
    
df_list = []

for car, result in zip(cars, results):
    df = pd.json_normalize(result)
    df['car'] = car
    df_list.append(df)

car_df = pd.concat(df_list, ignore_index=True)
car_matrix = car_df.pivot(index='car', columns='token_str', values='score')

fig = px.imshow(car_matrix, color_continuous_scale='cividis',
                title='[CAR_BRAND] is a [MASK] car.<br>Bert: please fill in the blank',
                width=1000, height=800, template='none')
fig.layout.xaxis.side = 'top'
fig.layout.yaxis.title = ''
fig.layout.xaxis.title = ''
fig.layout.xaxis.tickangle = 300
fig

Produces this chart: