from pat.models.pat_tokenizer import PATTokenizer
import torch


action_tokenizer = PATTokenizer.from_pretrained(
    "data/oat/oat1", device="cuda" if torch.cuda.is_available() else "cpu"
)

actions = torch.randn([1, 32, 20])
tokens = action_tokenizer.encode(actions=actions)

actions_recon = action_tokenizer.decode(tokens=tokens)

print(f"ori actions: {actions}")
print(f"tokens: {tokens}")
print(f"recon actions: {actions_recon}")
