from transformers import AutoModelForCausalLM from datasets import load_dataset, Image as HfImage from transformers import AutoProcessor import torch import json import time from tqdm import tqdm import subprocess import platform import sys import os import random from PIL import Image import numpy as np SEED = 42 os.environ["PYTHONHASHSEED"] = str(SEED) os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED) torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True try: torch.set_float32_matmul_precision("high") except Exception: pass try: from itertools import batched except ImportError: def batched(iterable, n): import itertools iterator = iter(iterable) while batch := list(itertools.islice(iterator, n)): yield batch from evaluate import load bleu = load("bleu") rouge = load("rouge") meteor = load("meteor") ds = load_dataset("SimulaMet/Kvasir-VQA-x1")["test"] ds_shuffled = ds.shuffle(seed=42) # Shuffle with fixed seed for reproducibility val_dataset = ds_shuffled.select(range(1500)) # Select first 1500 after shuffle val_dataset = val_dataset.cast_column("image", HfImage()) predictions = [] # List to store predictions gpu_name = torch.cuda.get_device_name( 0) if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu" def get_mem(): return torch.cuda.memory_allocated(device) / \ (1024 ** 2) if torch.cuda.is_available() else 0 initial_mem = get_mem() # ✏️✏️--------EDIT SECTION 1: SUBMISISON DETAILS and MODEL LOADING --------✏️✏️# SUBMISSION_INFO = { "Participant_Names": "Minh Quang Nguyen", "Affiliations": "Independent", "Contact_emails": ["nmquang04072005@gmail.com"], "Team_Name": "Sweet&Sour", "Country": "Vietnam", "Notes_to_organizers": ''' Task 1 submission using Qwen2.5-3B-Instruct + QLoRA r16 + Curriculum-Gated full structural Deep TopoAdapter. Trained for 2 epochs on 30k with warmup/gate regularization then 1 epoch full-data continuation. Uses exact batched required template structure. ''' } # --- Deep TopoAdapter runtime injection and loading logic --- from pathlib import Path from typing import Any, Dict, Iterable, List, Optional REPO_DIR = Path(__file__).resolve().parent HF_REPO_ID = "minhquang47/temp" if str(REPO_DIR) not in sys.path: sys.path.insert(0, str(REPO_DIR)) def _ensure_full_repo_available() -> Path: if (REPO_DIR / "src").exists() and (REPO_DIR / "checkpoints").exists(): return REPO_DIR try: from huggingface_hub import snapshot_download full_repo = Path( snapshot_download( repo_id=HF_REPO_ID, repo_type="model", local_files_only=False, ) ) if str(full_repo) not in sys.path: sys.path.insert(0, str(full_repo)) return full_repo except Exception as exc: print("Warning: could not download full repo snapshot:", repr(exc)) return REPO_DIR RUNTIME_DIR = _ensure_full_repo_available() if str(RUNTIME_DIR) not in sys.path: sys.path.insert(0, str(RUNTIME_DIR)) try: from src.models.structural_vqa_generative import build_structural_generative_vqa except Exception as exc: print("Warning: model import failed lazily:", repr(exc)) build_structural_generative_vqa = None CHECKPOINT_CANDIDATES = [ RUNTIME_DIR / "checkpoints" / "final.pt", RUNTIME_DIR / "checkpoints" / "epoch_3.pt", RUNTIME_DIR / "checkpoints" / "epoch_4.pt", RUNTIME_DIR / "checkpoints" / "last.pt", RUNTIME_DIR / "model.pt", REPO_DIR / "checkpoints" / "final.pt", REPO_DIR / "model.pt", ] DEFAULT_IMAGE_SIZE = (224, 224) DEFAULT_GRID_SIZE = (14, 14) DEFAULT_TOPO_FEATURE_DIM = 12 DEFAULT_GLOBAL_FEATURE_DIM = 8 class TopoAdapter(torch.nn.Module): def __init__(self, hidden_dim: int, topo_dim: int, bottleneck_dim: int = 32): super().__init__() self.topo_projector = torch.nn.Sequential( torch.nn.Linear(topo_dim, bottleneck_dim), torch.nn.SiLU(), torch.nn.Linear(bottleneck_dim, bottleneck_dim), ) self.down_proj = torch.nn.Linear(hidden_dim, bottleneck_dim, bias=False) self.gating = torch.nn.Linear(bottleneck_dim, bottleneck_dim) self.up_proj = torch.nn.Linear(bottleneck_dim, hidden_dim, bias=False) torch.nn.init.zeros_(self.up_proj.weight) def forward(self, hidden_states: torch.Tensor, topo_condition: torch.Tensor) -> torch.Tensor: if topo_condition.ndim == 2: topo_condition = topo_condition.unsqueeze(1).expand(-1, hidden_states.shape[1], -1) topo_condition = topo_condition.to( device=hidden_states.device, dtype=self.topo_projector[0].weight.dtype, ) topo_condition = torch.nan_to_num(topo_condition, nan=0.0, posinf=0.0, neginf=0.0) hs = hidden_states.to(dtype=self.down_proj.weight.dtype) t_emb = self.topo_projector(topo_condition) gate = torch.sigmoid(self.gating(t_emb)) delta = self.up_proj(self.down_proj(hs) * gate).to(dtype=hidden_states.dtype) return hidden_states + delta class DecoderLayerWithTopoAdapter(torch.nn.Module): def __init__(self, base_layer: torch.nn.Module, adapter: TopoAdapter): super().__init__() self.base_layer = base_layer self.topo_adapter = adapter self.topo_condition: Optional[torch.Tensor] = None if hasattr(base_layer, "attention_type"): self.attention_type = base_layer.attention_type def set_topo_condition(self, condition: Optional[torch.Tensor]): self.topo_condition = condition def forward(self, *args, **kwargs): out = self.base_layer(*args, **kwargs) if self.topo_condition is None: return out if isinstance(out, tuple): hs = self.topo_adapter(out[0], self.topo_condition) return (hs,) + out[1:] return self.topo_adapter(out, self.topo_condition) def _find_decoder_layers(llm) -> torch.nn.ModuleList: candidates = [ "model.layers", "base_model.model.model.layers", "base_model.model.layers", "model.model.layers", ] for path in candidates: obj = llm ok = True for part in path.split("."): if not hasattr(obj, part): ok = False break obj = getattr(obj, part) if ok and isinstance(obj, torch.nn.ModuleList): return obj raise RuntimeError("Cannot locate Qwen decoder layers for TopoAdapter injection") def _install_topo_adapters( model, topo_dim: int = 47, bottleneck_dim: int = 32, last_n_layers: int = 8, every_n_layers: int = 0, ) -> List[DecoderLayerWithTopoAdapter]: layers = _find_decoder_layers(model.llm) hidden_dim = int(model.llm.get_input_embeddings().embedding_dim) n_layers = len(layers) selected = set(range(max(0, n_layers - last_n_layers), n_layers)) if last_n_layers > 0 else set() if every_n_layers > 0: selected.update(range(0, n_layers, every_n_layers)) wrappers: List[DecoderLayerWithTopoAdapter] = [] for idx in sorted(selected): if isinstance(layers[idx], DecoderLayerWithTopoAdapter): wrappers.append(layers[idx]) continue wrapper = DecoderLayerWithTopoAdapter( layers[idx], TopoAdapter(hidden_dim, topo_dim, bottleneck_dim), ) layers[idx] = wrapper wrappers.append(wrapper) print( f"Installed {len(wrappers)} TopoAdapters / {n_layers} decoder layers " f"| hidden={hidden_dim} topo_dim={topo_dim}" ) return wrappers def _topo_condition( prior_mask: torch.Tensor, topo_features: torch.Tensor, global_features: torch.Tensor, mode: str = "all", ) -> torch.Tensor: # Perform calculations in float32 for maximum numerical safety tf = topo_features.float() flat = tf.flatten(1, 2) parts = [flat.mean(1), flat.std(1), flat.amax(1)] if mode == "all": pm = prior_mask.float().flatten(1) parts.extend([ pm.mean(1, keepdim=True), pm.std(1, keepdim=True), pm.amax(1, keepdim=True), ]) parts.append(global_features.float()) cond = torch.cat(parts, dim=-1) return torch.nan_to_num(cond, nan=0.0, posinf=0.0, neginf=0.0) def _set_topo_condition(wrappers: List[DecoderLayerWithTopoAdapter], condition: Optional[torch.Tensor]): for wrapper in wrappers: wrapper.set_topo_condition(condition) def _patch_model_generate_with_topo(model, wrappers: List[DecoderLayerWithTopoAdapter], mode: str = "all"): original_generate = model.generate @torch.no_grad() def generate_with_topo( self, image, prior_mask, topo_features, global_features, question_text, max_new_tokens=64, ): condition = _topo_condition(prior_mask, topo_features, global_features, mode).to(image.device) _set_topo_condition(wrappers, condition) try: return original_generate( image, prior_mask, topo_features, global_features, question_text, max_new_tokens, ) finally: _set_topo_condition(wrappers, None) import types model.generate = types.MethodType(generate_with_topo, model) def _find_checkpoint() -> Path: for path in CHECKPOINT_CANDIDATES: if path.exists(): return path raise FileNotFoundError( "No checkpoint found. Expected one of: " + ", ".join(str(path.relative_to(REPO_DIR)) for path in CHECKPOINT_CANDIDATES) ) def _select_device() -> torch.device: return torch.device("cuda" if torch.cuda.is_available() else "cpu") def _safe_load_checkpoint(checkpoint_path: Path) -> Dict[str, Any]: try: return torch.load(checkpoint_path, map_location="cpu", weights_only=False) except TypeError: loaded = torch.load(checkpoint_path, map_location="cpu", weights_only=True) if isinstance(loaded, dict) and "model_state_dict" in loaded: return loaded return { "model_state_dict": loaded, "args": { "llm_name_or_path": "Qwen/Qwen2.5-3B-Instruct", "vision_backend": "timm", "freeze_llm": True, "use_lora": True, "lora_r": 16, "lora_alpha": 32, "lora_dropout": 0.05, "lora_target_modules": "q_proj,v_proj", "use_ot": True, "use_ot_fusion": True, "ot_fusion_mode": "prefix", "ot_fusion_dropout": 0.10, "use_prior_as_ot_target": True, "use_topological_loss": True, "use_prior_align_loss": True, "use_global_topo_loss": True, "use_patch_topo_loss": True, "topo_mode": "all", "vision_pretrained": True, "freeze_vision_backbone": True, }, } def _load_model_from_checkpoint(checkpoint_path: Path, device: torch.device): if build_structural_generative_vqa is None: raise ImportError( "Could not import src.models.structural_vqa_generative. " "Copy the project src/ directory into the HF repo root." ) print(f"Loading checkpoint from: {checkpoint_path}") checkpoint = _safe_load_checkpoint(checkpoint_path) train_args = checkpoint.get("args", {}) config = {**train_args} config.update( { "llm_name_or_path": train_args.get("llm_name_or_path", "Qwen/Qwen2.5-3B-Instruct"), "vision_pretrained": train_args.get("vision_pretrained", True), "vision_backend": train_args.get("vision_backend", "timm"), "freeze_vision_backbone": train_args.get("freeze_vision_backbone", True), "freeze_llm": train_args.get("freeze_llm", True), "use_lora": train_args.get("use_lora", True), "lora_r": train_args.get("lora_r", 16), "lora_alpha": train_args.get("lora_alpha", 32), "lora_dropout": train_args.get("lora_dropout", 0.05), "lora_target_modules": train_args.get( "lora_target_modules", "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", ), "use_prior_as_ot_target": train_args.get("use_prior_as_ot_target", True), "use_prior_align_loss": train_args.get("use_prior_align_loss", True), "use_global_topo_loss": train_args.get("use_global_topo_loss", True), "use_patch_topo_loss": train_args.get("use_patch_topo_loss", True), } ) allowed = { "llm_name_or_path", "vision_pretrained", "vision_backend", "freeze_vision_backbone", "freeze_llm", "max_question_length", "max_answer_length", "use_lora", "lora_r", "lora_alpha", "lora_dropout", "lora_target_modules", "ot_loss_weight", "use_ot", "use_ot_fusion", "ot_fusion_mode", "ot_fusion_dropout", "use_prior_as_ot_target", "prior_ot_global_mass", "use_topological_loss", "use_prior_align_loss", "use_global_topo_loss", "use_patch_topo_loss", "prior_loss_weight", "global_topo_loss_weight", "patch_topo_loss_weight", } build_kwargs = {key: value for key, value in config.items() if key in allowed} topo_mode = train_args.get("topo_mode", "all") topo_dim = 36 if topo_mode == "tda_only" else 47 bottleneck_dim = int(train_args.get("bottleneck_dim", train_args.get("adapter_bottleneck_dim", 32))) last_n_layers = int(train_args.get("adapter_last_n_layers", 8)) every_n_layers = int(train_args.get("adapter_every_n_layers", 0)) print( "Runtime config: " f"topo_mode={topo_mode} topo_dim={topo_dim} " f"vision_pretrained={build_kwargs.get('vision_pretrained')} " f"use_patch_topo_loss={build_kwargs.get('use_patch_topo_loss')}" ) model = build_structural_generative_vqa(**build_kwargs).to(device) model.tokenizer.padding_side = "left" # Set padding side to left for batch generation safety wrappers = _install_topo_adapters( model, topo_dim=topo_dim, bottleneck_dim=bottleneck_dim, last_n_layers=last_n_layers, every_n_layers=every_n_layers, ) for wrapper in wrappers: wrapper.to(device) _patch_model_generate_with_topo(model, wrappers, mode=topo_mode) state = checkpoint.get("model_state_dict", checkpoint) missing, unexpected = model.load_state_dict(state, strict=False) critical_keywords = ( "lora_", "visual_projector", "ot_visual_projector", "ot_text_projector", "topo_adapter", "global_topo_head", "patch_topo_head", ) critical_missing = [key for key in missing if any(word in key for word in critical_keywords)] unexpected_topo = [key for key in unexpected if "topo_adapter" in key] if any("topo_adapter" in key for key in critical_missing) and unexpected_topo: remapped_state = dict(state) prefix_pairs = [ ("llm.base_model.model.model.layers.", "llm.model.layers."), ("llm.model.layers.", "llm.base_model.model.model.layers."), ] for old_prefix, new_prefix in prefix_pairs: for key, value in state.items(): if key.startswith(old_prefix): remapped_state[new_prefix + key[len(old_prefix):]] = value missing, unexpected = model.load_state_dict(remapped_state, strict=False) critical_missing = [key for key in missing if any(word in key for word in critical_keywords)] # Detailed trainable validation trainable_missing = [ name for name, param in model.named_parameters() if param.requires_grad and name in missing ] print("\n" + "="*50) print("MODEL LOAD DIAGNOSTIC REPORT") print("="*50) print(f"Total missing keys: {len(missing)}") print(f"Total unexpected keys: {len(unexpected)}") print(f"Trainable missing keys: {len(trainable_missing)}") if trainable_missing: print(f"Warning: Missing trainable keys: {trainable_missing[:15]}") if critical_missing: print(f"Error: Critical missing fine-tuned keys: {critical_missing[:15]}") print("="*50 + "\n") if trainable_missing: raise RuntimeError( f"Checkpoint is missing trainable weights required for fine-tuning/inference! " f"Missing {len(trainable_missing)} keys, including: {trainable_missing[:15]}" ) if critical_missing: raise RuntimeError(f"Missing critical fine-tuned checkpoint keys: {critical_missing[:20]}") if unexpected: print(f"Warning: unexpected checkpoint keys: {unexpected[:20]}") print(f"Loaded checkpoint successfully. Status: OK") model.eval() model.llm.eval() model.vision_encoder.eval() return model def _pil_to_tensor(image: Image.Image) -> torch.Tensor: import cv2 image_rgb = np.asarray(image.convert("RGB")).astype(np.uint8) image_rgb = cv2.resize(image_rgb, DEFAULT_IMAGE_SIZE[::-1], interpolation=cv2.INTER_AREA) arr = image_rgb.astype(np.float32) / 255.0 arr = np.transpose(arr, (2, 0, 1)) return torch.from_numpy(arr).float().unsqueeze(0) _STRUCTURAL_EXTRACTORS = None def _get_structural_extractors(): global _STRUCTURAL_EXTRACTORS if _STRUCTURAL_EXTRACTORS is not None: return _STRUCTURAL_EXTRACTORS from src.topology.lesion_prior import LesionPriorExtractor from src.topology.tda_morphology import TopologicalExtractor as MorphologyTopologicalExtractor prior_extractor = LesionPriorExtractor( image_size=DEFAULT_IMAGE_SIZE, grid_size=DEFAULT_GRID_SIZE, use_morphology=False, ) morpho_extractor = MorphologyTopologicalExtractor( grid_size=DEFAULT_GRID_SIZE, image_size=DEFAULT_IMAGE_SIZE, ) _STRUCTURAL_EXTRACTORS = (prior_extractor, morpho_extractor) return _STRUCTURAL_EXTRACTORS def _extract_online_structural_tensors(image: Image.Image) -> Dict[str, torch.Tensor]: import cv2 prior_extractor, morpho_extractor = _get_structural_extractors() image_rgb = np.asarray(image.convert("RGB")).astype(np.uint8) image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR) prior_output = prior_extractor.extract_prior(image_bgr) morpho_output = morpho_extractor.extract_features(image_bgr) prior_mask = prior_output["prior_mask"].astype(np.float32) topo_features = morpho_output["topo_features"].astype(np.float32) global_features = morpho_output["global_features"].astype(np.float32) if prior_mask.shape != DEFAULT_GRID_SIZE: raise ValueError(f"prior_mask shape {prior_mask.shape} != {DEFAULT_GRID_SIZE}") expected_topo_shape = (*DEFAULT_GRID_SIZE, DEFAULT_TOPO_FEATURE_DIM) if topo_features.shape != expected_topo_shape: raise ValueError(f"topo_features shape {topo_features.shape} != {expected_topo_shape}") if global_features.shape != (DEFAULT_GLOBAL_FEATURE_DIM,): raise ValueError( f"global_features shape {global_features.shape} != {(DEFAULT_GLOBAL_FEATURE_DIM,)}" ) return { "prior_mask": torch.from_numpy(prior_mask).float().unsqueeze(0), "topo_features": torch.from_numpy(topo_features).float().unsqueeze(0), "global_features": torch.from_numpy(global_features).float().unsqueeze(0), } def predict_one(image: Image.Image, question: str, max_new_tokens: int = 48) -> str: model = load_model() device = _DEVICE or _select_device() image_tensor = _pil_to_tensor(image).to(device) structural = {key: value.to(device) for key, value in _extract_online_structural_tensors(image).items()} with torch.inference_mode(): prediction = model.generate( image=image_tensor, prior_mask=structural["prior_mask"], topo_features=structural["topo_features"], global_features=structural["global_features"], question_text=[question], max_new_tokens=max_new_tokens, )[0] return prediction.strip() def predict(batch: Iterable[Dict[str, Any]]) -> List[Dict[str, Any]]: outputs = [] for record in batch: image = record.get("image") if isinstance(image, (str, os.PathLike)): image = Image.open(image) if not isinstance(image, Image.Image): raise TypeError("record['image'] must be a PIL.Image.Image or image path") question = str(record.get("question", "")) answer = predict_one(image, question) out = {key: record[key] for key in ("id", "img_id", "question") if key in record} out["answer"] = answer outputs.append(out) return outputs class MedicoTask1Model: def __call__(self, image: Image.Image, question: str) -> str: return predict_one(image, question) def predict(self, image: Image.Image, question: str) -> str: return predict_one(image, question) def load(): load_model() return MedicoTask1Model() model_hf = _load_model_from_checkpoint(_find_checkpoint(), device) processor = None BATCH_SIZE = 32 # 🏁----------------END SUBMISISON DETAILS and MODEL LOADING -----------------🏁# start_time, post_model_mem = time.time(), get_mem() total_time, final_mem = round( time.time() - start_time, 4), round(get_mem() - post_model_mem, 2) model_mem_used = round(post_model_mem - initial_mem, 2) with tqdm(total=len(val_dataset), desc="Validating", unit="samples") as pbar: for batch in batched(enumerate(val_dataset), BATCH_SIZE): pbar.update(len(batch)) idxs, exs = zip(*batch) # ✏️✏️___________EDIT SECTION 2: ANSWER GENERATION___________✏️✏️# batch_images = [] batch_prior_masks = [] batch_topo_features = [] batch_global_features = [] batch_questions = [] for e in exs: img = e["image"] q = e["question"] img_tensor = _pil_to_tensor(img).to(device) batch_images.append(img_tensor) struct = _extract_online_structural_tensors(img) batch_prior_masks.append(struct["prior_mask"].to(device)) batch_topo_features.append(struct["topo_features"].to(device)) batch_global_features.append(struct["global_features"].to(device)) batch_questions.append(q) batch_img_tensor = torch.cat(batch_images, dim=0) batch_pm_tensor = torch.cat(batch_prior_masks, dim=0) batch_tf_tensor = torch.cat(batch_topo_features, dim=0) batch_gf_tensor = torch.cat(batch_global_features, dim=0) with torch.inference_mode(): answers = model_hf.generate( image=batch_img_tensor, prior_mask=batch_pm_tensor, topo_features=batch_tf_tensor, global_features=batch_gf_tensor, question_text=batch_questions, max_new_tokens=48, ) answers = [a.strip() for a in answers] # ___________END SECTION 2: ANSWER GENERATION___________ # # ⛔ DO NOT EDIT any lines below from here, can edit only upto decoding step above as required. ⛔ assert all(isinstance(a, str) for a in answers), next(f"Non-string answer at index {i}" for i, a in zip(idxs, answers) if not isinstance(a, str)) # safe test predictions.extend( {"index": i, "img_id": e["img_id"], "question": e["question"], "answer": a.strip()} for i, e, a in zip(idxs, exs, answers) ) # Ensure all predictions match dataset length assert len(predictions) == len( val_dataset), "Mismatch between predictions and dataset length" total_time, final_mem = round( time.time() - start_time, 4), round(get_mem() - post_model_mem, 2) model_mem_used = round(post_model_mem - initial_mem, 2) # caulcualtes metrics references = [[e] for e in val_dataset['answer']] preds = [pred['answer'] for pred in predictions] bleu_result = bleu.compute(predictions=preds, references=references) rouge_result = rouge.compute(predictions=preds, references=references) meteor_result = meteor.compute(predictions=preds, references=references) bleu_score = round(bleu_result['bleu'], 4) rouge1_score = round(float(rouge_result['rouge1']), 4) rouge2_score = round(float(rouge_result['rouge2']), 4) rougeL_score = round(float(rouge_result['rougeL']), 4) meteor_score = round(float(meteor_result['meteor']), 4) public_scores = { 'bleu': bleu_score, 'rouge1': rouge1_score, 'rouge2': rouge2_score, 'rougeL': rougeL_score, 'meteor': meteor_score } print("✨Public scores: ", public_scores) # Saves predictions to a JSON file output_data = {"submission_info": SUBMISSION_INFO, "public_scores": public_scores, "predictions": predictions, "total_time": total_time, "time_per_item": total_time / len(val_dataset), "memory_used_mb": final_mem, "model_memory_mb": model_mem_used, "gpu_name": gpu_name, "debug": { "packages": json.loads(subprocess.check_output([sys.executable, "-m", "pip", "list", "--format=json"])), "system": { "python": platform.python_version(), "os": platform.system(), "platform": platform.platform(), "arch": platform.machine() }}} with open("predictions_1.json", "w") as f: json.dump(output_data, f, indent=4) print(f"Time: {total_time}s | Mem: {final_mem}MB | Model Load Mem: {model_mem_used}MB | GPU: {gpu_name}") print("✅ Scripts Looks Good! Generation process completed successfully. Results saved to 'predictions_1.json'.") print("Next Step:\n 1) Upload this submission_task1.py script file to HuggingFace model repository.") print('''\n 2) Make a submission to the competition:\n Run:: medvqa validate_and_submit --competition=medico-2026 --task=1 --repo_id=...''')