bert-ensemble / demo.py
Jinwchoi06's picture
original commit
a5f918f verified
Raw
History Blame Contribute Delete
1.74 kB
#high level way
from transformers import pipeline
from transformers import AutoConfig
from transformers import AutoModelForSequenceClassification
#pipeline(task, model, tokenizer, device (0 being gpu))
#['any-to-any', 'audio-classification', 'automatic-speech-recognition', 'depth-estimation',
# 'document-question-answering', 'feature-extraction', 'fill-mask', 'image-classification',
# 'image-feature-extraction', 'image-segmentation', 'image-text-to-text', 'keypoint-matching',
# 'mask-generation', 'ner'(named entity recognition), 'object-detection', 'sentiment-analysis', 'table-question-answering',
# 'text-classification', 'text-generation', 'text-to-audio', 'text-to-speech', 'token-classification',
# 'video-classification', 'zero-shot-audio-classification', 'zero-shot-classification',
# 'zero-shot-image-classification', 'zero-shot-object-detection']
# Load the model (this downloads weights automatically on first run)
#classifier = pipeline("text-classification", model="distilbert-base-uncased")
#for mask gen
predictions = []
unmasker = pipeline('fill-mask', model='distilbert-base-uncased', device=0)
results = unmasker("Obama is [MASK] today.", top_k=1)
predictions.append(results[0])
unmasker1 = pipeline('fill-mask', model='roberta-base', device=0)
results1 = unmasker1("Obama is <mask> today.", top_k=1)
predictions.append(results1[0])
unmasker2 = pipeline('fill-mask', model='bert-base-uncased', top_k=1)
results2 = unmasker2("Obama is [MASK] today.")
predictions.append(results2[0])
best_overall = max(predictions, key=lambda x: x['score'])
print(f"The winning phrase was: {best_overall['sequence']}")
print(f"With a confidence score of: {best_overall['score']:.4f}")