File size: 807 Bytes
1be5b40 |
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 |
import numpy as np
from openpi.models import tokenizer as _tokenizer
def test_tokenize():
tokenizer = _tokenizer.PaligemmaTokenizer(max_len=10)
tokens, masks = tokenizer.tokenize("Hello, world!")
assert tokens.shape == (10,)
assert masks.shape == (10,)
def test_fast_tokenizer():
prompt = "Hello, world!"
state = np.random.rand(5).astype(np.float32)
action = np.random.rand(3, 2).astype(np.float32)
tokenizer = _tokenizer.FASTTokenizer(max_len=256)
tokens, token_masks, ar_masks, loss_masks = tokenizer.tokenize(prompt, state, action)
assert tokens.shape == (256,)
assert token_masks.shape == (256,)
assert ar_masks.shape == (256,)
assert loss_masks.shape == (256,)
act = tokenizer.extract_actions(tokens, 3, 2)
assert act.shape == (3, 2)
|