| """Vision tower + TinyVLM tests (CPU-only, no downloads).""" |
|
|
| import sys |
| from pathlib import Path |
|
|
| import pytest |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| torch = pytest.importorskip("torch") |
|
|
| from qwenjev.vision import ( |
| BOS_ID, |
| EOT_ID, |
| PAD_ID, |
| TinyVLM, |
| TinyVLMConfig, |
| grid_to_tensor, |
| render_grid_pil, |
| text_to_byte_ids, |
| ) |
|
|
|
|
| GRID = [ |
| [1, 1, 1, 1, 1], |
| [1, 2, 0, 3, 1], |
| [1, 0, 4, 0, 1], |
| [1, 1, 1, 1, 1], |
| ] |
|
|
|
|
| def test_text_to_byte_ids_roundtrip(): |
| text = "STATE:\nqwenjev-state-v1" |
| ids = text_to_byte_ids(text, max_len=512) |
| assert ids[0] == BOS_ID and ids[-1] == EOT_ID |
| body = bytes(ids[1:-1]).decode("utf-8") |
| assert body == text |
|
|
|
|
| def test_text_truncation(): |
| ids = text_to_byte_ids("x" * 999, max_len=64) |
| assert len(ids) == 64 and ids[-1] == EOT_ID |
|
|
|
|
| def test_grid_to_tensor_values(): |
| t = grid_to_tensor(GRID, image_size=64) |
| assert t.shape == (3, 64, 64) |
| assert 0.0 <= t.min().item() and t.max().item() <= 1.0 |
| |
| assert t.sum().item() > 0 |
|
|
|
|
| def test_grid_to_tensor_distinguishes_values(): |
| g1 = [[0, 0], [0, 0]] |
| g2 = [[0, 1], [0, 0]] |
| t1, t2 = grid_to_tensor(g1, 32), grid_to_tensor(g2, 32) |
| assert not torch.allclose(t1, t2) |
|
|
|
|
| def test_render_grid_pil(): |
| pil = pytest.importorskip("PIL") |
| img = render_grid_pil(GRID, cell_px=4) |
| assert isinstance(img, pil.Image.Image) |
| assert img.size == (5 * 4, 4 * 4) |
|
|
|
|
| def test_tinyvlm_forward_shapes(): |
| m = TinyVLM(TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64)) |
| px = torch.stack([grid_to_tensor(GRID, 64), grid_to_tensor(GRID, 64)]) |
| ids = torch.tensor([ |
| text_to_byte_ids("abcd", 32), |
| text_to_byte_ids("defg", 32), |
| ]) |
| out, hs = m(px, ids, output_hidden_states=True) |
| n_vis = m.n_vision_tokens |
| assert out.shape == (2, n_vis + ids.shape[1], 64) |
| |
| assert len(hs) == 3 |
| assert all(h.shape == out.shape for h in hs) |
|
|
|
|
| def test_tinyvlm_all_finite_with_padding(): |
| """Padding + causal masks must never produce NaN/Inf (regression).""" |
| torch.manual_seed(0) |
| m = TinyVLM(TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64)) |
| px = torch.stack([grid_to_tensor(GRID, 64)]) |
| long_ids = text_to_byte_ids("a long state text " * 3, 128) |
| short_ids = text_to_byte_ids("tiny", 128) |
| ids = torch.tensor([long_ids, short_ids + [PAD_ID] * (len(long_ids) - len(short_ids))]) |
| attn = torch.tensor([ |
| [1] * len(long_ids), |
| [1] * len(short_ids) + [0] * (len(long_ids) - len(short_ids)), |
| ]) |
| out, hs = m(px.expand(2, -1, -1, -1), ids, attention_mask=attn, |
| output_hidden_states=True) |
| for h in hs: |
| assert torch.isfinite(h).all(), "non-finite hidden states (mask bug)" |
| assert torch.isfinite(out).all() |
|
|
|
|
| def test_tinyvlm_pooled_uses_last_text_position(): |
| m = TinyVLM(TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64)) |
| px = torch.stack([grid_to_tensor(GRID, 64)]) |
| ids = torch.tensor([text_to_byte_ids("hello", 32)]) |
| out, _ = m(px, ids) |
| n_vis = m.n_vision_tokens |
| pooled = m.pooled(out, torch.tensor([ids.shape[1]]) + n_vis) |
| assert torch.allclose(pooled[0], out[0, n_vis + ids.shape[1] - 1]) |
|
|
|
|
| def test_tinyvlm_vision_changes_representation(): |
| """The vision tower MUST matter: different frames -> different pooled.""" |
| torch.manual_seed(3) |
| m = TinyVLM(TinyVLMConfig(image_size=64, patch_size=8, lm_dim=64)).eval() |
| g1 = [[0] * 8 for _ in range(8)] |
| g2 = [[0] * 8 for _ in range(8)] |
| g2[3][3] = 7 |
| ids = torch.tensor([text_to_byte_ids("state", 32)]) |
| with torch.no_grad(): |
| o1, _ = m(torch.stack([grid_to_tensor(g1, 64)]), ids) |
| o2, _ = m(torch.stack([grid_to_tensor(g2, 64)]), ids) |
| assert not torch.allclose(o1[0, 0], o2[0, 0]) |
|
|
