bert-ensemble / google.py
Jinwchoi06's picture
original commit
a5f918f verified
Raw
History Blame Contribute Delete
386 Bytes
from transformers import pipeline
# google bert tasks
# fill-mask
# feature extraction
# have to use more specific models/other models for other tasks
# for mask gen
unmasker = pipeline('fill-mask', model='bert-base-uncased')
results = unmasker("Obama is [MASK].")
for prediction in results:
print(f"Score: {prediction['score']:.4f} | Prediction: {prediction['sequence']}")