| """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 |
| ) |
|
|
| |
| 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 |
| |
| 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)) |
|
|
| |
| lora_target_modules = args.lora_target_modules.split(",") |
| lora_module_names = find_proj_layers(model, lora_target_modules) |
|
|
| |
| 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 |
|
|
