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