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)