"""Shared PRIMA model construction for inference and benchmark evaluation.""" import torch import transformers from pathlib import Path from peft import LoraConfig, get_peft_model from model.PRIMA import PRIMAForCausalLM from tools.checkpoint import load_checkpoint from tools.config import validate_paths DEFAULT_HUB_MODEL = 'PLAN-Lab/PRIMA' MODEL_DEFAULTS = { 'merged_ckpt_path': '', 'hub_revision': 'main', 'version': '/path/to/base_model', 'ckpt_path': '/path/to/prima_weights', 'q_former_model': '/path/to/qformer_weights', 'pretrained_q_former_path': '', 'qformer_bert_path': '/path/to/bert_model', 'eva_weights': '/path/to/eva_weights', 'pretrain_mm_mlp_adapter': '', 'precision': 'bf16', 'image_size': 1024, 'vit_image_size': 224, 'model_max_length': 3400, 'max_num_images': 5, 'num_query_tokens': 32, 'lora_r': 8, 'lora_alpha': 16, 'lora_dropout': 0.05, 'lora_target_modules': 'q_proj,v_proj', 'conv_type': 'llava_v1', 'local_rank': 0, } INTERNAL_DEFAULTS = {'use_mm_start_end': True, 'img_emb_len': None} def add_model_arguments(parser, *, merged_checkpoint=''): for key, value in MODEL_DEFAULTS.items(): parser.add_argument('--' + key, default=value, type=type(value)) parser.set_defaults(merged_ckpt_path=merged_checkpoint, **INTERNAL_DEFAULTS) def normalize_model_args(args): for key, value in {**MODEL_DEFAULTS, **INTERNAL_DEFAULTS}.items(): if not hasattr(args, key): setattr(args, key, value) if args.precision != 'bf16': raise ValueError('The release currently supports the validated bf16 execution path.') if not args.pretrained_q_former_path or args.pretrained_q_former_path.startswith('/path/to/'): args.pretrained_q_former_path = args.q_former_model return args def setup_tokenizer_and_special_tokens(args): """Load tokenizer and add special tokens.""" source = resolve_merged_checkpoint(args) if args.merged_ckpt_path else args.version tokenizer = transformers.AutoTokenizer.from_pretrained( source, model_max_length=args.model_max_length, padding_side="left", use_fast=False, ) tokenizer.pad_token = tokenizer.unk_token args.seg_token_idx = tokenizer.convert_tokens_to_ids("[SEG]") if args.seg_token_idx == tokenizer.unk_token_id: raise ValueError("The base tokenizer must contain the [SEG] special token.") args.img_emb_len = args.num_query_tokens max_images = args.model_max_length // args.img_emb_len max_images = min(max_images, args.max_num_images) args.seg_image_tokens = [] for i in range(max_images): args.seg_image_tokens.append( tokenizer(f"[SEG] (IMAGE{i + 1}", add_special_tokens=False).input_ids ) args.seg_image_tokens = [pattern[pattern.index(args.seg_token_idx):] for pattern in args.seg_image_tokens] return tokenizer def setup_qformer_tokenizer(args): source = (Path(resolve_merged_checkpoint(args)) / 'qformer_tokenizer' if args.merged_ckpt_path else args.qformer_bert_path) qformer_tokenizer = transformers.BertTokenizer.from_pretrained( source, truncation_side="left" ) qformer_tokenizer.add_special_tokens({"bos_token": "[DEC]"}) return qformer_tokenizer def initialize_model(args, tokenizer): """Initialize the PRIMA model.""" model_args = { k: getattr(args, k) for k in [ "seg_token_idx", "seg_image_tokens", "max_num_images", "pretrain_mm_mlp_adapter", ] } model_args["img_emb_len"] = args.img_emb_len for name in ('ce_loss_weight', 'dice_loss_weight', 'bce_loss_weight', 'train_mask_decoder'): if hasattr(args, name): model_args[name] = getattr(args, name) source = resolve_merged_checkpoint(args) if args.merged_ckpt_path else args.version model = PRIMAForCausalLM.from_pretrained( source, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, **model_args ) # Configure model tokens model.config.eos_token_id = tokenizer.eos_token_id model.config.bos_token_id = tokenizer.bos_token_id model.config.pad_token_id = tokenizer.pad_token_id if args.merged_ckpt_path: if not getattr(model.config, 'self_contained', False): raise ValueError('merged_ckpt_path must point to a complete PRIMA checkpoint.') else: model.config.eva_weights = args.eva_weights model.config.qformer_bert_path = args.qformer_bert_path return model def prepare_model_for_inference(model, args, qformer_tokenizer): if args.merged_ckpt_path: if args.ckpt_path and not args.ckpt_path.startswith('/path/to/'): return get_peft_model(model, setup_lora_config(model, args)) return model # Initialize vision tower model.get_model().initialize_qformer(model.get_model().config, pretrained_q_former_path=args.pretrained_q_former_path, vit_image_size=args.vit_image_size, precision=args.precision, q_former_model=args.q_former_model, num_query_tokens=args.num_query_tokens, qformer_tokenizer=qformer_tokenizer, pretrain_mm_mlp_adapter=args.pretrain_mm_mlp_adapter) model.get_model().Qformer.to(dtype=torch.bfloat16, device=args.local_rank) model.get_model().mm_projector.to(dtype=torch.bfloat16, device=args.local_rank) model.get_model().visual_encoder.to(dtype=torch.bfloat16, device=args.local_rank) model.get_model().ln_vision.to(dtype=torch.bfloat16, device=args.local_rank) model.get_model().query_tokens.to(dtype=torch.bfloat16, device=args.local_rank) if args.lora_r > 0: lora_config = setup_lora_config(model, args) model = get_peft_model(model, lora_config) return model def setup_lora_config(model, args): """Configure LoRA settings for the model.""" def find_proj_layers(model, target_modules): """Identify projection layers in the model for LoRA adaptation.""" linear_cls = torch.nn.Linear lora_module_names = set() for name, module in model.named_modules(): if ( isinstance(module, linear_cls) and all( x not in name for x in [ "grounding_encoder", "vision_tower", "mm_projector", "text_hidden_fcs", "visual_encoder", "ln_vision", "Qformer", "query_tokens", ] ) and any(x in name for x in target_modules) ): lora_module_names.add(name) return sorted(list(lora_module_names)) # Extracting LoRA target modules lora_target_modules = args.lora_target_modules.split(",") lora_module_names = find_proj_layers(model, lora_target_modules) # Configuring LoRA lora_config = LoraConfig( r=args.lora_r, lora_alpha=args.lora_alpha, target_modules=lora_module_names, lora_dropout=args.lora_dropout, bias="none", task_type="CAUSAL_LM", ) return lora_config def load_model_from_checkpoint(model, args): overlay = args.ckpt_path and not args.ckpt_path.startswith('/path/to/') report = load_checkpoint(model, args.ckpt_path) if overlay or not args.merged_ckpt_path else {} model.to(device=torch.device(f'cuda:{args.local_rank}'), dtype=torch.bfloat16) return report def resolve_merged_checkpoint(args): """Resolve a local complete model or a pinned Hub snapshot once per command.""" if not getattr(args, '_merged_snapshot', None): source = Path(args.merged_ckpt_path).expanduser() if source.is_dir(): args._merged_snapshot = str(source.resolve()) else: from huggingface_hub import snapshot_download args._merged_snapshot = snapshot_download( args.merged_ckpt_path, revision=args.hub_revision) return args._merged_snapshot def load_bundle(args): normalize_model_args(args) validate_paths(args) tokenizer = setup_tokenizer_and_special_tokens(args) qformer_tokenizer = setup_qformer_tokenizer(args) model = initialize_model(args, tokenizer) model = prepare_model_for_inference(model, args, qformer_tokenizer) load_model_from_checkpoint(model, args) return model.eval(), tokenizer, qformer_tokenizer