CheXagent-8b / modeling_chexagent.py

Upload 3 files

b09211c 6 months ago

No virus

55.7 kB

	# coding=utf-8
	# Copyright 2023 The CheXagent Authors, The Salesforce Authors and The HuggingFace Team. All rights reserved.
	#
	# 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.

	import math
	from dataclasses import dataclass
	from typing import Any, Optional, Tuple, Union

	import torch
	import torch.utils.checkpoint
	from einops import rearrange
	from torch import nn
	from torch.nn import CrossEntropyLoss
	from transformers.activations import ACT2FN
	from transformers.modeling_outputs import (
	BaseModelOutput,
	BaseModelOutputWithPastAndCrossAttentions,
	BaseModelOutputWithPooling,
	BaseModelOutputWithPoolingAndCrossAttentions,
	)
	from transformers.modeling_utils import PreTrainedModel
	from transformers.models.auto import AutoModelForCausalLM, AutoModelForSeq2SeqLM
	from transformers.pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
	from transformers.utils import ModelOutput, logging

	from .configuration_chexagent import CheXagentConfig, CheXagentQFormerConfig, CheXagentVisionConfig

	logger = logging.get_logger(__name__)


	@dataclass
	class CheXagentForConditionalGenerationModelOutput(ModelOutput):
	loss: Optional[Tuple[torch.FloatTensor]] = None
	logits: Optional[Tuple[torch.FloatTensor]] = None
	vision_outputs: Optional[torch.FloatTensor] = None
	qformer_outputs: Optional[Tuple[torch.FloatTensor]] = None
	language_model_outputs: Optional[Tuple[torch.FloatTensor]] = None

	def to_tuple(self) -> Tuple[Any]:
	return tuple(
	self[k]
	if k not in ["vision_outputs", "qformer_outputs", "language_model_outputs"]
	else getattr(self, k).to_tuple()
	for k in self.keys()
	)


	class CheXagentVisionEmbeddings(nn.Module):
	def __init__(self, config: CheXagentVisionConfig):
	super().__init__()
	self.config = config
	self.embed_dim = config.hidden_size
	self.image_size = config.image_size
	self.patch_size = config.patch_size

	self.class_embedding = nn.Parameter(torch.randn(1, 1, self.embed_dim))

	self.patch_embedding = nn.Conv2d(
	in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size
	)

	self.num_patches = (self.image_size // self.patch_size) ** 2
	self.num_positions = self.num_patches + 1

	self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim))

	def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
	batch_size = pixel_values.shape[0]
	target_dtype = self.patch_embedding.weight.dtype
	patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid]
	patch_embeds = patch_embeds.flatten(2).transpose(1, 2)

	class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype)
	embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
	embeddings = embeddings + self.position_embedding[:, : embeddings.size(1), :].to(target_dtype)
	return embeddings


	class CheXagentAttention(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.config = config
	self.embed_dim = config.hidden_size
	self.num_heads = config.num_attention_heads
	self.head_dim = self.embed_dim // self.num_heads
	if self.head_dim * self.num_heads != self.embed_dim:
	raise ValueError(
	f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
	f" {self.num_heads})."
	)
	self.scale = self.head_dim ** -0.5
	self.dropout = nn.Dropout(config.attention_dropout)

	# small tweak here compared to CLIP, no bias here
	self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=False)

	if config.qkv_bias:
	q_bias = nn.Parameter(torch.zeros(self.embed_dim))
	v_bias = nn.Parameter(torch.zeros(self.embed_dim))
	else:
	q_bias = None
	v_bias = None

	if q_bias is not None:
	qkv_bias = torch.cat((q_bias, torch.zeros_like(v_bias, requires_grad=False), v_bias))
	self.qkv.bias = nn.Parameter(qkv_bias)

	self.projection = nn.Linear(self.embed_dim, self.embed_dim)

	def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
	return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

	def forward(
	self,
	hidden_states: torch.Tensor,
	head_mask: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
	bsz, tgt_len, embed_dim = hidden_states.size()

	mixed_qkv = self.qkv(hidden_states)

	mixed_qkv = mixed_qkv.reshape(bsz, tgt_len, 3, self.num_heads, embed_dim // self.num_heads).permute(
	2, 0, 3, 1, 4
	)
	query_states, key_states, value_states = mixed_qkv[0], mixed_qkv[1], mixed_qkv[2]

	# Take the dot product between "query" and "key" to get the raw attention scores.
	attention_scores = torch.matmul(query_states, key_states.transpose(-1, -2))

	attention_scores = attention_scores * self.scale

	# Normalize the attention scores to probabilities.
	attention_probs = nn.functional.softmax(attention_scores, dim=-1)

	# This is actually dropping out entire tokens to attend to, which might
	# seem a bit unusual, but is taken from the original Transformer paper.
	attention_probs = self.dropout(attention_probs)

	# Mask heads if we want to
	if head_mask is not None:
	attention_probs = attention_probs * head_mask

	context_layer = torch.matmul(attention_probs, value_states).permute(0, 2, 1, 3)

	new_context_layer_shape = context_layer.size()[:-2] + (self.embed_dim,)
	context_layer = context_layer.reshape(new_context_layer_shape)

	output = self.projection(context_layer)

	outputs = (output, attention_probs) if output_attentions else (output, None)

	return outputs


	class CheXagentMLP(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.config = config
	self.activation_fn = ACT2FN[config.hidden_act]
	self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
	self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	hidden_states = self.fc1(hidden_states)
	hidden_states = self.activation_fn(hidden_states)
	hidden_states = self.fc2(hidden_states)
	return hidden_states


	class CheXagentEncoderLayer(nn.Module):
	def __init__(self, config: CheXagentConfig):
	super().__init__()
	self.embed_dim = config.hidden_size
	self.self_attn = CheXagentAttention(config)
	self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
	self.mlp = CheXagentMLP(config)
	self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: torch.Tensor,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.FloatTensor]:
	residual = hidden_states
	hidden_states = self.layer_norm1(hidden_states)
	hidden_states, attn_weights = self.self_attn(
	hidden_states=hidden_states,
	head_mask=attention_mask,
	output_attentions=output_attentions,
	)
	hidden_states = hidden_states + residual
	residual = hidden_states
	hidden_states = self.layer_norm2(hidden_states)
	hidden_states = self.mlp(hidden_states)

	hidden_states = hidden_states + residual

	outputs = (hidden_states,)

	if output_attentions:
	outputs += (attn_weights,)

	return outputs


	class CheXagentPreTrainedModel(PreTrainedModel):
	config_class = CheXagentConfig
	base_model_prefix = "chexagent"
	supports_gradient_checkpointing = True
	_no_split_modules = [
	"CheXagentQFormerEmbeddings",
	"CheXagentAttention",
	"CheXagentQFormerMultiHeadAttention",
	"CheXagentQFormerSelfOutput",
	]
	_keep_in_fp32_modules = []

	def _init_weights(self, module):
	"""Initialize the weights"""
	factor = self.config.initializer_range
	if isinstance(module, nn.Conv2d) or isinstance(module, nn.Embedding) or isinstance(module, nn.Linear):
	module.weight.data.normal_(mean=0.0, std=factor)
	if hasattr(module, "bias") and module.bias is not None:
	module.bias.data.zero_()

	if isinstance(module, CheXagentVisionEmbeddings):
	if hasattr(self.config, "vision_config"):
	factor = self.config.vision_config.initializer_range
	nn.init.trunc_normal_(module.position_embedding, mean=0.0, std=factor)
	nn.init.trunc_normal_(module.class_embedding, mean=0.0, std=factor)

	elif isinstance(module, nn.LayerNorm):
	module.bias.data.zero_()
	module.weight.data.fill_(1.0)
	elif isinstance(module, nn.Linear) and module.bias is not None:
	module.bias.data.zero_()


	class CheXagentEncoder(nn.Module):
	def __init__(self, config: CheXagentConfig):
	super().__init__()
	self.config = config
	self.layers = nn.ModuleList([CheXagentEncoderLayer(config) for _ in range(config.num_hidden_layers)])
	self.gradient_checkpointing = False

	def forward(
	self,
	inputs_embeds,
	attention_mask: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutput]:
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	encoder_states = () if output_hidden_states else None
	all_attentions = () if output_attentions else None
	hidden_states = inputs_embeds
	for idx, encoder_layer in enumerate(self.layers):
	if output_hidden_states:
	encoder_states = encoder_states + (hidden_states,)
	if self.gradient_checkpointing and self.training:
	layer_outputs = self._gradient_checkpointing_func(
	encoder_layer.__call__,
	hidden_states,
	attention_mask,
	output_attentions,
	)
	else:
	layer_outputs = encoder_layer(hidden_states, attention_mask, output_attentions=output_attentions, )

	hidden_states = layer_outputs[0]

	if output_attentions:
	all_attentions = all_attentions + (layer_outputs[1],)

	if output_hidden_states:
	encoder_states = encoder_states + (hidden_states,)

	if not return_dict:
	return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
	return BaseModelOutput(
	last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
	)


	class CheXagentVisionModel(CheXagentPreTrainedModel):
	main_input_name = "pixel_values"
	config_class = CheXagentVisionConfig

	def __init__(self, config: CheXagentVisionConfig):
	super().__init__(config)
	self.config = config
	embed_dim = config.hidden_size

	self.embeddings = CheXagentVisionEmbeddings(config)
	self.encoder = CheXagentEncoder(config)
	self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

	self.post_init()

	def forward(
	self,
	pixel_values: Optional[torch.FloatTensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPooling]:
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	if pixel_values is None:
	raise ValueError("You have to specify pixel_values")
	hidden_states = self.embeddings(pixel_values)

	encoder_outputs = self.encoder(
	inputs_embeds=hidden_states,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	last_hidden_state = encoder_outputs[0]
	last_hidden_state = self.post_layernorm(last_hidden_state)

	pooled_output = last_hidden_state[:, 0, :]
	pooled_output = self.post_layernorm(pooled_output)

	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,
	)

	def get_input_embeddings(self):
	return self.embeddings


	class CheXagentQFormerMultiHeadAttention(nn.Module):
	def __init__(self, config, is_cross_attention=False):
	super().__init__()
	self.config = config
	if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
	raise ValueError(
	"The hidden size (%d) is not a multiple of the number of attention heads (%d)"
	% (config.hidden_size, config.num_attention_heads)
	)

	self.num_attention_heads = config.num_attention_heads
	self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
	self.all_head_size = self.num_attention_heads * self.attention_head_size

	self.query = nn.Linear(config.hidden_size, self.all_head_size)
	if is_cross_attention:
	self.key = nn.Linear(config.encoder_hidden_size, self.all_head_size)
	self.value = nn.Linear(config.encoder_hidden_size, self.all_head_size)
	else:
	self.key = nn.Linear(config.hidden_size, self.all_head_size)
	self.value = nn.Linear(config.hidden_size, self.all_head_size)

	self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
	self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
	if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
	self.max_position_embeddings = config.max_position_embeddings
	self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
	self.save_attention = False

	def save_attn_gradients(self, attn_gradients):
	self.attn_gradients = attn_gradients

	def get_attn_gradients(self):
	return self.attn_gradients

	def save_attention_map(self, attention_map):
	self.attention_map = attention_map

	def get_attention_map(self):
	return self.attention_map

	def transpose_for_scores(self, x):
	new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
	x = x.view(*new_x_shape)
	return x.permute(0, 2, 1, 3)

	def forward(
	self,
	hidden_states,
	attention_mask=None,
	head_mask=None,
	encoder_hidden_states=None,
	encoder_attention_mask=None,
	past_key_value=None,
	output_attentions=False,
	):
	# If this is instantiated as a cross-attention module, the keys
	# and values come from an encoder; the attention mask needs to be
	# such that the encoder's padding tokens are not attended to.
	is_cross_attention = encoder_hidden_states is not None

	if is_cross_attention:
	key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
	value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
	attention_mask = encoder_attention_mask
	elif past_key_value is not None:
	key_layer = self.transpose_for_scores(self.key(hidden_states))
	value_layer = self.transpose_for_scores(self.value(hidden_states))
	key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
	value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
	else:
	key_layer = self.transpose_for_scores(self.key(hidden_states))
	value_layer = self.transpose_for_scores(self.value(hidden_states))

	mixed_query_layer = self.query(hidden_states)

	query_layer = self.transpose_for_scores(mixed_query_layer)

	past_key_value = (key_layer, value_layer)

	# Take the dot product between "query" and "key" to get the raw attention scores.
	attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

	if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
	seq_length = hidden_states.size()[1]
	position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
	position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
	distance = position_ids_l - position_ids_r
	positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
	positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility

	if self.position_embedding_type == "relative_key":
	relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
	attention_scores = attention_scores + relative_position_scores
	elif self.position_embedding_type == "relative_key_query":
	relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
	relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
	attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key

	attention_scores = attention_scores / math.sqrt(self.attention_head_size)

	if attention_mask is not None:
	# Apply the attention mask is (precomputed for all layers in BertModel forward() function)
	attention_scores = attention_scores + attention_mask

	# Normalize the attention scores to probabilities.
	attention_probs = nn.Softmax(dim=-1)(attention_scores)

	if is_cross_attention and self.save_attention:
	self.save_attention_map(attention_probs)
	attention_probs.register_hook(self.save_attn_gradients)

	# This is actually dropping out entire tokens to attend to, which might
	# seem a bit unusual, but is taken from the original Transformer paper.
	attention_probs_dropped = self.dropout(attention_probs)

	# Mask heads if we want to
	if head_mask is not None:
	attention_probs_dropped = attention_probs_dropped * head_mask

	context_layer = torch.matmul(attention_probs_dropped, value_layer)

	context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
	new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
	context_layer = context_layer.view(*new_context_layer_shape)

	outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

	outputs = outputs + (past_key_value,)
	return outputs


	class CheXagentQFormerSelfOutput(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.dense = nn.Linear(config.hidden_size, config.hidden_size)
	self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)

	def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
	hidden_states = self.dense(hidden_states)
	hidden_states = self.dropout(hidden_states)
	hidden_states = self.LayerNorm(hidden_states + input_tensor)
	return hidden_states


	class CheXagentQFormerAttention(nn.Module):
	def __init__(self, config, is_cross_attention=False):
	super().__init__()
	self.attention = CheXagentQFormerMultiHeadAttention(config, is_cross_attention)
	self.output = CheXagentQFormerSelfOutput(config)
	self.pruned_heads = set()

	def prune_heads(self, heads):
	if len(heads) == 0:
	return
	heads, index = find_pruneable_heads_and_indices(
	heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
	)

	# Prune linear layers
	self.attention.query = prune_linear_layer(self.attention.query, index)
	self.attention.key = prune_linear_layer(self.attention.key, index)
	self.attention.value = prune_linear_layer(self.attention.value, index)
	self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

	# Update hyper params and store pruned heads
	self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
	self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
	self.pruned_heads = self.pruned_heads.union(heads)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.FloatTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	encoder_hidden_states: Optional[torch.FloatTensor] = None,
	encoder_attention_mask: Optional[torch.FloatTensor] = None,
	past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.Tensor]:
	self_outputs = self.attention(
	hidden_states,
	attention_mask,
	head_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	past_key_value,
	output_attentions,
	)
	attention_output = self.output(self_outputs[0], hidden_states)
	outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
	return outputs


	class CheXagentQFormerIntermediate(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
	if isinstance(config.hidden_act, str):
	self.intermediate_act_fn = ACT2FN[config.hidden_act]
	else:
	self.intermediate_act_fn = config.hidden_act

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	hidden_states = self.dense(hidden_states)
	hidden_states = self.intermediate_act_fn(hidden_states)
	return hidden_states


	class CheXagentQFormerOutput(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
	self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)

	def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
	hidden_states = self.dense(hidden_states)
	hidden_states = self.dropout(hidden_states)
	hidden_states = self.LayerNorm(hidden_states + input_tensor)
	return hidden_states


	class CheXagentQFormerLayer(nn.Module):
	def __init__(self, config, layer_idx):
	super().__init__()
	self.chunk_size_feed_forward = config.chunk_size_feed_forward
	self.seq_len_dim = 1
	self.attention = CheXagentQFormerAttention(config)

	self.layer_idx = layer_idx

	if layer_idx % config.cross_attention_frequency == 0:
	self.crossattention = CheXagentQFormerAttention(config, is_cross_attention=True)
	self.has_cross_attention = True
	else:
	self.has_cross_attention = False

	self.intermediate_query = CheXagentQFormerIntermediate(config)
	self.output_query = CheXagentQFormerOutput(config)

	def forward(
	self,
	hidden_states,
	attention_mask=None,
	head_mask=None,
	encoder_hidden_states=None,
	encoder_attention_mask=None,
	past_key_value=None,
	output_attentions=False,
	query_length=0,
	):
	# decoder uni-directional self-attention cached key/values tuple is at positions 1,2
	self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
	self_attention_outputs = self.attention(
	hidden_states,
	attention_mask,
	head_mask,
	output_attentions=output_attentions,
	past_key_value=self_attn_past_key_value,
	)
	attention_output = self_attention_outputs[0]
	outputs = self_attention_outputs[1:-1]

	present_key_value = self_attention_outputs[-1]

	if query_length > 0:
	query_attention_output = attention_output[:, :query_length, :]

	if self.has_cross_attention:
	if encoder_hidden_states is None:
	raise ValueError("encoder_hidden_states must be given for cross-attention layers")
	cross_attention_outputs = self.crossattention(
	query_attention_output,
	attention_mask,
	head_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	output_attentions=output_attentions,
	)
	query_attention_output = cross_attention_outputs[0]
	# add cross attentions if we output attention weights
	outputs = outputs + cross_attention_outputs[1:-1]

	layer_output = apply_chunking_to_forward(
	self.feed_forward_chunk_query,
	self.chunk_size_feed_forward,
	self.seq_len_dim,
	query_attention_output,
	)

	if attention_output.shape[1] > query_length:
	layer_output_text = apply_chunking_to_forward(
	self.feed_forward_chunk,
	self.chunk_size_feed_forward,
	self.seq_len_dim,
	attention_output[:, query_length:, :],
	)
	layer_output = torch.cat([layer_output, layer_output_text], dim=1)
	else:
	layer_output = apply_chunking_to_forward(
	self.feed_forward_chunk,
	self.chunk_size_feed_forward,
	self.seq_len_dim,
	attention_output,
	)
	outputs = (layer_output,) + outputs

	outputs = outputs + (present_key_value,)

	return outputs

	def feed_forward_chunk(self, attention_output):
	intermediate_output = self.intermediate(attention_output)
	layer_output = self.output(intermediate_output, attention_output)
	return layer_output

	def feed_forward_chunk_query(self, attention_output):
	intermediate_output = self.intermediate_query(attention_output)
	layer_output = self.output_query(intermediate_output, attention_output)
	return layer_output


	class CheXagentQFormerEncoder(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.config = config
	self.layer = nn.ModuleList(
	[CheXagentQFormerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
	)
	self.gradient_checkpointing = False

	def forward(
	self,
	hidden_states,
	attention_mask=None,
	head_mask=None,
	encoder_hidden_states=None,
	encoder_attention_mask=None,
	past_key_values=None,
	use_cache=None,
	output_attentions=False,
	output_hidden_states=False,
	return_dict=True,
	query_length=0,
	):
	all_hidden_states = () if output_hidden_states else None
	all_self_attentions = () if output_attentions else None
	all_cross_attentions = () if output_attentions else None

	next_decoder_cache = () if use_cache else None

	for i in range(self.config.num_hidden_layers):
	layer_module = self.layer[i]
	if output_hidden_states:
	all_hidden_states = all_hidden_states + (hidden_states,)

	layer_head_mask = head_mask[i] if head_mask is not None else None
	past_key_value = past_key_values[i] if past_key_values is not None else None

	if getattr(self.config, "gradient_checkpointing", False) and self.training:
	if use_cache:
	logger.warning(
	"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
	)
	use_cache = False
	layer_outputs = self._gradient_checkpointing_func(
	layer_module.__call__,
	hidden_states,
	attention_mask,
	layer_head_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	)
	else:
	layer_outputs = layer_module(
	hidden_states,
	attention_mask,
	layer_head_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	past_key_value,
	output_attentions,
	query_length,
	)

	hidden_states = layer_outputs[0]
	if use_cache:
	next_decoder_cache += (layer_outputs[-1],)
	if output_attentions:
	all_self_attentions = all_self_attentions + (layer_outputs[1],)
	if layer_module.has_cross_attention:
	all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

	if output_hidden_states:
	all_hidden_states = all_hidden_states + (hidden_states,)

	if not return_dict:
	return tuple(
	v
	for v in [
	hidden_states,
	next_decoder_cache,
	all_hidden_states,
	all_self_attentions,
	all_cross_attentions,
	]
	if v is not None
	)
	return BaseModelOutputWithPastAndCrossAttentions(
	last_hidden_state=hidden_states,
	past_key_values=next_decoder_cache,
	hidden_states=all_hidden_states,
	attentions=all_self_attentions,
	cross_attentions=all_cross_attentions,
	)


	class CheXagentQFormerEmbeddings(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
	self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)

	self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)

	# position_ids (1, len position emb) is contiguous in memory and exported when serialized
	self.register_buffer(
	"position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
	)
	self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

	self.config = config

	def forward(
	self,
	input_ids=None,
	position_ids=None,
	query_embeds=None,
	past_key_values_length=0,
	):
	if input_ids is not None:
	seq_length = input_ids.size()[1]
	else:
	seq_length = 0

	if position_ids is None:
	position_ids = self.position_ids[:, past_key_values_length: seq_length + past_key_values_length].clone()

	if input_ids is not None:
	embeddings = self.word_embeddings(input_ids)
	if self.position_embedding_type == "absolute":
	position_embeddings = self.position_embeddings(position_ids.to(embeddings.device))
	embeddings = embeddings + position_embeddings

	if query_embeds is not None:
	embeddings = torch.cat((query_embeds, embeddings), dim=1)
	else:
	embeddings = query_embeds

	embeddings = embeddings.to(self.layernorm.weight.dtype)
	embeddings = self.layernorm(embeddings)
	embeddings = self.dropout(embeddings)
	return embeddings


	class CheXagentQFormerEncoder(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.config = config
	self.layer = nn.ModuleList(
	[CheXagentQFormerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
	)
	self.gradient_checkpointing = False

	def forward(
	self,
	hidden_states,
	attention_mask=None,
	head_mask=None,
	encoder_hidden_states=None,
	encoder_attention_mask=None,
	past_key_values=None,
	use_cache=None,
	output_attentions=False,
	output_hidden_states=False,
	return_dict=True,
	query_length=0,
	):
	all_hidden_states = () if output_hidden_states else None
	all_self_attentions = () if output_attentions else None
	all_cross_attentions = () if output_attentions else None

	next_decoder_cache = () if use_cache else None

	for i in range(self.config.num_hidden_layers):
	layer_module = self.layer[i]
	if output_hidden_states:
	all_hidden_states = all_hidden_states + (hidden_states,)

	layer_head_mask = head_mask[i] if head_mask is not None else None
	past_key_value = past_key_values[i] if past_key_values is not None else None

	if getattr(self.config, "gradient_checkpointing", False) and self.training:
	if use_cache:
	logger.warning(
	"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
	)
	use_cache = False
	layer_outputs = self._gradient_checkpointing_func(
	layer_module.__call__,
	hidden_states,
	attention_mask,
	layer_head_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	)
	else:
	layer_outputs = layer_module(
	hidden_states,
	attention_mask,
	layer_head_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	past_key_value,
	output_attentions,
	query_length,
	)

	hidden_states = layer_outputs[0]
	if use_cache:
	next_decoder_cache += (layer_outputs[-1],)
	if output_attentions:
	all_self_attentions = all_self_attentions + (layer_outputs[1],)
	if layer_module.has_cross_attention:
	all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

	if output_hidden_states:
	all_hidden_states = all_hidden_states + (hidden_states,)

	if not return_dict:
	return tuple(
	v
	for v in [
	hidden_states,
	next_decoder_cache,
	all_hidden_states,
	all_self_attentions,
	all_cross_attentions,
	]
	if v is not None
	)
	return BaseModelOutputWithPastAndCrossAttentions(
	last_hidden_state=hidden_states,
	past_key_values=next_decoder_cache,
	hidden_states=all_hidden_states,
	attentions=all_self_attentions,
	cross_attentions=all_cross_attentions,
	)


	class CheXagentQFormerModel(CheXagentPreTrainedModel):
	def __init__(self, config: CheXagentQFormerConfig):
	super().__init__(config)
	self.config = config

	self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)

	self.encoder = CheXagentQFormerEncoder(config)

	self.post_init()

	def get_input_embeddings(self):
	return self.embeddings.word_embeddings

	def set_input_embeddings(self, value):
	self.embeddings.word_embeddings = value

	def _prune_heads(self, heads_to_prune):
	for layer, heads in heads_to_prune.items():
	self.encoder.layer[layer].attention.prune_heads(heads)

	def get_extended_attention_mask(
	self,
	attention_mask: torch.Tensor,
	input_shape: Tuple[int],
	device: torch.device,
	has_query: bool = False,
	) -> torch.Tensor:
	# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
	# ourselves in which case we just need to make it broadcastable to all heads.
	if attention_mask.dim() == 3:
	extended_attention_mask = attention_mask[:, None, :, :]
	elif attention_mask.dim() == 2:
	# Provided a padding mask of dimensions [batch_size, seq_length]
	# - the model is an encoder, so make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
	extended_attention_mask = attention_mask[:, None, None, :]
	else:
	raise ValueError(
	"Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
	input_shape, attention_mask.shape
	)
	)

	# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
	# masked positions, this operation will create a tensor which is 0.0 for
	# positions we want to attend and -10000.0 for masked positions.
	# Since we are adding it to the raw scores before the softmax, this is
	# effectively the same as removing these entirely.
	extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility
	extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
	return extended_attention_mask

	def forward(
	self,
	query_embeds: torch.FloatTensor,
	attention_mask: Optional[torch.FloatTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	encoder_hidden_states: Optional[torch.FloatTensor] = None,
	encoder_attention_mask: Optional[torch.FloatTensor] = None,
	past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	# past_key_values_length
	past_key_values_length = (
	past_key_values[0][0].shape[2] - self.config.query_length if past_key_values is not None else 0
	)

	query_length = query_embeds.shape[1] if query_embeds is not None else 0

	embedding_output = self.layernorm(query_embeds)
	embedding_output = self.dropout(embedding_output)

	input_shape = embedding_output.size()[:-1]
	batch_size, seq_length = input_shape
	device = embedding_output.device

	if attention_mask is None:
	attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)

	# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
	# ourselves in which case we just need to make it broadcastable to all heads.
	extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, device)

	# If a 2D or 3D attention mask is provided for the cross-attention
	# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
	if encoder_hidden_states is not None:
	if type(encoder_hidden_states) == list:
	encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
	else:
	encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
	encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)

	if type(encoder_attention_mask) == list:
	encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
	elif encoder_attention_mask is None:
	encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
	encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
	else:
	encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
	else:
	encoder_extended_attention_mask = None

	# Prepare head mask if needed
	# 1.0 in head_mask indicate we keep the head
	# attention_probs has shape bsz x n_heads x N x N
	# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
	# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
	head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

	encoder_outputs = self.encoder(
	embedding_output,
	attention_mask=extended_attention_mask,
	head_mask=head_mask,
	encoder_hidden_states=encoder_hidden_states,
	encoder_attention_mask=encoder_extended_attention_mask,
	past_key_values=past_key_values,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	query_length=query_length,
	)
	sequence_output = encoder_outputs[0]
	pooled_output = sequence_output[:, 0, :]

	if not return_dict:
	return (sequence_output, pooled_output) + encoder_outputs[1:]

	return BaseModelOutputWithPoolingAndCrossAttentions(
	last_hidden_state=sequence_output,
	pooler_output=pooled_output,
	past_key_values=encoder_outputs.past_key_values,
	hidden_states=encoder_outputs.hidden_states,
	attentions=encoder_outputs.attentions,
	cross_attentions=encoder_outputs.cross_attentions,
	)


	class CheXagentForConditionalGeneration(CheXagentPreTrainedModel):
	config_class = CheXagentConfig
	main_input_name = "pixel_values"

	def __init__(self, config: CheXagentConfig):
	super().__init__(config)

	self.vision_model = CheXagentVisionModel(config.vision_config)

	self.query_tokens = nn.Parameter(torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size))
	self.qformer = CheXagentQFormerModel(config.qformer_config)

	self.language_projection = nn.Linear(config.qformer_config.hidden_size, config.text_config.hidden_size)
	if config.use_decoder_only_language_model:
	language_model = AutoModelForCausalLM.from_config(config.text_config)
	else:
	language_model = AutoModelForSeq2SeqLM.from_config(config.text_config)

	# Update _tied_weights_keys using the base model used.
	if language_model._tied_weights_keys is not None:
	self._tied_weights_keys = [f"language_model.{k}" for k in language_model._tied_weights_keys]

	self.language_model = language_model

	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self):
	return self.language_model.get_input_embeddings()

	def set_input_embeddings(self, value):
	self.language_model.set_input_embeddings(value)

	def set_output_embeddings(self, new_embeddings):
	self.language_model.set_output_embeddings(new_embeddings)

	def get_output_embeddings(self) -> nn.Module:
	return self.language_model.get_output_embeddings()

	def get_encoder(self):
	return self.language_model.get_encoder()

	def get_decoder(self):
	return self.language_model.get_decoder()

	def _tie_weights(self):
	if not self.config.use_decoder_only_language_model:
	self.language_model.encoder.embed_tokens = self.language_model.shared
	self.language_model.decoder.embed_tokens = self.language_model.shared

	def _preprocess_accelerate(self):
	hf_device_map = self.hf_device_map
	if len(hf_device_map) > 1 and "language_model" not in hf_device_map and torch.cuda.device_count() > 1:
	# warn users about unexpected behavior when using multi-GPU + BLIP-2 + `accelerate`.
	logger.warning(
	"The `language_model` is not in the `hf_device_map` dictionary and you are running your script"
	" in a multi-GPU environment. this may lead to unexpected behavior when using `accelerate`."
	" Please pass a `device_map` that contains `language_model` to remove this warning."
	" Please refer to https://github.com/huggingface/blog/blob/main/accelerate-large-models.md for"
	" more details on creating a `device_map` for large models.",
	)
	if hasattr(self.language_model, "_hf_hook"):
	self.language_model._hf_hook.io_same_device = True # For `generate` compatibility

	def forward(
	self,
	pixel_values: torch.FloatTensor = None,
	input_ids: torch.FloatTensor = None,
	attention_mask: Optional[torch.LongTensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	labels: Optional[torch.LongTensor] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, CheXagentForConditionalGenerationModelOutput]:
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	vision_outputs, query_outputs = None, None
	if pixel_values is not None:
	# step 1: forward the images through the vision encoder,
	# to get image embeddings of shape (batch_size, seq_len, hidden_size)
	image_mask = pixel_values.sum(dim=(2, 3, 4)) != 0
	vision_outputs = self.vision_model(
	pixel_values=pixel_values[image_mask],
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)
	tmp = vision_outputs[0]
	image_embeds = tmp.new_zeros((image_mask.shape, tmp.shape[1:]))
	image_embeds[image_mask] = tmp

	# step 2: forward the query tokens through the QFormer, using the image embeddings for cross-attention
	image_attention_mask = torch.zeros(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
	image_attention_mask[image_mask] = 1

	image_embeds = rearrange(image_embeds, "b i n d -> b (i n) d")
	image_attention_mask = rearrange(image_attention_mask, "b i n -> b (i n)")
	query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
	query_outputs = self.qformer(
	query_embeds=query_tokens,
	encoder_hidden_states=image_embeds,
	encoder_attention_mask=image_attention_mask,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)
	query_output = query_outputs[0]

	# step 3: project vision to language
	input_vis = self.language_projection(query_output)
	vis_atts = torch.ones(input_vis.size()[:-1], dtype=torch.long, device=input_vis.device)

	# step 4: get the embeddings of the prompt
	inputs_lang = self.language_model.get_input_embeddings()(input_ids)
	lang_atts = attention_mask
	if lang_atts is None:
	lang_atts = torch.ones_like(input_ids)

	# step 5: conditioned on the images and/or prompts
	if pixel_values is not None:
	inputs_embeds = torch.cat([input_vis, inputs_lang], dim=1)
	attention_mask = torch.cat([vis_atts, lang_atts], dim=1)
	else:
	inputs_embeds = inputs_lang
	attention_mask = lang_atts

	outputs = self.language_model(
	inputs_embeds=inputs_embeds,
	attention_mask=attention_mask,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict
	)
	logits = outputs.logits if return_dict else outputs[0]

	loss = None
	# we compute the loss here since we need to take into account the sequence length of the query embeds
	if labels is not None:
	# make target
	empty_labels = torch.ones(vis_atts.size(), dtype=torch.long, device=input_ids.device).fill_(-100)
	labels = torch.cat([empty_labels, labels], dim=1)
	labels = labels.to(logits.device)
	logits = logits[:, -labels.size(1):, :]
	# Shift so that tokens < n predict n
	shift_logits = logits[..., :-1, :].contiguous()
	shift_labels = labels[..., 1:].contiguous().to(logits.device)
	# Flatten the tokens
	loss_fct = CrossEntropyLoss(reduction="mean")
	loss = loss_fct(shift_logits.view(-1, self.config.text_config.vocab_size), shift_labels.view(-1))

	if not return_dict:
	output = (logits, vision_outputs, query_outputs, outputs)
	return ((loss,) + output) if loss is not None else output

	return CheXagentForConditionalGenerationModelOutput(
	loss=loss,
	logits=logits,
	vision_outputs=vision_outputs,
	qformer_outputs=query_outputs,
	language_model_outputs=outputs,
	)

	@torch.no_grad()
	def generate(
	self,
	pixel_values: torch.FloatTensor = None,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.LongTensor] = None,
	**generate_kwargs,
	) -> torch.LongTensor:
	if hasattr(self, "hf_device_map"):
	# preprocess for `accelerate`
	self._preprocess_accelerate()

	batch_size = pixel_values.shape[0] if pixel_values is not None else input_ids.shape[0]
	if pixel_values is not None:
	# step 1: forward the images through the vision encoder
	image_mask = pixel_values.sum(dim=(2, 3, 4)) != 0
	vision_outputs = self.vision_model(pixel_values[image_mask], return_dict=True)
	tmp = vision_outputs[0]
	image_embeds = tmp.new_zeros((image_mask.shape, tmp.shape[1:]))
	image_embeds[image_mask] = tmp

	# step 2: forward the query tokens through the QFormer, using the image embeddings for cross-attention
	image_attention_mask = torch.zeros(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
	image_attention_mask[image_mask] = 1
	image_embeds = rearrange(image_embeds, "b i n d -> b (i n) d")
	image_attention_mask = rearrange(image_attention_mask, "b i n -> b (i n)")
	query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
	query_outputs = self.qformer(
	query_embeds=query_tokens,
	encoder_hidden_states=image_embeds,
	encoder_attention_mask=image_attention_mask,
	return_dict=True,
	)
	query_output = query_outputs.last_hidden_state

	# step 3: project vision to language
	input_vis = self.language_projection(query_output)
	vis_atts = torch.ones(input_vis.size()[:-1], dtype=torch.long, device=input_vis.device)

	# step 4: get the embeddings of the prompt
	if input_ids is None:
	input_ids = (
	torch.LongTensor([[self.config.text_config.bos_token_id]])
	.repeat(batch_size, 1)
	.to(next(self.parameters()).device)
	)
	inputs_lang = self.language_model.get_input_embeddings()(input_ids)
	lang_atts = attention_mask
	if lang_atts is None:
	lang_atts = torch.ones_like(input_ids)

	# step 5: conditioned on the images and/or prompts
	if pixel_values is not None:
	inputs_embeds = torch.cat([input_vis, inputs_lang], dim=1)
	attention_mask = torch.cat([vis_atts, lang_atts], dim=1)
	else:
	inputs_embeds = inputs_lang
	attention_mask = lang_atts

	outputs = self.language_model.generate(
	inputs_embeds=inputs_embeds,
	attention_mask=attention_mask,
	**generate_kwargs,
	)
	return outputs