103 lines
5.1 KiB
Python
103 lines
5.1 KiB
Python
import torch
|
|
from transformers import CLIPTextModel
|
|
from typing import Any, Callable, Dict, Optional, Tuple, Union, List
|
|
from transformers.modeling_outputs import BaseModelOutputWithPooling
|
|
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
|
|
|
|
|
|
class CLIPTextModelWrapper(CLIPTextModel):
|
|
# Adapted from https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/clip/modeling_clip.py#L812
|
|
# Modified to accept precomputed token embeddings "input_token_embs" as input or calculate them from input_ids and return them.
|
|
def forward(
|
|
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,
|
|
input_token_embs: Optional[torch.Tensor] = None,
|
|
return_token_embs: Optional[bool] = False,
|
|
) -> Union[Tuple, torch.Tensor, BaseModelOutputWithPooling]:
|
|
|
|
if return_token_embs:
|
|
return self.text_model.embeddings.token_embedding(input_ids)
|
|
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
output_attentions = output_attentions if output_attentions is not None else self.text_model.config.output_attentions
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else self.text_model.config.output_hidden_states
|
|
)
|
|
return_dict = return_dict if return_dict is not None else self.text_model.config.use_return_dict
|
|
|
|
if input_ids is None:
|
|
raise ValueError("You have to specify input_ids")
|
|
|
|
input_shape = input_ids.size()
|
|
input_ids = input_ids.view(-1, input_shape[-1])
|
|
|
|
hidden_states = self.text_model.embeddings(input_ids=input_ids, position_ids=position_ids, inputs_embeds=input_token_embs)
|
|
|
|
# CLIP's text model uses causal mask, prepare it here.
|
|
# https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
|
|
attn_mask_converter = AttentionMaskConverter(is_causal=True)
|
|
causal_attention_mask = attn_mask_converter.to_causal_4d(
|
|
batch_size=input_shape[0],
|
|
query_length=input_shape[1],
|
|
key_value_length=input_shape[1],
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device
|
|
)
|
|
# expand attention_mask
|
|
if attention_mask is not None:
|
|
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
|
attention_mask = attn_mask_converter.to_4d(
|
|
attention_mask,
|
|
input_shape[1],
|
|
key_value_length=input_shape[1],
|
|
dtype=hidden_states.dtype
|
|
)
|
|
|
|
encoder_outputs = self.text_model.encoder(
|
|
inputs_embeds=hidden_states,
|
|
attention_mask=attention_mask,
|
|
causal_attention_mask=causal_attention_mask,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
)
|
|
|
|
last_hidden_state = encoder_outputs[0]
|
|
last_hidden_state = self.text_model.final_layer_norm(last_hidden_state)
|
|
|
|
if self.text_model.eos_token_id == 2:
|
|
# The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here.
|
|
# A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added
|
|
# ------------------------------------------------------------
|
|
# text_embeds.shape = [batch_size, sequence_length, transformer.width]
|
|
# take features from the eot embedding (eot_token is the highest number in each sequence)
|
|
# casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
|
|
pooled_output = last_hidden_state[
|
|
torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
|
|
input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
|
|
]
|
|
else:
|
|
# The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible)
|
|
pooled_output = last_hidden_state[
|
|
torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
|
|
# We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
|
|
(input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.text_model.eos_token_id)
|
|
.int()
|
|
.argmax(dim=-1),
|
|
]
|
|
|
|
if not return_dict:
|
|
return (last_hidden_state, pooled_output) + encoder_outputs[1:]
|
|
|
|
return BaseModelOutputWithPooling(
|
|
last_hidden_state=last_hidden_state,
|
|
pooler_output=pooled_output,
|
|
hidden_states=encoder_outputs.hidden_states,
|
|
attentions=encoder_outputs.attentions,
|
|
) |