File size: 1,965 Bytes
c095142 a82e5d8 c095142 | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | from collections import UserDict
import torch
from torch import nn
from cortex_ai.vision.native import build_native_fusion
class TinyCore(nn.Module):
def __init__(self):
super().__init__()
self.embedding = nn.Embedding(32, 8)
self.last_inputs_embeds = None
def get_input_embeddings(self):
return self.embedding
def forward(self, inputs_embeds, attention_mask):
self.last_inputs_embeds = inputs_embeds
return {"hidden_states": inputs_embeds, "attention_mask": attention_mask}
class TinyVisionEncoder(nn.Module):
def forward(self, image):
return torch.ones(image.shape[0], 3, 4)
def test_native_image_embedding_fusion_executes_real_tensor_path():
core = TinyCore()
fusion = build_native_fusion(TinyVisionEncoder(), core, vision_hidden_size=4, cortex_hidden_size=8)
result = fusion(input_ids=torch.tensor([[1, 2]]), image=torch.zeros(1, 3, 16, 16))
assert result.image_token_count == 3
assert result.text_token_count == 2
assert core.last_inputs_embeds.shape == (1, 5, 8)
assert result.outputs["attention_mask"].shape == (1, 5)
class BatchFeatureLikeEncoder(nn.Module):
"""Encoder whose processor output is a UserDict, as Transformers processors return."""
def forward(self, pixel_values):
return torch.ones(pixel_values.shape[0], 2, 4)
class BatchFeatureLikeProcessor:
def __call__(self, image, return_tensors=None):
return UserDict({"pixel_values": torch.zeros(1, 3, 16, 16)})
def test_processor_output_is_unpacked_by_keyword_not_positionally():
core = TinyCore()
fusion = build_native_fusion(
BatchFeatureLikeEncoder(),
core,
vision_hidden_size=4,
cortex_hidden_size=8,
image_processor=BatchFeatureLikeProcessor(),
)
result = fusion(input_ids=torch.tensor([[1, 2]]), image="unused")
assert result.image_token_count == 2
assert result.text_token_count == 2
|