from __future__ import annotations import hashlib from abc import ABC, abstractmethod import torch from torch import nn class FrozenBackbone(nn.Module, ABC): vision_dim: int text_dim: int @abstractmethod def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: ... @abstractmethod def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: ... def train(self, mode: bool = True) -> "FrozenBackbone": # The extractor trains around the backbone; pretrained backbone behavior stays fixed. return super().train(False) class TinyFrozenBackbone(FrozenBackbone): """Offline backbone used to exercise the complete pipeline in CPU CI.""" def __init__(self, hidden_dim: int = 64, patch_size: int = 16): super().__init__() generator = torch.Generator().manual_seed(1234) self.patch = nn.Conv2d(3, hidden_dim, patch_size, stride=patch_size, bias=False) self.byte_embedding = nn.Embedding(257, hidden_dim, padding_idx=0) with torch.no_grad(): self.patch.weight.normal_(generator=generator, std=0.04) self.byte_embedding.weight.normal_(generator=generator, std=0.04) self.byte_embedding.weight[0].zero_() self.requires_grad_(False) self.vision_dim = hidden_dim self.text_dim = hidden_dim def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: with torch.no_grad(): tokens = self.patch(images).flatten(2).transpose(1, 2) mask = torch.ones(tokens.shape[:2], dtype=torch.bool, device=tokens.device) return tokens, mask def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: encoded = [] for text in texts: # Hash whole whitespace-delimited tokens so the offline representation preserves words. ids = [] for token in text.lower().replace(".", "").split(): digest = hashlib.blake2b(token.encode("utf-8"), digest_size=2).digest() ids.append((int.from_bytes(digest, "little") % 256) + 1) encoded.append(ids or [1]) length = max(len(row) for row in encoded) input_ids = torch.zeros( len(encoded), length, dtype=torch.long, device=self.patch.weight.device ) mask = torch.zeros_like(input_ids, dtype=torch.bool) for index, row in enumerate(encoded): input_ids[index, : len(row)] = torch.tensor(row, device=input_ids.device) mask[index, : len(row)] = True with torch.no_grad(): tokens = self.byte_embedding(input_ids) return tokens, mask class CLIPFrozenBackbone(FrozenBackbone): """Frozen CLIP token encoder for the first real GPU experiment.""" def __init__(self, model_name: str): super().__init__() try: from transformers import AutoTokenizer, CLIPModel except ImportError as error: raise RuntimeError("Install the 'hf' extra to use the CLIP backbone.") from error # Keep training and the post-training benchmark on the same cached artifact. # Otherwise Hugging Face may fetch both the legacy .bin and safetensors weights. self.model = CLIPModel.from_pretrained(model_name, use_safetensors=False) self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model.requires_grad_(False).eval() self.vision_dim = self.model.config.vision_config.hidden_size self.text_dim = self.model.config.text_config.hidden_size mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]) std = torch.tensor([0.26862954, 0.26130258, 0.27577711]) self.register_buffer("image_mean", mean.view(1, 3, 1, 1), persistent=False) self.register_buffer("image_std", std.view(1, 3, 1, 1), persistent=False) def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: pixels = (images - self.image_mean) / self.image_std with torch.no_grad(): tokens = self.model.vision_model(pixel_values=pixels).last_hidden_state mask = torch.ones(tokens.shape[:2], dtype=torch.bool, device=tokens.device) return tokens, mask def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: values = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt") values = {key: value.to(self.image_mean.device) for key, value in values.items()} with torch.no_grad(): tokens = self.model.text_model(**values).last_hidden_state return tokens, values["attention_mask"].bool() class SigLIP2FrozenBackbone(FrozenBackbone): """Billion-scale paired image/text encoder used for formal experiments.""" def __init__(self, model_name: str): super().__init__() try: from transformers import AutoModel, AutoProcessor except ImportError as error: raise RuntimeError("Install the 'hf' extra to use SigLIP2.") from error dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 self.model = AutoModel.from_pretrained( model_name, torch_dtype=dtype, low_cpu_mem_usage=True, ) self.processor = AutoProcessor.from_pretrained(model_name) self.model.requires_grad_(False).eval() self.vision_dim = self.model.config.vision_config.hidden_size self.text_dim = self.model.config.text_config.hidden_size def _to_model_device(self, values: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: device = next(self.model.parameters()).device dtype = next(self.model.parameters()).dtype moved = {} for key, value in values.items(): moved[key] = ( value.to(device=device, dtype=dtype) if value.is_floating_point() else value.to(device=device) ) return moved def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: values = self.processor(images=images.detach().cpu(), return_tensors="pt") allowed = {"pixel_values", "pixel_attention_mask", "spatial_shapes"} vision_values = self._to_model_device( {key: value for key, value in values.items() if key in allowed} ) with torch.no_grad(): tokens = self.model.vision_model(**vision_values).last_hidden_state mask = torch.ones(tokens.shape[:2], dtype=torch.bool, device=tokens.device) return tokens, mask def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: max_length = self.model.config.text_config.max_position_embeddings values = self.processor( text=texts, padding="max_length", truncation=True, max_length=max_length, return_tensors="pt", ) text_values = self._to_model_device( {key: value for key, value in values.items() if key in {"input_ids", "attention_mask"}} ) with torch.no_grad(): tokens = self.model.text_model(**text_values).last_hidden_state return tokens, text_values["attention_mask"].bool() def build_backbone(name: str, model_name: str) -> FrozenBackbone: if name == "tiny": return TinyFrozenBackbone() if name == "clip": return CLIPFrozenBackbone(model_name) if name == "siglip2": return SigLIP2FrozenBackbone(model_name) raise ValueError(f"Unknown backbone: {name}")