add modeling files

align_transformers.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.models.dinov2.modeling_dinov2 import Dinov2Encoder
+
+from .configuration import AlignTransformerConfig
+
+
+def build_align_transformer(config):
+    if config.model_type == "align_transformer":
+        model = AlignTransformer(config)
+    else:
+        raise NotImplementedError()
+
+    return model
+
+
+class AlignTransformer(PreTrainedModel):
+    def __init__(self, config: AlignTransformerConfig):
+        super().__init__(config)
+
+        self.projector = None
+
+        if config.num_hidden_layers:
+            self.transformer_layers = Dinov2Encoder(config)
+        else:
+            self.transformer_layers = None
+
+        if config.use_layer_norm:
+            self.layer_norm = nn.LayerNorm(config.hidden_size)
+        else:
+            self.layer_norm = None
+
+    def forward(self, vision_tokens):
+
+        if self.projector is not None:
+
+            cls_token = vision_tokens[:, :1]
+            patch_tokens = vision_tokens[:, 1:]
+
+            patch_tokens = self.projector(patch_tokens)["last_hidden_state"]
+            vision_tokens = torch.cat([cls_token, patch_tokens], dim=1)
+
+        if self.transformer_layers is not None:
+            vision_tokens = self.transformer_layers(vision_tokens)["last_hidden_state"]
+
+        if self.layer_norm is not None:
+            vision_tokens = self.layer_norm(vision_tokens)
+
+        return vision_tokens

common_layers.py

