File size: 420 Bytes
e9d1a5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
from transformers import BartTokenizer, BartModel
def main():
tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
model = BartModel.from_pretrained('facebook/bart-large')
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
outputs = model(**inputs)
H_all = outputs.last_hidden_state # noqa
print(H_all.shape) # (1, 8, 1024)
if __name__ == '__main__':
main()
|