bert-ensemble / distilbert.py
Jinwchoi06's picture
original commit
a5f918f verified
Raw
History Blame Contribute Delete
667 Bytes
from transformers import pipeline
import numpy
# distilbert tasks
# fill-mask
# feature extraction
# have to use more specific models/other models for other tasks
#for mask gen
unmasker = pipeline('fill-mask', model='distilbert-base-uncased', device=0)
results = unmasker("I [MASK] today.", top_k=10)
for prediction in results:
print(f"Score: {prediction['score']:.4f} | Prediction: {prediction['sequence']}")
#for feature extraction
extractor = pipeline('feature-extraction', model='distilbert-base-uncased', device=0)
results = extractor("Obama")
#idk what this really does
features = numpy.array(results)
print(f"Shape: {features.shape}")