seeker
62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
import torch
|
|
|
|
from pi05mini.model import Pi05Config, Pi05Mini
|
|
|
|
|
|
def tiny_model() -> Pi05Mini:
|
|
return Pi05Mini(Pi05Config(
|
|
vocab_size=300,
|
|
image_size=16,
|
|
patch_size=16,
|
|
d_model=32,
|
|
n_heads=4,
|
|
n_layers=1,
|
|
max_text_tokens=16,
|
|
max_action_tokens=4,
|
|
))
|
|
|
|
|
|
def test_pi05_has_no_continuous_state_projection() -> None:
|
|
model = tiny_model()
|
|
assert not hasattr(model, "state_projection")
|
|
output = model(
|
|
torch.randn(2, 3, 16, 16),
|
|
torch.randint(0, 300, (2, 7)),
|
|
torch.randn(2, 4, 2),
|
|
torch.rand(2),
|
|
generated_token_start=4,
|
|
)
|
|
assert output.text_logits.shape == (2, 7, 300)
|
|
assert output.vector_field.shape == (2, 4, 2)
|
|
|
|
|
|
def test_observation_generated_and_flow_attention_partition() -> None:
|
|
model = tiny_model()
|
|
# One image patch followed by four observation tokens and three generated
|
|
# tokens. Flow may read observation plus the first generated (subtask)
|
|
# token, but not the two teacher-forced FAST tokens.
|
|
action_prefix_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 0, 0]], dtype=torch.bool)
|
|
prefix, action = model._masks(
|
|
batch=1,
|
|
prefix_length=8,
|
|
action_length=2,
|
|
text_offset=1,
|
|
attention_mask=torch.ones(1, 7, dtype=torch.bool),
|
|
action_prefix_mask=action_prefix_mask,
|
|
generated_token_start=4,
|
|
device=torch.device("cpu"),
|
|
)
|
|
|
|
generated_start = 5 # image patch + four observation text tokens
|
|
# Observation block is fully bidirectional and cannot see generated tokens.
|
|
assert prefix[0, 0, :generated_start].all()
|
|
assert not prefix[0, 0, generated_start:].any()
|
|
assert prefix[0, 4, :generated_start].all()
|
|
# Generated tokens see all observation tokens and only prior/self generated tokens.
|
|
assert prefix[0, 5, :6].all() and not prefix[0, 5, 6:].any()
|
|
assert prefix[0, 6, :7].all() and not prefix[0, 6, 7]
|
|
# Flow reads allowed prefix keys and all flow-action keys bidirectionally.
|
|
assert action[0, :, :6].all()
|
|
assert not action[0, :, 6:8].any()
|
|
assert action[0, :, 8:].all()
|