from dataclasses import dataclass from typing import Dict, Optional, Sequence, Tuple import torch import torch.nn as nn import torch.nn.functional as F from src.alignment.sinkhorn_ot import SinkhornOTConfig, SinkhornOTRouter from src.models.text_encoder import BasicVQATokenizer, SimpleTextEncoder, SimpleTextEncoderConfig from src.models.vision_encoder import StructuralVisionEncoder, build_structural_vision_encoder class AnswerVocabulary: """ Map answer text <-> class id cho prototype VQA classification. Đây là baseline classification, không phải generative decoder. """ def __init__(self, answers: Optional[Sequence[str]] = None, min_freq: int = 1, max_answers: Optional[int] = None): self.answer_to_id: Dict[str, int] = {} self.id_to_answer: Dict[int, str] = {} if answers is not None: self.build(answers, min_freq=min_freq, max_answers=max_answers) def __len__(self) -> int: return len(self.answer_to_id) def build(self, answers: Sequence[str], min_freq: int = 1, max_answers: Optional[int] = None) -> None: from collections import Counter counter = Counter(self.normalize(answer) for answer in answers) self.answer_to_id.clear() self.id_to_answer.clear() for answer, freq in counter.most_common(max_answers): if freq < min_freq: continue idx = len(self.answer_to_id) self.answer_to_id[answer] = idx self.id_to_answer[idx] = answer @staticmethod def normalize(answer: str) -> str: return " ".join(str(answer or "").strip().lower().split()) def encode(self, answer: str) -> int: key = self.normalize(answer) if key not in self.answer_to_id: raise KeyError(f"Answer ngoài vocabulary: {answer}") return self.answer_to_id[key] def batch_encode(self, answers: Sequence[str]) -> torch.Tensor: return torch.tensor([self.encode(answer) for answer in answers], dtype=torch.long) def decode(self, answer_id: int) -> str: return self.id_to_answer[int(answer_id)] @dataclass class StructuralVQAConfig: d_model: int = 768 num_answers: int = 1000 text_max_length: int = 64 dropout: float = 0.10 ot_loss_weight: float = 0.01 use_ot: bool = True 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 use_prior_as_ot_target: bool = True prior_ot_global_mass: float = 0.05 classifier_hidden_dim: Optional[int] = None def __post_init__(self): if self.d_model <= 0: raise ValueError("d_model phải > 0") if self.num_answers <= 0: raise ValueError("num_answers phải > 0") if self.topo_feature_dim <= 0: raise ValueError("topo_feature_dim phải > 0") if self.global_feature_dim <= 0: raise ValueError("global_feature_dim phải > 0") 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.classifier_hidden_dim is None: self.classifier_hidden_dim = self.d_model * 2 class StructuralVQAPrototype(nn.Module): """ Prototype VQA classification model: image + structural features -> StructuralVisionEncoder -> visual_context question_text/input_ids -> SimpleTextEncoder -> text_tokens text_tokens + visual_context -> SinkhornOTRouter -> aligned_visual pooled multimodal representation -> answer logits Topological losses optional: - prior alignment: OT transport should align with lesion prior. - global topology reconstruction: global token should reconstruct global features. - patch topology reconstruction: fused tokens should reconstruct patch topo features. """ def __init__( self, config: StructuralVQAConfig, vision_encoder: StructuralVisionEncoder, text_encoder: SimpleTextEncoder, tokenizer: Optional[BasicVQATokenizer] = None, ot_router: Optional[SinkhornOTRouter] = None, ): super().__init__() self.config = config self.vision_encoder = vision_encoder self.text_encoder = text_encoder self.tokenizer = tokenizer self.ot_router = ot_router or SinkhornOTRouter(SinkhornOTConfig()) fusion_dim = config.d_model * 4 self.classifier = nn.Sequential( nn.LayerNorm(fusion_dim), nn.Linear(fusion_dim, config.classifier_hidden_dim), nn.GELU(), nn.Dropout(config.dropout), nn.Linear(config.classifier_hidden_dim, config.num_answers), ) self.global_topo_head = nn.Sequential( nn.LayerNorm(config.d_model), nn.Linear(config.d_model, config.classifier_hidden_dim), nn.GELU(), nn.Dropout(config.dropout), nn.Linear(config.classifier_hidden_dim, config.global_feature_dim), ) self.patch_topo_head = nn.Sequential( nn.LayerNorm(config.d_model), nn.Linear(config.d_model, config.classifier_hidden_dim), nn.GELU(), nn.Dropout(config.dropout), nn.Linear(config.classifier_hidden_dim, config.topo_feature_dim), ) def forward( self, image: torch.Tensor, prior_mask: torch.Tensor, topo_features: torch.Tensor, global_features: torch.Tensor, question_text: Optional[Sequence[str]] = None, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, answer_ids: Optional[torch.Tensor] = None, return_diagnostics: bool = True, ): device = image.device if input_ids is None: if question_text is None: raise ValueError("Cần truyền question_text hoặc input_ids") if self.tokenizer is None: raise ValueError("question_text được truyền nhưng model.tokenizer=None") tokenized = self.tokenizer.batch_encode( question_text, max_length=self.config.text_max_length, ) input_ids = tokenized["input_ids"] attention_mask = tokenized["attention_mask"] input_ids = input_ids.to(device=device) if attention_mask is not None: attention_mask = attention_mask.to(device=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"] text_out = self.text_encoder( input_ids=input_ids, attention_mask=attention_mask, return_dict=True, ) text_tokens = text_out["text_tokens"] pooled_text = text_out["pooled_text"] attention_mask = text_out["attention_mask"] if self.config.use_ot: visual_target_mass = None if self.config.use_prior_as_ot_target: visual_target_mass = self.build_prior_visual_target_mass( prior_mask=prior_mask, num_visual_tokens=visual_context.shape[1], global_mass=self.config.prior_ot_global_mass, ) ot_out = self.ot_router( visual_tokens=visual_context, text_tokens=text_tokens, text_mask=attention_mask, visual_target_mass=visual_target_mass, return_diagnostics=True, ) aligned_visual_tokens = ot_out["aligned_visual"] pooled_aligned_visual = SimpleTextEncoder.masked_mean_pool( aligned_visual_tokens, attention_mask, ) ot_cost = ot_out["ot_cost"] else: ot_out = None pooled_aligned_visual = visual_context.mean(dim=1) ot_cost = visual_context.new_tensor(0.0) pooled_visual = visual_context.mean(dim=1) multimodal = torch.cat( [ pooled_text, pooled_visual, pooled_aligned_visual, pooled_text * pooled_aligned_visual, ], dim=-1, ) logits = self.classifier(multimodal) topo_loss_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=attention_mask, ) loss = None ce_loss = None if answer_ids is not None: answer_ids = answer_ids.to(device=device) ce_loss = F.cross_entropy(logits, answer_ids) loss = ce_loss + self.config.ot_loss_weight * ot_cost + topo_loss_out["topological_loss"] if not return_diagnostics: return logits output = { "logits": logits, "loss": loss, "ce_loss": ce_loss, "ot_cost": ot_cost, "visual_context": visual_context, "text_tokens": text_tokens, "pooled_text": pooled_text, "pooled_visual": pooled_visual, "pooled_aligned_visual": pooled_aligned_visual, "vision_out": vision_out, **topo_loss_out, } if ot_out is not None: output["ot_out"] = ot_out return output 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 global_topo_pred = None patch_topo_pred = None if self.config.use_topological_loss: if self.config.use_prior_align_loss and ot_out is not None: prior_align_loss = self.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: global_token = vision_out.get("global_token") if global_token is not None: global_topo_pred = self.global_topo_head(global_token.squeeze(1)) global_topo_loss = F.mse_loss(global_topo_pred, global_features.float()) if self.config.use_patch_topo_loss: fused_tokens = vision_out["fused_tokens"] patch_topo_pred = self.patch_topo_head(fused_tokens) topo_target = topo_features.flatten(start_dim=1, end_dim=2).float() patch_topo_loss = F.mse_loss(patch_topo_pred, topo_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), "global_topo_pred": global_topo_pred, "patch_topo_pred": patch_topo_pred, } @staticmethod def compute_prior_alignment_loss( transport_plan: torch.Tensor, prior_mask: torch.Tensor, attention_mask: torch.Tensor, eps: float = 1e-8, ) -> torch.Tensor: if transport_plan.shape[-1] < 2: return transport_plan.new_tensor(0.0) orig_dtype = transport_plan.dtype # Force float32 to prevent underflow in float16 autocast transport_plan = transport_plan.float() prior_mask = prior_mask.float() attention_mask = attention_mask.float() patch_transport = transport_plan[:, :, 1:] text_mask = attention_mask.to(device=transport_plan.device, dtype=transport_plan.dtype).unsqueeze(-1) patch_attention = (patch_transport * text_mask).sum(dim=1) prior = prior_mask.flatten(start_dim=1).to(device=transport_plan.device, dtype=transport_plan.dtype) if prior.shape[1] != patch_attention.shape[1]: raise ValueError( f"prior patches={prior.shape[1]} không khớp transport patches={patch_attention.shape[1]}" ) patch_attention = patch_attention / patch_attention.sum(dim=1, keepdim=True).clamp_min(eps) prior = prior / prior.sum(dim=1, keepdim=True).clamp_min(eps) cosine = F.cosine_similarity(patch_attention, prior, dim=1, eps=eps) loss = (1.0 - cosine).mean() return loss.to(dtype=orig_dtype) @staticmethod def build_prior_visual_target_mass( prior_mask: torch.Tensor, num_visual_tokens: int, global_mass: float = 0.05, eps: float = 1e-8, ) -> torch.Tensor: """ Tạo target marginal cho visual side của Sinkhorn từ lesion prior. visual_context có dạng: token 0 = global structural token token 1..196 = patch tokens Nếu Sinkhorn dùng target marginal uniform, tổng mass trên patch tokens bị ép gần đều nên prior_align_loss gần như không thể giảm. Hàm này đổi target marginal thành: global token mass = global_mass patch token mass = (1 - global_mass) * normalized prior_mask """ if num_visual_tokens < 2: raise ValueError("Cần ít nhất 1 global token + patch tokens để dùng prior target mass") orig_dtype = prior_mask.dtype # Force float32 to prevent underflow in float16 autocast prior_mask = prior_mask.float() prior = prior_mask.flatten(start_dim=1) expected_patches = num_visual_tokens - 1 if prior.shape[1] != expected_patches: raise ValueError( f"prior patches={prior.shape[1]} không khớp visual patch tokens={expected_patches}" ) prior = prior.clamp_min(0.0) fallback = torch.ones_like(prior) / max(expected_patches, 1) prior_sum = prior.sum(dim=1, keepdim=True) prior_dist = torch.where(prior_sum > eps, prior / prior_sum.clamp_min(eps), fallback) global_column = prior.new_full((prior.shape[0], 1), global_mass) patch_mass = (1.0 - global_mass) * prior_dist res = torch.cat([global_column, patch_mass], dim=1) return res.to(dtype=orig_dtype) @torch.no_grad() def predict(self, *args, **kwargs) -> torch.Tensor: self.eval() out = self.forward(*args, **kwargs) logits = out["logits"] if isinstance(out, dict) else out return logits.argmax(dim=-1) def build_structural_vqa_prototype( questions: Sequence[str], answers: Sequence[str], d_model: int = 768, num_answers: Optional[int] = None, text_max_length: int = 64, vision_pretrained: bool = False, vision_backend: str = "timm", freeze_vision_backbone: bool = True, max_text_vocab_size: int = 30000, max_answer_vocab_size: Optional[int] = None, ot_loss_weight: float = 0.01, 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, use_prior_as_ot_target: bool = True, prior_ot_global_mass: float = 0.05, ) -> Tuple[StructuralVQAPrototype, BasicVQATokenizer, AnswerVocabulary]: tokenizer = BasicVQATokenizer.build_from_texts( questions, max_vocab_size=max_text_vocab_size, ) answer_vocab = AnswerVocabulary( answers, max_answers=max_answer_vocab_size, ) if num_answers is None: num_answers = len(answer_vocab) if num_answers <= 0: raise ValueError("Không build được answer vocabulary; num_answers=0") vision_encoder = build_structural_vision_encoder( d_model=d_model, pretrained=vision_pretrained, backend=vision_backend, freeze_backbone=freeze_vision_backbone, ) text_config = SimpleTextEncoderConfig( vocab_size=len(tokenizer), d_model=d_model, max_length=text_max_length, pad_token_id=tokenizer.pad_token_id, ) text_encoder = SimpleTextEncoder(text_config) config = StructuralVQAConfig( d_model=d_model, num_answers=num_answers, text_max_length=text_max_length, ot_loss_weight=ot_loss_weight, use_topological_loss=use_topological_loss, use_prior_align_loss=use_prior_align_loss, use_global_topo_loss=use_global_topo_loss, use_patch_topo_loss=use_patch_topo_loss, prior_loss_weight=prior_loss_weight, global_topo_loss_weight=global_topo_loss_weight, patch_topo_loss_weight=patch_topo_loss_weight, use_prior_as_ot_target=use_prior_as_ot_target, prior_ot_global_mass=prior_ot_global_mass, ) model = StructuralVQAPrototype( config=config, vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer, ) return model, tokenizer, answer_vocab