File size: 617 Bytes
a8639ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
import torch
from transformers import AutoTokenizer
from architecture import Transformer
t = Transformer()
t.to("mps")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
codes = ["def func(a, b):", "if x > 0:", "for i in range(10):"]
encoding = tokenizer(codes, padding=True, truncation=True, return_tensors="pt")
input_ids = encoding["input_ids"]
attention_mask = encoding["attention_mask"]
print("Input IDs:")
print(input_ids)
print("Attention Mask:")
print(attention_mask)
output = t(input_ids.to("mps"), padding_mask=attention_mask.to("mps"))
print("Transformer output:")
print(output)
|