File size: 5,324 Bytes
f5eae9c ba3bbc7 f5eae9c ba3bbc7 f5eae9c ba3bbc7 f5eae9c ba3bbc7 f5eae9c | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """Modular vision adaptation interfaces for CORTEX AI.
The module intentionally does not ship vision weights. A real encoder must be
provided by the caller (for example, a Transformers vision model loaded from a
separate approved checkpoint). The adapter can therefore be tested with a
small local encoder while remaining honest about deployment readiness.
"""
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Protocol, Sequence
import torch
from torch import Tensor, nn
@dataclass(frozen=True)
class VisionConfig:
"""Configuration for the image-to-language adaptation layer."""
vision_hidden_size: int
cortex_hidden_size: int = 7168
projector_type: str = "linear"
image_token_count: int | None = None
status: str = "experimental / in development"
@dataclass
class MultimodalInput:
"""A language prompt plus projected image tokens for a future fusion path."""
input_ids: Tensor | None
image_embeddings: Tensor
attention_mask: Tensor | None = None
metadata: dict[str, Any] | None = None
class VisionEncoder(Protocol):
"""Minimal callable contract for an image encoder."""
def __call__(self, images: Any) -> Tensor: ...
class ImageProcessor(Protocol):
"""Minimal preprocessing contract, compatible with HF processors."""
def __call__(self, images: Any, **kwargs: Any) -> Any: ...
class VisionProjector(nn.Module):
"""Map encoder features into CORTEX hidden space.
``projector_type='linear'`` is a single linear map. ``projector_type='mlp'`` is the
LLaVA-style two-layer gated projection, whose intermediate width equals the target
hidden size. Only this module is trained in the first experimental stage.
"""
def __init__(self, config: VisionConfig) -> None:
super().__init__()
if config.projector_type == "linear":
self.projection: nn.Module = nn.Linear(
config.vision_hidden_size, config.cortex_hidden_size
)
elif config.projector_type == "mlp":
self.projection = nn.Sequential(
nn.Linear(config.vision_hidden_size, config.cortex_hidden_size),
nn.GELU(),
nn.Linear(config.cortex_hidden_size, config.cortex_hidden_size),
)
else:
raise ValueError(
f"Unknown projector_type={config.projector_type!r}; "
"expected 'linear' or 'mlp'"
)
self.config = config
def forward(self, features: Tensor) -> Tensor:
if features.ndim != 3:
raise ValueError("vision features must have shape [batch, tokens, hidden]")
return self.projection(features)
class VisionAdapter(nn.Module):
"""Image -> encoder -> projector, without modifying the CORTEX core."""
def __init__(self, encoder: VisionEncoder, projector: VisionProjector,
processor: ImageProcessor | None = None) -> None:
super().__init__()
self.encoder = encoder
self.projector = projector
self.processor = processor
def encode_images(self, images: Any) -> Tensor:
model_inputs = self.processor(images, return_tensors="pt") if self.processor else images
# Same BatchFeature caveat as native.py: HF processors return a UserDict, which is
# not a dict subclass, so it must be tested as a Mapping before unpacking by keyword.
if isinstance(model_inputs, Mapping):
features = self.encoder(**model_inputs)
else:
features = self.encoder(model_inputs)
if hasattr(features, "last_hidden_state"):
features = features.last_hidden_state
if not isinstance(features, Tensor):
raise TypeError("vision encoder must return a Tensor or an object with last_hidden_state")
return self.projector(features)
def prepare_multimodal_input(self, images: Any, input_ids: Tensor | None = None,
attention_mask: Tensor | None = None) -> MultimodalInput:
projected = self.encode_images(images)
return MultimodalInput(input_ids, projected, attention_mask,
metadata={"vision_status": "experimental / in development"})
class MultimodalFusion(Protocol):
"""Future contract for inserting projected image tokens into the LM context."""
def fuse(self, text_embeddings: Tensor, image_embeddings: Tensor) -> Tensor: ...
class PrefixFusion(nn.Module):
"""Reference fusion strategy: prepend image tokens to text embeddings."""
def fuse(self, text_embeddings: Tensor, image_embeddings: Tensor) -> Tensor:
if text_embeddings.ndim != 3 or image_embeddings.ndim != 3:
raise ValueError("text and image embeddings must be rank-3 tensors")
if text_embeddings.shape[0] != image_embeddings.shape[0]:
raise ValueError("text and image batches must match")
return torch.cat((image_embeddings, text_embeddings), dim=1)
def build_vision_adapter(config: VisionConfig, encoder: VisionEncoder,
processor: ImageProcessor | None = None) -> VisionAdapter:
"""Build the adaptation layer around an externally supplied encoder."""
return VisionAdapter(encoder, VisionProjector(config), processor)
|