from dataclasses import dataclass from typing import Dict, Optional, Sequence import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer try: from peft import LoraConfig, get_peft_model except ImportError: # PEFT is optional unless LoRA is enabled. LoraConfig = None get_peft_model = None from src.alignment.sinkhorn_ot import SinkhornOTConfig, SinkhornOTRouter from src.models.structural_vqa import StructuralVQAPrototype from src.models.vision_encoder import StructuralVisionEncoder, build_structural_vision_encoder @dataclass class StructuralGenerativeVQAConfig: llm_name_or_path: str = "distilgpt2" vision_d_model: int = 768 freeze_llm: bool = True ot_loss_weight: float = 0.05 use_ot: bool = True use_ot_fusion: bool = True ot_fusion_mode: str = "prefix" ot_fusion_dropout: float = 0.10 use_lora: bool = False lora_r: int = 8 lora_alpha: int = 16 lora_dropout: float = 0.05 lora_target_modules: str = "q_proj,v_proj" use_prior_as_ot_target: bool = True prior_ot_global_mass: float = 0.05 use_topological_loss: bool = True use_prior_align_loss: bool = True use_global_topo_loss: bool = True use_patch_topo_loss: bool = False prior_loss_weight: float = 0.05 global_topo_loss_weight: float = 0.01 patch_topo_loss_weight: float = 0.005 topo_feature_dim: int = 12 global_feature_dim: int = 8 max_question_length: int = 128 max_answer_length: int = 128 dropout: float = 0.10 def __post_init__(self): if not 0.0 <= self.prior_ot_global_mass < 1.0: raise ValueError("prior_ot_global_mass phải nằm trong [0, 1)") if self.ot_fusion_mode not in {"none", "prefix"}: raise ValueError("ot_fusion_mode hiện chỉ hỗ trợ 'none' hoặc 'prefix'") if self.use_lora and (LoraConfig is None or get_peft_model is None): raise ImportError("use_lora=True nhưng chưa cài package 'peft'. Hãy chạy: pip install peft") if self.lora_r <= 0: raise ValueError("lora_r phải > 0") class VisualToLLMProjector(nn.Module): def __init__(self, vision_dim: int, llm_dim: int, dropout: float = 0.10): super().__init__() self.net = nn.Sequential( nn.LayerNorm(vision_dim), nn.Linear(vision_dim, llm_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(llm_dim, llm_dim), ) def forward(self, visual_context: torch.Tensor) -> torch.Tensor: return self.net(visual_context) class StructuralGenerativeVQA(nn.Module): """ Generative VQA model: StructuralVisionEncoder -> visual prefix embeddings -> causal LM. Nếu `use_ot_fusion=True`, Sinkhorn OT aligned visual features cũng được project thành prefix tokens và đưa trực tiếp vào LLM input. """ def __init__( self, config: StructuralGenerativeVQAConfig, vision_encoder: StructuralVisionEncoder, ): super().__init__() self.config = config self.vision_encoder = vision_encoder self.tokenizer = AutoTokenizer.from_pretrained(config.llm_name_or_path) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token self.llm = AutoModelForCausalLM.from_pretrained(config.llm_name_or_path) self.llm.config.pad_token_id = self.tokenizer.pad_token_id llm_dim = self.llm.get_input_embeddings().embedding_dim if config.freeze_llm: for param in self.llm.parameters(): param.requires_grad = False if config.use_lora: target_modules = [m.strip() for m in config.lora_target_modules.split(",") if m.strip()] lora_config = LoraConfig( r=config.lora_r, lora_alpha=config.lora_alpha, lora_dropout=config.lora_dropout, target_modules=target_modules, bias="none", task_type="CAUSAL_LM", ) self.llm = get_peft_model(self.llm, lora_config) self.llm.print_trainable_parameters() self.visual_projector = VisualToLLMProjector( vision_dim=config.vision_d_model, llm_dim=llm_dim, dropout=config.dropout, ) self.ot_visual_projector = VisualToLLMProjector( vision_dim=config.vision_d_model, llm_dim=llm_dim, dropout=config.ot_fusion_dropout, ) self.ot_text_projector = nn.Linear(llm_dim, config.vision_d_model) self.ot_router = SinkhornOTRouter(SinkhornOTConfig()) hidden = config.vision_d_model * 2 self.global_topo_head = nn.Sequential( nn.LayerNorm(config.vision_d_model), nn.Linear(config.vision_d_model, hidden), nn.GELU(), nn.Dropout(config.dropout), nn.Linear(hidden, config.global_feature_dim), ) self.patch_topo_head = nn.Sequential( nn.LayerNorm(config.vision_d_model), nn.Linear(config.vision_d_model, hidden), nn.GELU(), nn.Dropout(config.dropout), nn.Linear(hidden, config.topo_feature_dim), ) def forward( self, image: torch.Tensor, prior_mask: torch.Tensor, topo_features: torch.Tensor, global_features: torch.Tensor, question_text: Sequence[str], answer_text: Optional[Sequence[str]] = None, return_diagnostics: bool = True, ) -> Dict[str, torch.Tensor]: device = image.device prior_mask = prior_mask.to(device=device, dtype=image.dtype) topo_features = topo_features.to(device=device, dtype=image.dtype) global_features = global_features.to(device=device, dtype=image.dtype) vision_out = self.vision_encoder( image=image, prior_mask=prior_mask, topo_features=topo_features, global_features=global_features, return_diagnostics=True, ) visual_context = vision_out["visual_context"] prompt_text = [self._format_prompt(q, use_chat_template=True) for q in question_text] prompt = self.tokenizer( prompt_text, padding=True, truncation=True, max_length=self.config.max_question_length, return_tensors="pt", ).to(device) prompt_embeds = self.llm.get_input_embeddings()(prompt.input_ids) llm_dtype = self.llm.get_input_embeddings().weight.dtype ot_out = self._compute_ot( visual_context=visual_context, prompt_embeds=prompt_embeds, prompt_attention_mask=prompt.attention_mask, prior_mask=prior_mask, ) ot_cost = ot_out["ot_cost"] if ot_out is not None else visual_context.new_tensor(0.0) visual_prefix = self.visual_projector(visual_context).to(dtype=llm_dtype) prefix_parts = [visual_prefix] prefix_masks = [torch.ones(visual_prefix.shape[:2], device=device, dtype=prompt.attention_mask.dtype)] ot_aligned_prefix = None if self._use_ot_prefix(ot_out): ot_aligned_prefix = self.ot_visual_projector(ot_out["aligned_visual"]).to(dtype=llm_dtype) prefix_parts.append(ot_aligned_prefix) prefix_masks.append(prompt.attention_mask) prefix_embeds = torch.cat(prefix_parts, dim=1) prefix_attention = torch.cat(prefix_masks, dim=1) full_embeds = torch.cat([prefix_embeds, prompt_embeds], dim=1) full_attention = torch.cat([prefix_attention, prompt.attention_mask], dim=1) labels = None if answer_text is not None: answer = self.tokenizer( [str(a) + self.tokenizer.eos_token for a in answer_text], padding=True, truncation=True, max_length=self.config.max_answer_length, return_tensors="pt", ).to(device) answer_embeds = self.llm.get_input_embeddings()(answer.input_ids) full_embeds = torch.cat([full_embeds, answer_embeds], dim=1) full_attention = torch.cat([full_attention, answer.attention_mask], dim=1) ignore_prefix = torch.full( (image.shape[0], prefix_embeds.shape[1] + prompt.input_ids.shape[1]), -100, device=device, dtype=torch.long, ) answer_labels = answer.input_ids.masked_fill(answer.attention_mask == 0, -100) labels = torch.cat([ignore_prefix, answer_labels], dim=1) llm_out = self.llm( inputs_embeds=full_embeds, attention_mask=full_attention, labels=labels, return_dict=True, ) lm_loss = llm_out.loss if labels is not None else visual_context.new_tensor(0.0) topo_out = self.compute_topological_losses( vision_out=vision_out, ot_out=ot_out, prior_mask=prior_mask, topo_features=topo_features, global_features=global_features, attention_mask=prompt.attention_mask, ) loss = lm_loss + self.config.ot_loss_weight * ot_cost + topo_out["topological_loss"] if labels is not None else None if not return_diagnostics: return {"loss": loss, "logits": llm_out.logits} out = { "loss": loss, "lm_loss": lm_loss, "ot_cost": ot_cost, "logits": llm_out.logits, "vision_out": vision_out, "visual_prefix_shape": torch.tensor(visual_prefix.shape, device=device), "ot_fusion_enabled": torch.tensor(self._use_ot_prefix(ot_out), device=device), **topo_out, } if ot_aligned_prefix is not None: out["ot_aligned_prefix_shape"] = torch.tensor(ot_aligned_prefix.shape, device=device) if ot_out is not None: out["ot_out"] = ot_out return out @torch.no_grad() def generate( self, image: torch.Tensor, prior_mask: torch.Tensor, topo_features: torch.Tensor, global_features: torch.Tensor, question_text: Sequence[str], max_new_tokens: int = 64, ) -> Sequence[str]: self.eval() device = image.device prior_mask = prior_mask.to(device=device, dtype=image.dtype) topo_features = topo_features.to(device=device, dtype=image.dtype) global_features = global_features.to(device=device, dtype=image.dtype) vision_out = self.vision_encoder( image=image, prior_mask=prior_mask, topo_features=topo_features, global_features=global_features, return_diagnostics=True, ) visual_context = vision_out["visual_context"] prompt = self.tokenizer( [self._format_prompt(q, use_chat_template=True) for q in question_text], padding=True, truncation=True, max_length=self.config.max_question_length, return_tensors="pt", ).to(device) prompt_embeds = self.llm.get_input_embeddings()(prompt.input_ids) llm_dtype = self.llm.get_input_embeddings().weight.dtype ot_out = self._compute_ot( visual_context=visual_context, prompt_embeds=prompt_embeds, prompt_attention_mask=prompt.attention_mask, prior_mask=prior_mask, ) visual_prefix = self.visual_projector(visual_context).to(dtype=llm_dtype) prefix_parts = [visual_prefix] prefix_masks = [torch.ones(visual_prefix.shape[:2], device=device, dtype=prompt.attention_mask.dtype)] if self._use_ot_prefix(ot_out): prefix_parts.append(self.ot_visual_projector(ot_out["aligned_visual"]).to(dtype=llm_dtype)) prefix_masks.append(prompt.attention_mask) prefix_embeds = torch.cat(prefix_parts, dim=1) prefix_attention = torch.cat(prefix_masks, dim=1) inputs_embeds = torch.cat([prefix_embeds, prompt_embeds], dim=1) attention_mask = torch.cat([prefix_attention, prompt.attention_mask], dim=1) generated = self.llm.generate( inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=max_new_tokens, min_new_tokens=2, do_sample=False, num_beams=1, repetition_penalty=1.15, no_repeat_ngram_size=3, pad_token_id=self.tokenizer.pad_token_id, eos_token_id=self._generation_eos_token_ids(), ) decoded = self.tokenizer.batch_decode(generated, skip_special_tokens=True) return [text.strip() for text in decoded] def _compute_ot( self, visual_context: torch.Tensor, prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor, prior_mask: torch.Tensor, ) -> Optional[Dict[str, torch.Tensor]]: if not self.config.use_ot: return None prompt_token_embeds = prompt_embeds.detach() if self.config.freeze_llm else prompt_embeds prompt_token_embeds = prompt_token_embeds.to(dtype=self.ot_text_projector.weight.dtype) text_tokens = self.ot_text_projector(prompt_token_embeds) visual_target_mass = None if self.config.use_prior_as_ot_target: visual_target_mass = StructuralVQAPrototype.build_prior_visual_target_mass( prior_mask=prior_mask, num_visual_tokens=visual_context.shape[1], global_mass=self.config.prior_ot_global_mass, ) return self.ot_router( visual_tokens=visual_context, text_tokens=text_tokens, text_mask=prompt_attention_mask, visual_target_mass=visual_target_mass, return_diagnostics=True, ) def _use_ot_prefix(self, ot_out: Optional[Dict[str, torch.Tensor]]) -> bool: return ( self.config.use_ot and self.config.use_ot_fusion and self.config.ot_fusion_mode == "prefix" and ot_out is not None ) def compute_topological_losses( self, vision_out: Dict[str, torch.Tensor], ot_out: Optional[Dict[str, torch.Tensor]], prior_mask: torch.Tensor, topo_features: torch.Tensor, global_features: torch.Tensor, attention_mask: torch.Tensor, ) -> Dict[str, torch.Tensor]: zero = prior_mask.new_tensor(0.0) prior_align_loss = zero global_topo_loss = zero patch_topo_loss = zero if self.config.use_topological_loss: if self.config.use_prior_align_loss and ot_out is not None: prior_align_loss = StructuralVQAPrototype.compute_prior_alignment_loss( transport_plan=ot_out["transport_plan"], prior_mask=prior_mask, attention_mask=attention_mask, ) if self.config.use_global_topo_loss and vision_out.get("global_token") is not None: global_pred = self.global_topo_head(vision_out["global_token"].squeeze(1)) global_topo_loss = F.mse_loss(global_pred, global_features.float()) if self.config.use_patch_topo_loss: patch_pred = self.patch_topo_head(vision_out["fused_tokens"]) patch_target = topo_features.flatten(start_dim=1, end_dim=2).float() patch_topo_loss = F.mse_loss(patch_pred, patch_target) effective_prior_loss_weight = 0.0 if self.config.use_prior_as_ot_target else self.config.prior_loss_weight topological_loss = ( effective_prior_loss_weight * prior_align_loss + self.config.global_topo_loss_weight * global_topo_loss + self.config.patch_topo_loss_weight * patch_topo_loss ) return { "prior_align_loss": prior_align_loss, "global_topo_loss": global_topo_loss, "patch_topo_loss": patch_topo_loss, "topological_loss": topological_loss, "effective_prior_loss_weight": prior_mask.new_tensor(effective_prior_loss_weight), } def _generation_eos_token_ids(self): eos_ids = [] if self.tokenizer.eos_token_id is not None: eos_ids.append(self.tokenizer.eos_token_id) im_end_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>") if isinstance(im_end_id, int) and im_end_id >= 0 and im_end_id not in eos_ids: eos_ids.append(im_end_id) return eos_ids or self.tokenizer.eos_token_id def _format_prompt(self, question: str, use_chat_template: bool = True) -> str: question = str(question or "").replace("", "").strip() instruction = ( "You are a medical visual question answering assistant. " "Answer using only the image evidence. " "Give one concise sentence and do not repeat the question." ) if use_chat_template and getattr(self.tokenizer, "chat_template", None): messages = [ {"role": "system", "content": instruction}, {"role": "user", "content": question}, ] return self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) return f"System: {instruction}\nUser: {question}\nAssistant:" def build_structural_generative_vqa( llm_name_or_path: str = "distilgpt2", vision_pretrained: bool = False, vision_backend: str = "timm", freeze_vision_backbone: bool = True, freeze_llm: bool = True, **kwargs, ) -> StructuralGenerativeVQA: config = StructuralGenerativeVQAConfig( llm_name_or_path=llm_name_or_path, freeze_llm=freeze_llm, **kwargs, ) vision_encoder = build_structural_vision_encoder( d_model=config.vision_d_model, pretrained=vision_pretrained, backend=vision_backend, freeze_backbone=freeze_vision_backbone, ) return StructuralGenerativeVQA(config=config, vision_encoder=vision_encoder)