from typing import Dict, List, Tuple import timm import torch from torch import nn from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from .local_store import runtime_text_model_path def resolve_dtype(dtype_name: str) -> torch.dtype: if dtype_name == "float16": return torch.float16 if dtype_name == "float32": return torch.float32 raise ValueError(f"Unsupported dtype: {dtype_name}") def infer_hidden_size(llama: AutoModelForCausalLM) -> int: if hasattr(llama.config, "hidden_size"): return int(llama.config.hidden_size) if hasattr(llama.config, "n_embd"): return int(llama.config.n_embd) raise ValueError("Could not infer hidden size from the language model config") class ImageEmbeddingModel(nn.Module): def __init__(self, model_name: str, embedding_dim: int) -> None: super().__init__() self.embedding = timm.create_model(model_name, pretrained=True, num_classes=embedding_dim) def forward(self, images: torch.Tensor) -> torch.Tensor: return self.embedding(images) class LlamaTextProjector(nn.Module): def __init__(self, llama: AutoModelForCausalLM, embedding_dim: int) -> None: super().__init__() self.text_base = getattr(llama, "model", llama) hidden_size = infer_hidden_size(llama) self.projection = nn.Sequential( nn.Linear(hidden_size, 512, bias=False), nn.Dropout(0.5), nn.GELU(), nn.LayerNorm(512), nn.Linear(512, embedding_dim, bias=False), ) def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: outputs = self.text_base(input_ids=input_ids, attention_mask=attention_mask) hidden = outputs.last_hidden_state mask = attention_mask.unsqueeze(-1).to(hidden.dtype) pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) pooled = pooled.to(self.projection[0].weight.dtype) return self.projection(pooled) class FusionHead(nn.Module): def __init__(self, fusion: str, embedding_dim: int) -> None: super().__init__() self.fusion = fusion if fusion == "concat": fused_dim = embedding_dim * 2 elif fusion == "element": fused_dim = embedding_dim elif fusion == "avgpool": self.w_image = nn.Linear(embedding_dim, 64, bias=False) self.w_text = nn.Linear(embedding_dim, 64, bias=False) self.pool = nn.AvgPool1d(kernel_size=4) fused_dim = 16 elif fusion == "gated": self.u_text = nn.Linear(embedding_dim, 64, bias=False) self.v_image = nn.Linear(embedding_dim, 64, bias=False) self.w_fused = nn.Linear(64, 8, bias=False) self.sigmoid = nn.Sigmoid() fused_dim = 8 else: raise ValueError(f"Unsupported fusion method: {fusion}") hidden_dim = max(fused_dim // 2, 1) self.initial = nn.Linear(fused_dim, hidden_dim, bias=True) self.final = nn.Linear(hidden_dim, 1, bias=True) def forward(self, text_features: torch.Tensor, image_features: torch.Tensor) -> torch.Tensor: target_dtype = self.initial.weight.dtype text_features = text_features.to(target_dtype) image_features = image_features.to(target_dtype) if self.fusion == "concat": fused = torch.cat([text_features, image_features], dim=-1) elif self.fusion == "element": fused = text_features * image_features elif self.fusion == "avgpool": text_encoded = self.w_text(text_features) image_encoded = self.w_image(image_features) fused = self.pool((text_encoded * image_encoded).unsqueeze(1)).squeeze(1) else: gate = self.sigmoid(self.u_text(text_features)) image_encoded = self.v_image(image_features) fused = self.w_fused(gate * image_encoded) hidden = self.initial(fused) return self.final(hidden).squeeze(-1) class FuseMD(nn.Module): def __init__( self, text_encoder: LlamaTextProjector, image_encoder: ImageEmbeddingModel, fusion_head: FusionHead, ) -> None: super().__init__() self.text_encoder = text_encoder self.image_encoder = image_encoder self.fusion_head = fusion_head def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor, images: torch.Tensor, ) -> torch.Tensor: text_features = self.text_encoder(input_ids, attention_mask) image_features = self.image_encoder(images) return self.fusion_head(text_features, image_features) def load_tokenizer_and_llama( text_model_name: str, use_8bit: bool, dtype: torch.dtype, device: torch.device, ) -> Tuple[AutoTokenizer, AutoModelForCausalLM]: resolved_model_path = runtime_text_model_path(text_model_name) tokenizer = AutoTokenizer.from_pretrained(resolved_model_path, local_files_only=True) if tokenizer.pad_token is None: if tokenizer.eos_token is not None: tokenizer.pad_token = tokenizer.eos_token else: tokenizer.add_special_tokens({"pad_token": "[PAD]"}) quantization_config = None if use_8bit and device.type == "cuda": quantization_config = BitsAndBytesConfig(load_in_8bit=True) model_dtype = dtype if device.type == "cuda" else torch.float32 llama = AutoModelForCausalLM.from_pretrained( resolved_model_path, low_cpu_mem_usage=True, local_files_only=True, quantization_config=quantization_config, torch_dtype=model_dtype, ) if len(tokenizer) != llama.get_input_embeddings().num_embeddings: llama.resize_token_embeddings(len(tokenizer), mean_resizing=False) return tokenizer, llama def freeze_base_models(model: FuseMD) -> None: for param in model.text_encoder.text_base.parameters(): param.requires_grad = False for param in model.text_encoder.projection.parameters(): param.requires_grad = True for param in model.image_encoder.embedding.parameters(): param.requires_grad = False if hasattr(model.image_encoder.embedding, "head"): for param in model.image_encoder.embedding.head.parameters(): param.requires_grad = True for param in model.fusion_head.parameters(): param.requires_grad = True def trainable_parameters(model: FuseMD) -> List[nn.Parameter]: return [param for param in model.parameters() if param.requires_grad] def get_trainable_checkpoint(model: FuseMD) -> Dict[str, object]: vit_state = model.image_encoder.state_dict() vit_head_state = {key: value.cpu() for key, value in vit_state.items() if key.startswith("embedding.head")} return { "text_projection_state_dict": { key: value.cpu() for key, value in model.text_encoder.projection.state_dict().items() }, "vit_head_state_dict": vit_head_state, "fusion_head_state_dict": { key: value.cpu() for key, value in model.fusion_head.state_dict().items() }, } def load_trainable_checkpoint(model: FuseMD, checkpoint: Dict[str, object]) -> None: text_projection_state = checkpoint.get("text_projection_state_dict") if text_projection_state is None: text_projection_state = checkpoint.get("text_clf_head_state_dict") if text_projection_state is None: raise KeyError("Checkpoint is missing both 'text_projection_state_dict' and 'text_clf_head_state_dict'.") model.text_encoder.projection.load_state_dict(text_projection_state) vit_current = model.image_encoder.state_dict() vit_current.update(checkpoint["vit_head_state_dict"]) model.image_encoder.load_state_dict(vit_current, strict=False) model.fusion_head.load_state_dict(checkpoint["fusion_head_state_dict"]) def build_model( llama: AutoModelForCausalLM, image_model_name: str, fusion_method: str, embedding_dim: int, ) -> FuseMD: text_encoder = LlamaTextProjector(llama=llama, embedding_dim=embedding_dim) image_encoder = ImageEmbeddingModel(image_model_name, embedding_dim=embedding_dim) fusion_head = FusionHead(fusion=fusion_method, embedding_dim=embedding_dim) model = FuseMD(text_encoder=text_encoder, image_encoder=image_encoder, fusion_head=fusion_head) freeze_base_models(model) return model