@@ -0,0 +1,28 @@
+from torch import nn
+from transformers.modeling_utils import PreTrainedModel
+
+
+class BasePreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and
+    a simple interface for downloading and loading pretrained models.
+    """
+
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if (
+            isinstance(module, nn.Conv2d)  # noqa: SIM101
+            or isinstance(module, nn.Embedding)
+            or isinstance(module, nn.Linear)
+        ):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if hasattr(module, "bias") and module.bias is not None:
+                module.bias.data.zero_()
+
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.Parameter):
+            raise ValueError()

configuration.py

@@ -0,0 +1,129 @@
+from transformers import AutoConfig
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.dinov2.configuration_dinov2 import Dinov2Config
+
+
+class VisionConfig(PretrainedConfig):
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+    @staticmethod
+    def from_exp_config(vision_config: dict):
+
+        model_type = vision_config["model_type"]
+
+        if model_type in [
+            "siglip_vision_model",
+            "clip_vision_model",
+            "dinov2",
+            "sam",
+            "raddino",
+        ]:
+            config = AutoConfig.from_pretrained(
+                vision_config["pretrained_name_or_path"]
+            )
+            config = config.to_dict()
+            vision_config.update(config)
+        elif model_type == "xrayclip":
+            config = AutoConfig.from_pretrained(
+                vision_config["pretrained_name_or_path"]
+            )
+            config = config.to_dict()
+            config["model_type"] = "xrayclip"
+            vision_config.update(config)
+        elif model_type == "biomedclip":
+            pass
+        elif model_type == "m3ae":
+            pass
+
+        else:
+            raise NotImplementedError()
+
+        vision_config = VisionConfig(**vision_config)
+
+        return vision_config
+
+
+class TextConfig(PretrainedConfig):
+    def __init__(
+        self,
+        model_type,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_type = model_type
+
+    @staticmethod
+    def from_exp_config(
+        text_config: dict,
+    ):
+        model_type = text_config["model_type"]
+
+        if model_type in [
+            "siglip_text_model",
+            "clip_text_model",
+            "mpnet",
+            "biomedclip",
+            "bioclinicalmpbert",
+        ]:
+            text_config = TextConfig(**text_config)
+        else:
+            raise NotImplementedError()
+
+        return text_config
+
+
+class AlignTransformerConfig(PretrainedConfig):
+    def __init__(
+        self,
+        model_type: str = "align_transformer",
+        projector_config=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_type = model_type
+        self.projector_config = projector_config
+
+    @staticmethod
+    def from_exp_config(
+        align_transformer_config: dict,
+    ):
+        projector_config = align_transformer_config.pop("projector_config", None)
+
+        config = Dinov2Config(**align_transformer_config)
+        config = config.to_dict()
+
+        align_transformer_config = AlignTransformerConfig(
+            **(config | align_transformer_config),
+            projector_config=projector_config,
+        )
+
+        return align_transformer_config
+
+
+class CxrAlignConfig(PretrainedConfig):
+    is_composition = True
+
+    def __init__(
+        self,
+        vision_config: dict,
+        text_config: dict,
+        align_transformer_config: dict,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # Vision config
+        self.vision_config = VisionConfig.from_exp_config(vision_config)
+
+        # text config
+        self.text_config = TextConfig.from_exp_config(text_config)
+
+        self.align_transformer_config = AlignTransformerConfig.from_exp_config(
+            align_transformer_config
+        )
+
+        self.kwargs = kwargs

losses.py

@@ -0,0 +1,347 @@
+import math
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class KeyPhraseAlignmentLoss(nn.Module):
+
+    def __init__(
+        self,
+        hidden_dim=768,
+        use_vision_cls_token=True,
+        attn_temperature=None,
+        loss_temperature=0.07,
+        text_features_l2_norm=False,
+        mpnce_row_sum=False,
+        mpnce_col_sum=False,
+        sim_op="cos",
+        use_layer_norm=True,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+        self.layer_norm = nn.LayerNorm(hidden_dim) if use_layer_norm else None
+
+        self.use_vision_cls_token = use_vision_cls_token
+        self.loss_temperature = nn.Parameter(
+            torch.FloatTensor([np.log(loss_temperature)])
+        )
+        if attn_temperature is not None:
+            self.attn_temperature = nn.Parameter(
+                torch.FloatTensor([np.log(attn_temperature)])
+            )
+        else:
+            self.attn_temperature = None
+        self.text_features_l2_norm = text_features_l2_norm
+        self.sim_op = sim_op
+
+        self.similarity_logit = SimilarityLogit(sim_op)
+
+        self.mpnce_row_sum = mpnce_row_sum
+        self.mpnce_col_sum = mpnce_col_sum
+
+    def forward(
+        self,
+        key_phrases,
+        vision_tokens,
+        forward_text_model,
+        ddp_gather=True,
+        need_attn_weights=False,
+        compute_loss=True,
+        **kwargs,
+    ):
+        outputs = {}
+
+        text_features, group_map = self.compute_text_features(
+            key_phrases, forward_text_model, ddp_gather
+        )
+
+        if ddp_gather and dist.is_initialized():
+            vision_tokens = torch.cat(dist.nn.all_gather(vision_tokens), dim=0)
+
+        if self.layer_norm is not None:
+            vision_tokens = self.layer_norm(vision_tokens)
+
+        vision_patch_tokens = vision_tokens[:, 1:]
+
+        # text to image cross-attention
+        if not self.use_vision_cls_token:
+            vision_attn_tokens = vision_patch_tokens
+        else:
+            vision_attn_tokens = vision_tokens
+
+        t2i_logits, t2i_attn_weights_list = self.compute_t2i_logits(
+            text_features, vision_attn_tokens, need_attn_weights
+        )
+        outputs["t2i_logits"] = t2i_logits
+        outputs["t2i_attn_weights"] = t2i_attn_weights_list
+
+        if compute_loss:
+            losses = {}
+            loss = 0
+
+            # compute t2i loss
+            t2i_loss = multi_positive_nce_loss(
+                t2i_logits,
+                group_map,
+                temperature=self.loss_temperature.exp(),
+                row_sum=self.mpnce_row_sum,
+                col_sum=self.mpnce_col_sum,
+            )
+            loss += t2i_loss
+            losses["t2i_loss"] = t2i_loss
+
+            losses["loss"] = loss
+            outputs["losses"] = losses
+        return outputs
+
+    def compute_text_features(self, key_phrases, forward_text_model, ddp_gather=True):
+
+        key_text_features_list = list()
+        group_list = list()
+
+        B_local = len(key_phrases)
+        # Calculate offset by getting the rank of the current process when using DDP
+        local_rank = dist.get_rank() if (ddp_gather and dist.is_initialized()) else 0
+
+        for i, kp in enumerate(key_phrases):
+            feats = forward_text_model(kp)
+
+            # (N_i, D)
+            if self.text_features_l2_norm:
+                feat = feats["text_features"]
+            else:
+                feat = feats["text_features_wo_l2_norm"]
+
+            if feat.shape[-1] == 2 * self.hidden_dim:
+                feat = feat[:, self.hidden_dim :]
+
+            key_text_features_list.append(feat)
+
+            # Add local_rank * B_local offset to local index i
+            global_index = i + local_rank * B_local
+            group_list.extend([global_index] * feat.size(0))
+
+        text_features = torch.cat(key_text_features_list, dim=0)
+        group_map = torch.tensor(group_list, device=text_features.device)
+
+        if ddp_gather and dist.is_initialized():
+            # Gather text_features and image_features and group_map
+            text_features = pad_and_gather(text_features)
+
+            group_map = pad_and_gather(group_map)
+            group_map = group_map.long()
+
+        if self.layer_norm is not None:
+            text_features = self.layer_norm(text_features)
+
+        return text_features, group_map
+
+    def compute_t2i_logits(
+        self, text_features, vision_attn_tokens, need_attn_weights, repeat=True
+    ):
+
+        t2i_logits, t2i_attn_weights_list = self.similarity_logit(
+            text_features,
+            vision_attn_tokens,
+            need_attn_weights,
+            repeat=repeat,
+            temperature=(
+                self.attn_temperature.exp()
+                if self.attn_temperature is not None
+                else self.loss_temperature.exp()
+            ),
+        )
+
+        return t2i_logits, t2i_attn_weights_list
+
+
+class SimilarityLogit(nn.Module):
+    def __init__(self, sim_op="dot", **kwargs):
+        super().__init__()
+        self.sim_op = sim_op
+
+    def forward(
+        self,
+        queries: torch.Tensor,
+        local_tokens: torch.Tensor,
+        need_attn_weights: bool = False,
+        repeat: bool = True,
+        **kwargs,
+    ):
+        if repeat:
+            query_attn_features = queries.unsqueeze(0).expand(
+                local_tokens.shape[0], queries.shape[0], queries.shape[1]
+            )
+        else:
+            assert queries.dim() == 3
+            query_attn_features = queries
+
+        if self.sim_op == "cos":
+            temperature = kwargs.get("temperature")
+            assert temperature is not None
+            denominator = temperature
+            query_attn_features = F.normalize(query_attn_features, p=2, dim=-1)
+            local_tokens = F.normalize(local_tokens, p=2, dim=-1)
+        elif self.sim_op == "dot":
+            denominator = math.sqrt(local_tokens.size(-1))
+        else:
+            raise NotImplementedError
+
+        scores = (
+            torch.bmm(query_attn_features, local_tokens.permute(0, 2, 1)) / denominator
+        )
+        attn_weights = F.softmax(scores, dim=-1)
+
+        aggregated = torch.matmul(attn_weights, local_tokens)
+
+        query_attn_features = F.normalize(query_attn_features, p=2, dim=-1)
+        aggregated = F.normalize(aggregated, p=2, dim=-1)
+
+        logits = torch.matmul(
+            query_attn_features.unsqueeze(2), aggregated.unsqueeze(-1)
+        ).squeeze()
+
+        logits = logits.T
+
+        if need_attn_weights:
+            attn_scores = [scores]
+        else:
+            attn_scores = None
+
+        return logits, attn_scores
+
+
+def multi_positive_nce_loss(
+    logits: torch.Tensor,
+    group_map: torch.Tensor,
+    temperature: float = 1.0,
+    eps: float = 1e-8,
+    row_sum: bool = False,
+    col_sum: bool = False,
+):
+    """
+    Args:
+        logits: tensor of shape (N_total, B_global), each row is a logit between a key phrase and each candidate image.
+        group_map: tensor of shape (N_total,), source image index of each key phrase.
+        temperature: scaling factor.
+
+    For each key phrase row i, the positive is the candidate image index == group_map[i],
+    and the rest are treated as negatives.
+
+    For each column j, each positive for image j is considered independently.
+
+    Returns:
+        loss: scalar tensor.
+    """
+    scaled_logits = torch.exp(logits / temperature)  # (N_total, B_global)
+
+    pos_logits = scaled_logits[
+        torch.arange(scaled_logits.size(0)), group_map
+    ]  # (N_total,)
+
+    row_loss = get_row_loss(
+        scaled_logits,
+        pos_logits,
+        group_map,
+        eps,
+        row_sum,
+    )
+
+    neg_mask = torch.ones_like(scaled_logits)
+    neg_mask[torch.arange(scaled_logits.size(0)), group_map] = 0  # (N_total, B_global)
+
+    column_loss = get_col_loss(
+        scaled_logits,
+        pos_logits,
+        neg_mask,
+        group_map,
+        eps,
+        col_sum,
+    )
+
+    loss = (row_loss.mean() + column_loss.mean()) / 2
+
+    return loss
+
+
+def get_row_loss(
+    logits: torch.Tensor,
+    pos_logits: torch.Tensor,
+    group_map: torch.Tensor,
+    eps: float = 1e-8,
+    row_sum: bool = False,
+):
+    if row_sum:
+        # Create a tensor to hold the summed values
+        row_sum_logits = torch.zeros(
+            logits.shape[-1], device=logits.device
+        )  # (B_global)
+        row_pos_sum_logits = torch.zeros(
+            logits.shape[-1], device=logits.device
+        )  # (B_global)
+
+        # Use scatter_add to sum values based on group_map
+        row_sum_logits.scatter_add_(0, group_map, logits.sum(dim=1))  # (B_global)
+        row_pos_sum_logits.scatter_add_(0, group_map, pos_logits)  # (B_global)
+        p_row = row_pos_sum_logits / (row_sum_logits + eps)  # (B_global)
+    else:
+        row_sum_logits = logits.sum(dim=1)  # (N_total)
+        p_row = pos_logits / (row_sum_logits + eps)  # (N_total)
+
+    return -torch.log(p_row + eps)
+
+
+def get_col_loss(
+    logits: torch.Tensor,
+    pos_logits: torch.Tensor,
+    neg_mask: torch.Tensor,
+    group_map: torch.Tensor,
+    eps: float = 1e-8,
+    col_sum: bool = False,
+):
+    if col_sum:
+        # MIL-NCE loss
+        column_sum_logits = logits.sum(dim=0)  # (B_global,)
+        pos_mask = torch.ones_like(logits) - neg_mask  # (N_total, B_global)
+        column_pos_logits = (logits * pos_mask).sum(dim=0)  # (B_global,)
+        p_column = column_pos_logits / (column_sum_logits + eps)  # (B_global,)
+    else:
+        # MP-NCE loss (UniCLIP)
+        neg_logits = logits * neg_mask  # (N_total, B_global)
+        sum_neg_logits = neg_logits.sum(dim=0)  # (B_global,)
+        sum_neg_logits = sum_neg_logits[group_map]  # (N_total)
+        p_column = pos_logits / (pos_logits + sum_neg_logits + eps)  # (N_total)
+
+    return -torch.log(p_column + eps)
+
+
+def pad_and_gather(tensor):
+    # Determine the size of the tensor
+    local_size = torch.tensor(tensor.size(), device=tensor.device)
+
+    # Gather all sizes
+    all_sizes = [torch.zeros_like(local_size) for _ in range(dist.get_world_size())]
+    dist.all_gather(all_sizes, local_size)
+
+    # Determine the maximum size
+    max_size = torch.stack(all_sizes).max(dim=0)[0]
+
+    # Pad the tensor to the maximum size
+    padded_tensor = torch.zeros(max_size.tolist(), device=tensor.device)
+    padded_tensor[: local_size[0]] = tensor
+
+    # Gather all padded tensors
+    gathered_tensors = dist.nn.all_gather(padded_tensor)
+
+    # Trim the gathered tensors to their original sizes
+    gathered_tensors = [g[: s[0]] for g, s in zip(gathered_tensors, all_sizes)]
+
+    gathered_tensors = torch.cat(gathered_tensors, dim=0)
+
+    return gathered_tensors

radzero_modeling.py

@@ -0,0 +1,302 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoTokenizer, BertModel
+from transformers.models.clip.modeling_clip import CLIPTextModel
+from transformers.models.mpnet.modeling_mpnet import MPNetModel
+from transformers.trainer import logger
+
+from .align_transformers import build_align_transformer
+from .common_layers import BasePreTrainedModel
+from .configuration import CxrAlignConfig
+from .losses import KeyPhraseAlignmentLoss
+from .text_encoders import aggregate_tokens, build_text_encoder
+from .vision_encoders import MRM, Dinov2Model, build_vision_encoder
+
+
+class CxrAlignModel(BasePreTrainedModel):
+
+    config_class = CxrAlignConfig
+
+    def build_vision_model(self, config: CxrAlignConfig):
+        vision_config = config.vision_config
+        vision_config.pretrained_dir = config.pretrained_dir
+        vision_model = build_vision_encoder(vision_config)
+        return vision_model
+
+    def build_text_model(self, config: CxrAlignConfig):
+        text_config = config.text_config
+        text_model = build_text_encoder(text_config)
+
+        if text_config.model_type == "bioclinicalmpbert":
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                text_config.pretrained_tokenizer_name_or_path
+            )
+            self.idxtoword = {v: k for k, v in self.tokenizer.get_vocab().items()}
+
+        return text_model
+
+    def build_align_transformer_model(self, config: CxrAlignConfig):
+        align_transformer_config = config.align_transformer_config
+        align_transformer = build_align_transformer(align_transformer_config)
+
+        return align_transformer
+
+    def __init__(self, config: CxrAlignConfig):
+        super().__init__(config)
+
+        logger.info("Build vision model ...")
+        self.vision_model = self.build_vision_model(config)
+
+        logger.info("Build text model ...")
+        self.text_model = self.build_text_model(config)
+
+        if (
+            isinstance(self.text_model, CLIPTextModel)
+            or isinstance(self.text_model, MPNetModel)
+            or isinstance(self.text_model, BertModel)
+        ):
+            text_dim = self.text_model.config.hidden_size
+
+        self.hidden_size = config.align_transformer_config.hidden_size
+
+        if config.text_config.use_text_projection:
+            self.text_projector = nn.Linear(text_dim, 2 * self.hidden_size)
+        else:
+            self.text_projector = None
+
+        logger.info("Build align transformer model ...")
+        self.align_transformer = self.build_align_transformer_model(config)
+
+        logger.info("Build loss functions ...")
+        loss_cfg = config.kwargs["loss"]
+        self.loss_ratio = dict()
+        self.loss_fns = nn.ModuleDict()
+        for loss_type, ratio in zip(loss_cfg["apply"], loss_cfg["ratio"]):
+            logger.info(f"Build {loss_type} loss function ...")
+            if loss_cfg[loss_type] is None:
+                loss_cfg[loss_type] = dict()
+            if torch.distributed.is_available() and torch.distributed.is_initialized():
+                loss_cfg[loss_type]["rank"] = torch.distributed.get_rank()
+                loss_cfg[loss_type]["world_size"] = torch.distributed.get_world_size()
+            self.loss_fns[loss_type] = eval(loss_type)(**loss_cfg[loss_type])
+            self.loss_ratio[loss_type] = ratio
+
+        self.compute_logits_type = config.kwargs.get("compute_logits_type")
+        self.use_negative_logits = config.kwargs.get("use_negative_logits")
+
+        self.module_to_update = config.kwargs.get("module_to_update")
+
+        self.post_init()
+
+    def forward_vision_model(self, pixel_values):
+
+        if isinstance(self.vision_model, Dinov2Model):
+            vision_tokens = self.vision_model(pixel_values)["last_hidden_state"]
+        elif isinstance(self.vision_model, MRM):
+            img_emb_g, img_emb_l = self.vision_model(pixel_values)
+            img_emb_g = img_emb_g.unsqueeze(1)
+            img_emb_l = img_emb_l.view(img_emb_l.size(0), img_emb_l.size(1), -1)
+            img_emb_l = img_emb_l.permute(0, 2, 1)
+
+            vision_tokens = torch.cat([img_emb_g, img_emb_l], dim=1)
+        else:
+            raise NotImplementedError
+
+        vision_tokens = self.align_transformer(vision_tokens)
+
+        cls_token = vision_tokens[:, 0]
+        patch_tokens = vision_tokens[:, 1:]
+        image_features = torch.cat([cls_token, patch_tokens.mean(dim=1)], dim=1)
+        image_features = F.normalize(image_features, p=2, dim=1)
+
+        outputs = {}
+        outputs["vision_tokens"] = vision_tokens
+        outputs["image_cls_token"] = cls_token
+        outputs["image_patch_tokens"] = patch_tokens
+        outputs["image_features"] = image_features
+
+        return outputs
+
+    def forward_text_model(self, encoded_input):
+        text_outputs = {}
+
+        if isinstance(self.text_model, MPNetModel):
+            model_output = self.text_model(
+                input_ids=encoded_input["input_ids"],
+                attention_mask=encoded_input["attention_mask"],
+            )
+
+            token_embeddings = model_output[
+                0
+            ]  # First element of model_output contains all token embeddings
+
+            # text embedding projection
+            if self.text_projector is not None:
+                token_embeddings = self.text_projector(token_embeddings)
+
+            # token_embeddings = self.text_projector(token_embeddings)
+            if self.config.text_config.use_cls_token:
+                text_features = token_embeddings[:, 0, :]
+
+            else:
+                # mean pooling
+                input_mask_expanded = (
+                    encoded_input["attention_mask"]
+                    .unsqueeze(-1)
+                    .expand(token_embeddings.size())
+                    .float()
+                )
+                text_features = torch.sum(
+                    token_embeddings * input_mask_expanded, 1
+                ) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+
+        elif isinstance(self.text_model, BertModel):
+            # BioClinicalMPBERT
+
+            model_output = self.text_model(
+                input_ids=encoded_input["input_ids"],
+                attention_mask=encoded_input["attention_mask"],
+                token_type_ids=encoded_input.get("token_type_ids", None),
+            )
+
+            if self.config.text_config.use_cls_token:
+                text_features = model_output.last_hidden_state[:, 0, :]
+
+            elif self.config.text_config.use_aggregate_tokens:
+
+                all_embeddings = model_output[2]
+                embeddings = torch.stack(
+                    all_embeddings[-self.config.text_config.last_n_layers :]
+                )
+                embeddings = embeddings.permute(1, 0, 2, 3)
+
+                embeddings, sents = aggregate_tokens(
+                    embeddings, encoded_input["input_ids"], self.idxtoword
+                )
+                sent_embeddings = embeddings.mean(axis=2)
+
+                if self.config.text_config.aggregate_method == "sum":
+                    word_embeddings = embeddings.sum(axis=1)
+                    sent_embeddings = sent_embeddings.sum(axis=1)
+                elif self.config.text_config.aggregate_method == "mean":
+                    word_embeddings = embeddings.mean(axis=1)
+                    sent_embeddings = sent_embeddings.mean(axis=1)
+
+                word_embeddings = word_embeddings.permute(0, 2, 1)
+
+                text_features = sent_embeddings
+                text_outputs["word_embeddings"] = word_embeddings
+
+            else:
+                text_features = model_output.last_hidden_state
+                mask = encoded_input["attention_mask"].unsqueeze(-1).float()
+                text_features = torch.sum(text_features * mask, dim=1) / torch.clamp(
+                    mask.sum(dim=1), min=1e-9
+                )
+
+            if self.text_projector is not None:
+                text_features = self.text_projector(text_features)
+
+        else:
+            raise NotImplementedError
+
+        text_outputs["text_features_wo_l2_norm"] = text_features
+        text_outputs["text_features"] = F.normalize(text_features, p=2, dim=1)
+
+        return text_outputs
+
+    def forward(
+        self,
+        pixel_values,
+        encoded_key_phrases=None,
+        return_loss=True,
+        **kwargs,
+    ):
+        vision_outputs = self.forward_vision_model(pixel_values)
+
+        outputs = {}
+        outputs.update(vision_outputs)
+
+        # Trainer's self.can_return_loss is True if 'return_loss' is in model's forward function
+        if return_loss:
+            loss = 0
+            losses = {}
+
+            for loss_type, loss_fn in self.loss_fns.items():
+                if isinstance(loss_fn, KeyPhraseAlignmentLoss):
+                    loss_outputs = loss_fn(
+                        encoded_key_phrases,
+                        outputs["vision_tokens"],
+                        self.forward_text_model,
+                    )
+                    key_phrase_alignment_losses = loss_outputs["losses"]
+                    losses["key_phrase_alignment_loss"] = (
+                        key_phrase_alignment_losses.pop("loss")
+                    )
+                    for loss_name, loss_value in key_phrase_alignment_losses.items():
+                        losses[loss_name] = loss_value
+                    loop_loss = losses["key_phrase_alignment_loss"]
+                else:
+                    raise NotImplementedError
+
+                loss += loop_loss * self.loss_ratio[loss_type]
+
+            losses["loss"] = loss
+
+            outputs["losses"] = losses
+
+        return outputs
+
+    def compute_logits(
+        self,
+        pixel_values,
+        encoded_key_phrases,
+        **kwargs,
+    ):
+        vision_outputs = self.forward_vision_model(pixel_values)
+
+        outputs = {}
+
+        if self.compute_logits_type == "key_phrase_alignment":
+
+            splited_key_phrases = [
+                {
+                    "input_ids": encoded_key_phrases[0]["input_ids"][i : i + 1],
+                    "attention_mask": encoded_key_phrases[0]["attention_mask"][
+                        i : i + 1
+                    ],
+                }
+                for i in range(encoded_key_phrases[0]["input_ids"].size(0))
+            ]
+
+            loss_outputs = self.loss_fns["KeyPhraseAlignmentLoss"](
+                splited_key_phrases,
+                vision_outputs["vision_tokens"],
+                self.forward_text_model,
+                ddp_gather=False,
+                need_attn_weights=True,
+                compute_loss=False,
+            )
+            outputs.update(loss_outputs)
+
+            # mean attention weights from all layers
+            outputs["similarity_scores"] = torch.mean(
+                torch.stack(loss_outputs["t2i_attn_weights"]), dim=0
+            )
+
+            # remove attention score for cls token
+            if self.loss_fns["KeyPhraseAlignmentLoss"].use_vision_cls_token:
+                outputs["similarity_scores"] = outputs["similarity_scores"][:, :, 1:]
+
+            # compute logits
+            logits = loss_outputs["t2i_logits"]
+            logits = logits.T
+
+            logits = (
+                logits / self.loss_fns["KeyPhraseAlignmentLoss"].loss_temperature.exp()
+            )
+
+        outputs["logits"] = logits
+        return outputs

text_encoders.py

@@ -0,0 +1,27 @@
+import open_clip
+import torch
+from transformers import AutoModel
+from transformers.models.clip.modeling_clip import CLIPTextModel
+from transformers.models.siglip.modeling_siglip import SiglipTextModel
+
+
+def build_text_encoder(config):
+    if config.model_type == "mpnet":
+        model = AutoModel.from_pretrained(config.pretrained_name_or_path)
+    else:
+        raise NotImplementedError()
+
+    return model
+
+
+# Mean Pooling - Take attention mask into account for correct averaging
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[
+        0
+    ]  # First element of model_output contains all token embeddings
+    input_mask_expanded = (
+        attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    )
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
+        input_mask_expanded.sum(1), min=1e-9
+    )

vision_encoders.py

@@ -0,0 +1,11 @@
+from transformers import Dinov2Model
+
+
+def build_vision_encoder(config):
+    if config.model_type == "dinov2":
+        model = Dinov2Model.from_pretrained(config.pretrained_name_or_path)
+
+    else:
+        raise NotImplementedError()
+
+    return model