File size: 5,619 Bytes
071945c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
# coding=utf-8
# Copyright 2024 LY Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Optional
import torch
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.models.clip.modeling_clip import CLIPOutput
from .configuration_clyp import CLYPConfig, CLYPLossConfig
from .model import InfoNCELoss, create_text_encoder, create_vision_encoder
from .model_rinna import RinnaCLIPModel # noqa
@dataclass
class CLYPOutput(CLIPOutput):
...
class CLYPPreTrainedModel(PreTrainedModel):
config_class = CLYPConfig
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def _init_weights(self, module: Any) -> None:
pass
class CLYPModel(CLYPPreTrainedModel):
def __init__(self, config: CLYPConfig):
super().__init__(config)
self.vision_encoder = create_vision_encoder(config.vision_encoder_config)
self.text_encoder = create_text_encoder(config.text_encoder_config)
self.initialize_clip(
learn_temperature=config.learn_temperature,
temperature_init=config.temperature_init,
temperature_min=config.temperature_min,
temperature_max=config.temperature_max,
itc_loss_config=config.itc_loss_config,
)
def initialize_clip(
self,
learn_temperature: Optional[bool] = None,
temperature_init: Optional[float] = None,
temperature_min: Optional[float] = None,
temperature_max: Optional[float] = None,
itc_loss_config: Optional[CLYPLossConfig] = None,
) -> None:
# create contrastive loss function
if itc_loss_config:
raise NotImplementedError
else:
assert learn_temperature is not None
assert temperature_init is not None
self.itc_loss_fn = InfoNCELoss(
learn_temperature=learn_temperature,
init_temperature=temperature_init,
max_temperature=temperature_max,
min_temperature=temperature_min,
gather_with_grad=True,
)
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
pixel_values: Optional[torch.FloatTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
return_loss: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> tuple | CLYPOutput:
image_feats = self.vision_encoder(pixel_values)
text_feats = self.text_encoder(
{
"input_ids": input_ids,
"attention_mask": attention_mask,
"position_ids": position_ids,
}
)
loss = None
if return_loss:
loss = self.itc_loss_fn(image_feats, text_feats)
image_embeds = F.normalize(image_feats, dim=-1)
text_embeds = F.normalize(text_feats, dim=-1)
sim_i2t = image_embeds @ text_embeds.T
sim_t2i = text_embeds @ image_embeds.T
logits_per_image = sim_i2t / self.itc_loss_fn.temperature
logits_per_text = sim_t2i / self.itc_loss_fn.temperature
if not return_dict:
if loss is None:
return (logits_per_image, logits_per_text, text_embeds, image_embeds)
return (loss, logits_per_image, logits_per_text, text_embeds, image_embeds)
# TODO:
# - Support vision_model_output and text_model_output
# - Improve type: torch.Tensor -> torch.FloatTensor
return CLYPOutput(
loss=loss,
logits_per_image=logits_per_image, # type: ignore
logits_per_text=logits_per_text, # type: ignore
text_embeds=text_embeds, # type: ignore
image_embeds=image_embeds, # type: ignore
)
def get_text_features(
self,
input_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.Tensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> torch.FloatTensor:
text_feats = self.text_encoder(
{
"input_ids": input_ids,
"attention_mask": attention_mask,
"position_ids": position_ids,
}
)
return text_feats
def get_image_features(
self,
pixel_values: Optional[torch.FloatTensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> torch.FloatTensor:
image_feats = self.vision_encoder(pixel_values)
return image_feats
if __name__ == "__main__":
model = CLYPModel.from_pretrained(".")
|