distilbert.py

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}")