nvidia/C-RADIOv4-H

1# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8from argparse import Namespace9import string10from typing import List11 12import torch13from torch import nn14import torch.nn.functional as F15 16from .adaptor_registry import adaptor_registry, dict_t, state_t17 18from .adaptor_generic import GenericAdaptor19from .utils import rank_gate20 21 22_VERSION_MAP = {23    'siglip2-g-384': 'google/siglip2-giant-opt-patch16-384',24    'siglip2-so400m': 'google/siglip2-so400m-patch16-naflex',25}26 27 28class SigLIP2Adaptor(GenericAdaptor):29    def __init__(self, main_config: Namespace, adaptor_config: dict_t, state: state_t):30        super().__init__(main_config, adaptor_config, state)31 32        version = adaptor_config['model']33        version = _VERSION_MAP[version]34 35        from transformers import AutoModel, AutoProcessor36        with rank_gate():37            model = AutoModel.from_pretrained(version, trust_remote_code=True)38            proc = AutoProcessor.from_pretrained(version, trust_remote_code=True)39 40        self.tokenizer = SigLIP2WrappedTokenizer(proc)41        self.text_model = model.text_model42 43        del model44 45    def encode_text(self, text, normalize: bool = False):46        output = self.text_model(**text, return_dict=True)47        token = output.pooler_output48 49        if normalize:50            token = F.normalize(token, dim=-1)51 52        return token53 54 55class SigLIP2WrappedTokenizer:56    def __init__(self, proc):57        self._proc = proc58 59    def __call__(self, text: List[str]):60        text = [canonicalize_text(t) for t in text]61        ret = self._proc(text=text, return_tensors='pt', max_length=64, padding='max_length', truncation=True)62        return ret63 64 65def canonicalize_text(66    text: str,67    *,68    keep_punctuation_exact_string=None,69    trans_punctuation: dict = str.maketrans("", "", string.punctuation),70):71    """Returns canonicalized `text` (lowercase and punctuation removed).72 73    From: https://github.com/google-research/big_vision/blob/53f18caf27a9419231bbf08d3388b07671616d3d/big_vision/evaluators/proj/image_text/prompt_engineering.py#L9474 75    Args:76      text: string to be canonicalized.77      keep_punctuation_exact_string: If provided, then this exact string kept.78        For example providing '{}' will keep any occurrences of '{}' (but will79        still remove '{' and '}' that appear separately).80    """81    text = text.replace("_", " ")82    if keep_punctuation_exact_string:83        text = keep_punctuation_exact_string.join(84            part.translate(trans_punctuation)85            for part in text.split(keep_punctuation_exact_string)86        )87    else:88        text = text.translate(trans_punctuation)89    text = text.lower()90    text = " ".join(text.split())91    return text.strip()92 93 94@adaptor_registry.register_adaptor("siglip2")95def create_siglip2_adaptor(main_config: Namespace, adaptor_config: dict_t, state: state_t):96    return SigLIP2Adaptor(main_config, adaptor_config, state)97