968 lines
44 KiB
Python
968 lines
44 KiB
Python
|
# coding=utf-8
|
||
|
|
# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. 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.
|
||
|
|
""" PyTorch TrOCR decoder model (based on RoBERTa)."""
|
||
|
|
|
||
|
|
|
||
|
|
import copy
|
||
|
|
import math
|
||
|
|
from typing import Optional, Tuple, Union
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
from torch.nn import CrossEntropyLoss
|
||
|
|
|
||
|
|
from ...activations import ACT2FN
|
||
|
|
from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask
|
||
|
|
from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions
|
||
|
|
from ...modeling_utils import PreTrainedModel
|
||
|
|
from ...utils import add_start_docstrings, logging, replace_return_docstrings
|
||
|
|
from .configuration_trocr import TrOCRConfig
|
||
|
|
|
||
|
|
|
||
|
|
logger = logging.get_logger(__name__)
|
||
|
|
|
||
|
|
_CONFIG_FOR_DOC = "TrOCRConfig"
|
||
|
|
_CHECKPOINT_FOR_DOC = "microsoft/trocr-base-handwritten"
|
||
|
|
|
||
|
|
|
||
|
|
from ..deprecated._archive_maps import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
|
||
|
|
|
||
|
|
|
||
|
|
# Copied from transformers.models.bart.modeling_bart.BartLearnedPositionalEmbedding with Bart->TrOCR
|
||
|
|
class TrOCRLearnedPositionalEmbedding(nn.Embedding):
|
||
|
|
"""
|
||
|
|
This module learns positional embeddings up to a fixed maximum size.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, num_embeddings: int, embedding_dim: int):
|
||
|
|
# TrOCR is set up so that if padding_idx is specified then offset the embedding ids by 2
|
||
|
|
# and adjust num_embeddings appropriately. Other models don't have this hack
|
||
|
|
self.offset = 2
|
||
|
|
super().__init__(num_embeddings + self.offset, embedding_dim)
|
||
|
|
|
||
|
|
def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):
|
||
|
|
"""`input_ids' shape is expected to be [bsz x seqlen]."""
|
||
|
|
|
||
|
|
bsz, seq_len = input_ids.shape[:2]
|
||
|
|
positions = torch.arange(
|
||
|
|
past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
|
||
|
|
).expand(bsz, -1)
|
||
|
|
|
||
|
|
return super().forward(positions + self.offset)
|
||
|
|
|
||
|
|
|
||
|
|
class TrOCRSinusoidalPositionalEmbedding(nn.Module):
|
||
|
|
"""This module produces sinusoidal positional embeddings of any length."""
|
||
|
|
|
||
|
|
def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None):
|
||
|
|
super().__init__()
|
||
|
|
self.offset = 2
|
||
|
|
self.embedding_dim = embedding_dim
|
||
|
|
self.padding_idx = padding_idx
|
||
|
|
self.weights = self.get_embedding(num_positions, embedding_dim, padding_idx)
|
||
|
|
self.register_buffer("_float_tensor", torch.FloatTensor(1))
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None):
|
||
|
|
"""
|
||
|
|
Build sinusoidal embeddings. This matches the implementation in tensor2tensor, but differs slightly from the
|
||
|
|
description in Section 3.5 of "Attention Is All You Need".
|
||
|
|
"""
|
||
|
|
half_dim = embedding_dim // 2
|
||
|
|
emb = math.log(10000) / (half_dim - 1)
|
||
|
|
emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb)
|
||
|
|
emb = torch.arange(num_embeddings, dtype=torch.int64).float().unsqueeze(1) * emb.unsqueeze(0)
|
||
|
|
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
|
||
|
|
if embedding_dim % 2 == 1:
|
||
|
|
# zero pad
|
||
|
|
emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
|
||
|
|
if padding_idx is not None:
|
||
|
|
emb[padding_idx, :] = 0
|
||
|
|
|
||
|
|
return emb.to(torch.get_default_dtype())
|
||
|
|
|
||
|
|
@torch.no_grad()
|
||
|
|
def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):
|
||
|
|
bsz, seq_len = input_ids.size()
|
||
|
|
# Create the position ids from the input token ids. Any padded tokens remain padded.
|
||
|
|
position_ids = self.create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length).to(
|
||
|
|
input_ids.device
|
||
|
|
)
|
||
|
|
|
||
|
|
# expand embeddings if needed
|
||
|
|
max_pos = self.padding_idx + 1 + seq_len
|
||
|
|
if self.weights is None or max_pos > self.weights.size(0):
|
||
|
|
# recompute/expand embeddings if needed
|
||
|
|
self.weights = self.get_embedding(max_pos, self.embedding_dim, self.padding_idx)
|
||
|
|
self.weights = self.weights.to(self._float_tensor)
|
||
|
|
|
||
|
|
x = self.weights.index_select(0, position_ids.view(-1)).view(bsz, seq_len, -1).detach()
|
||
|
|
|
||
|
|
return x
|
||
|
|
|
||
|
|
def create_position_ids_from_input_ids(
|
||
|
|
self, input_ids: torch.Tensor, padding_idx: int, past_key_values_length: Optional[int] = 0
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding
|
||
|
|
symbols are ignored. This is modified from fairseq's `utils.make_positions`.
|
||
|
|
"""
|
||
|
|
# The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
|
||
|
|
mask = input_ids.ne(padding_idx).int()
|
||
|
|
incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
|
||
|
|
return incremental_indices.long() + padding_idx
|
||
|
|
|
||
|
|
|
||
|
|
class TrOCRAttention(nn.Module):
|
||
|
|
"""Multi-headed attention from 'Attention Is All You Need' paper."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
config,
|
||
|
|
embed_dim: int,
|
||
|
|
num_heads: int,
|
||
|
|
kdim: int = None,
|
||
|
|
vdim: int = None,
|
||
|
|
dropout: float = 0.0,
|
||
|
|
is_decoder: bool = False,
|
||
|
|
bias: bool = True,
|
||
|
|
is_cross_attention: bool = False,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
self.embed_dim = embed_dim
|
||
|
|
self.kdim = kdim if kdim is not None else embed_dim
|
||
|
|
self.vdim = vdim if vdim is not None else embed_dim
|
||
|
|
self.num_heads = num_heads
|
||
|
|
self.dropout = dropout
|
||
|
|
self.head_dim = embed_dim // num_heads
|
||
|
|
if not (self.head_dim * 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" {num_heads})."
|
||
|
|
)
|
||
|
|
self.scaling = self.head_dim**-0.5
|
||
|
|
self.is_decoder = is_decoder
|
||
|
|
|
||
|
|
self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)
|
||
|
|
self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
|
||
|
|
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
||
|
|
|
||
|
|
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
||
|
|
|
||
|
|
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,
|
||
|
|
key_value_states: Optional[torch.Tensor] = None,
|
||
|
|
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||
|
|
attention_mask: Optional[torch.Tensor] = None,
|
||
|
|
layer_head_mask: Optional[torch.Tensor] = None,
|
||
|
|
output_attentions: bool = False,
|
||
|
|
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
||
|
|
"""Input shape: Batch x Time x Channel"""
|
||
|
|
|
||
|
|
# if key_value_states are provided this layer is used as a cross-attention layer
|
||
|
|
# for the decoder
|
||
|
|
is_cross_attention = key_value_states is not None
|
||
|
|
bsz, tgt_len, embed_dim = hidden_states.size()
|
||
|
|
|
||
|
|
# get query proj
|
||
|
|
query_states = self.q_proj(hidden_states) * self.scaling
|
||
|
|
# get key, value proj
|
||
|
|
if is_cross_attention and past_key_value is not None:
|
||
|
|
# reuse k,v, cross_attentions
|
||
|
|
key_states = past_key_value[0]
|
||
|
|
value_states = past_key_value[1]
|
||
|
|
elif is_cross_attention:
|
||
|
|
# cross_attentions
|
||
|
|
key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
|
||
|
|
value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
|
||
|
|
elif past_key_value is not None:
|
||
|
|
# reuse k, v, self_attention
|
||
|
|
key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
|
||
|
|
value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
|
||
|
|
key_states = torch.cat([past_key_value[0], key_states], dim=2)
|
||
|
|
value_states = torch.cat([past_key_value[1], value_states], dim=2)
|
||
|
|
else:
|
||
|
|
# self_attention
|
||
|
|
key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
|
||
|
|
value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
|
||
|
|
|
||
|
|
if self.is_decoder:
|
||
|
|
# if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
|
||
|
|
# Further calls to cross_attention layer can then reuse all cross-attention
|
||
|
|
# key/value_states (first "if" case)
|
||
|
|
# if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
|
||
|
|
# all previous decoder key/value_states. Further calls to uni-directional self-attention
|
||
|
|
# can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
|
||
|
|
# if encoder bi-directional self-attention `past_key_value` is always `None`
|
||
|
|
past_key_value = (key_states, value_states)
|
||
|
|
|
||
|
|
proj_shape = (bsz * self.num_heads, -1, self.head_dim)
|
||
|
|
query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
|
||
|
|
key_states = key_states.view(*proj_shape)
|
||
|
|
value_states = value_states.view(*proj_shape)
|
||
|
|
|
||
|
|
src_len = key_states.size(1)
|
||
|
|
attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
|
||
|
|
|
||
|
|
if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
|
||
|
|
raise ValueError(
|
||
|
|
f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
|
||
|
|
f" {attn_weights.size()}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if attention_mask is not None:
|
||
|
|
if attention_mask.size() != (bsz, 1, tgt_len, src_len):
|
||
|
|
raise ValueError(
|
||
|
|
f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
|
||
|
|
)
|
||
|
|
attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
|
||
|
|
attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
|
||
|
|
|
||
|
|
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
|
||
|
|
|
||
|
|
if layer_head_mask is not None:
|
||
|
|
if layer_head_mask.size() != (self.num_heads,):
|
||
|
|
raise ValueError(
|
||
|
|
f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
|
||
|
|
f" {layer_head_mask.size()}"
|
||
|
|
)
|
||
|
|
attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
|
||
|
|
attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
|
||
|
|
|
||
|
|
if output_attentions:
|
||
|
|
# this operation is a bit awkward, but it's required to
|
||
|
|
# make sure that attn_weights keeps its gradient.
|
||
|
|
# In order to do so, attn_weights have to be reshaped
|
||
|
|
# twice and have to be reused in the following
|
||
|
|
attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
|
||
|
|
attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
|
||
|
|
else:
|
||
|
|
attn_weights_reshaped = None
|
||
|
|
|
||
|
|
attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
|
||
|
|
|
||
|
|
attn_output = torch.bmm(attn_probs, value_states)
|
||
|
|
|
||
|
|
if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
|
||
|
|
raise ValueError(
|
||
|
|
f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
|
||
|
|
f" {attn_output.size()}"
|
||
|
|
)
|
||
|
|
|
||
|
|
attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
|
||
|
|
attn_output = attn_output.transpose(1, 2)
|
||
|
|
attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)
|
||
|
|
|
||
|
|
attn_output = self.out_proj(attn_output)
|
||
|
|
|
||
|
|
return attn_output, attn_weights_reshaped, past_key_value
|
||
|
|
|
||
|
|
|
||
|
|
class TrOCRDecoderLayer(nn.Module):
|
||
|
|
def __init__(self, config: TrOCRConfig):
|
||
|
|
super().__init__()
|
||
|
|
self.embed_dim = config.hidden_size
|
||
|
|
|
||
|
|
self.self_attn = TrOCRAttention(
|
||
|
|
config,
|
||
|
|
embed_dim=self.embed_dim,
|
||
|
|
num_heads=config.decoder_attention_heads,
|
||
|
|
dropout=config.attention_dropout,
|
||
|
|
is_decoder=True,
|
||
|
|
)
|
||
|
|
self.dropout = config.dropout
|
||
|
|
self.activation_fn = ACT2FN[config.activation_function]
|
||
|
|
self.activation_dropout = config.activation_dropout
|
||
|
|
|
||
|
|
self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
|
||
|
|
|
||
|
|
if config.is_decoder:
|
||
|
|
self.encoder_attn = TrOCRAttention(
|
||
|
|
config,
|
||
|
|
embed_dim=self.embed_dim,
|
||
|
|
num_heads=config.decoder_attention_heads,
|
||
|
|
kdim=config.cross_attention_hidden_size,
|
||
|
|
vdim=config.cross_attention_hidden_size,
|
||
|
|
dropout=config.attention_dropout,
|
||
|
|
is_decoder=True,
|
||
|
|
is_cross_attention=True,
|
||
|
|
)
|
||
|
|
self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
|
||
|
|
|
||
|
|
self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
|
||
|
|
self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
|
||
|
|
self.final_layer_norm = nn.LayerNorm(self.embed_dim)
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
hidden_states: torch.Tensor,
|
||
|
|
attention_mask: Optional[torch.Tensor] = None,
|
||
|
|
encoder_hidden_states: Optional[torch.Tensor] = None,
|
||
|
|
encoder_attention_mask: Optional[torch.Tensor] = None,
|
||
|
|
layer_head_mask: Optional[torch.Tensor] = None,
|
||
|
|
cross_attn_layer_head_mask: Optional[torch.Tensor] = None,
|
||
|
|
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||
|
|
output_attentions: Optional[bool] = False,
|
||
|
|
use_cache: Optional[bool] = True,
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Args:
|
||
|
|
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
|
||
|
|
attention_mask (`torch.FloatTensor`): attention mask of size
|
||
|
|
`(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
|
||
|
|
encoder_hidden_states (`torch.FloatTensor`):
|
||
|
|
cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
|
||
|
|
encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
|
||
|
|
`(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
|
||
|
|
layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
|
||
|
|
`(encoder_attention_heads,)`.
|
||
|
|
cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of
|
||
|
|
size *(decoder_attention_heads,)*.
|
||
|
|
past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states
|
||
|
|
output_attentions (`bool`, *optional*):
|
||
|
|
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
|
||
|
|
returned tensors for more detail.
|
||
|
|
"""
|
||
|
|
residual = hidden_states
|
||
|
|
|
||
|
|
# Self Attention
|
||
|
|
# 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
|
||
|
|
# add present self-attn cache to positions 1,2 of present_key_value tuple
|
||
|
|
hidden_states, self_attn_weights, present_key_value = self.self_attn(
|
||
|
|
hidden_states=hidden_states,
|
||
|
|
past_key_value=self_attn_past_key_value,
|
||
|
|
attention_mask=attention_mask,
|
||
|
|
layer_head_mask=layer_head_mask,
|
||
|
|
output_attentions=output_attentions,
|
||
|
|
)
|
||
|
|
|
||
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
|
||
|
|
hidden_states = residual + hidden_states
|
||
|
|
hidden_states = self.self_attn_layer_norm(hidden_states)
|
||
|
|
|
||
|
|
# Cross-Attention Block
|
||
|
|
cross_attn_present_key_value = None
|
||
|
|
cross_attn_weights = None
|
||
|
|
|
||
|
|
if encoder_hidden_states is not None:
|
||
|
|
residual = hidden_states
|
||
|
|
|
||
|
|
# cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple
|
||
|
|
cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
|
||
|
|
hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
|
||
|
|
hidden_states=hidden_states,
|
||
|
|
key_value_states=encoder_hidden_states,
|
||
|
|
attention_mask=encoder_attention_mask,
|
||
|
|
layer_head_mask=cross_attn_layer_head_mask,
|
||
|
|
past_key_value=cross_attn_past_key_value,
|
||
|
|
output_attentions=output_attentions,
|
||
|
|
)
|
||
|
|
|
||
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
|
||
|
|
hidden_states = residual + hidden_states
|
||
|
|
hidden_states = self.encoder_attn_layer_norm(hidden_states)
|
||
|
|
|
||
|
|
# add cross-attn to positions 3,4 of present_key_value tuple
|
||
|
|
present_key_value = present_key_value + cross_attn_present_key_value
|
||
|
|
|
||
|
|
# Fully Connected
|
||
|
|
residual = hidden_states
|
||
|
|
hidden_states = self.activation_fn(self.fc1(hidden_states))
|
||
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
|
||
|
|
hidden_states = self.fc2(hidden_states)
|
||
|
|
|
||
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
|
||
|
|
hidden_states = residual + hidden_states
|
||
|
|
hidden_states = self.final_layer_norm(hidden_states)
|
||
|
|
|
||
|
|
outputs = (hidden_states,)
|
||
|
|
|
||
|
|
if output_attentions:
|
||
|
|
outputs += (self_attn_weights, cross_attn_weights)
|
||
|
|
|
||
|
|
if use_cache:
|
||
|
|
outputs += (present_key_value,)
|
||
|
|
|
||
|
|
return outputs
|
||
|
|
|
||
|
|
|
||
|
|
class TrOCRPreTrainedModel(PreTrainedModel):
|
||
|
|
config_class = TrOCRConfig
|
||
|
|
base_model_prefix = "model"
|
||
|
|
supports_gradient_checkpointing = True
|
||
|
|
|
||
|
|
def _init_weights(self, module):
|
||
|
|
std = self.config.init_std
|
||
|
|
if isinstance(module, (nn.Linear, nn.Conv1d)):
|
||
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
||
|
|
if module.bias is not None:
|
||
|
|
module.bias.data.zero_()
|
||
|
|
elif isinstance(module, nn.Embedding):
|
||
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
||
|
|
if module.padding_idx is not None:
|
||
|
|
module.weight.data[module.padding_idx].zero_()
|
||
|
|
|
||
|
|
|
||
|
|
TROCR_START_DOCSTRING = r"""
|
||
|
|
This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
|
||
|
|
library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
|
||
|
|
etc.)
|
||
|
|
|
||
|
|
This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
|
||
|
|
Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
|
||
|
|
and behavior.
|
||
|
|
|
||
|
|
Parameters:
|
||
|
|
config ([`TrOCRConfig`]):
|
||
|
|
Model configuration class with all the parameters of the model. Initializing with a config file does not
|
||
|
|
load the weights associated with the model, only the configuration. Check out the
|
||
|
|
[`~PreTrainedModel.from_pretrained`] method to load the model weights.
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
class TrOCRDecoder(TrOCRPreTrainedModel):
|
||
|
|
"""
|
||
|
|
Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`TrOCRDecoderLayer`]
|
||
|
|
|
||
|
|
Args:
|
||
|
|
config: TrOCRConfig
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, config: TrOCRConfig):
|
||
|
|
super().__init__(config)
|
||
|
|
self.dropout = config.dropout
|
||
|
|
self.layerdrop = config.decoder_layerdrop
|
||
|
|
self.padding_idx = config.pad_token_id
|
||
|
|
self.embed_scale = math.sqrt(config.hidden_size) if config.scale_embedding else 1.0
|
||
|
|
|
||
|
|
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
||
|
|
|
||
|
|
if config.use_learned_position_embeddings:
|
||
|
|
self.embed_positions = TrOCRLearnedPositionalEmbedding(config.max_position_embeddings, config.hidden_size)
|
||
|
|
else:
|
||
|
|
self.embed_positions = TrOCRSinusoidalPositionalEmbedding(
|
||
|
|
config.max_position_embeddings + self.padding_idx + 1,
|
||
|
|
config.hidden_size,
|
||
|
|
self.padding_idx,
|
||
|
|
)
|
||
|
|
|
||
|
|
if config.layernorm_embedding:
|
||
|
|
self.layernorm_embedding = nn.LayerNorm(config.hidden_size)
|
||
|
|
else:
|
||
|
|
self.layernorm_embedding = None
|
||
|
|
|
||
|
|
self.layers = nn.ModuleList([TrOCRDecoderLayer(config) for _ in range(config.decoder_layers)])
|
||
|
|
|
||
|
|
self.gradient_checkpointing = False
|
||
|
|
# Initialize weights and apply final processing
|
||
|
|
self.post_init()
|
||
|
|
|
||
|
|
def get_input_embeddings(self):
|
||
|
|
return self.embed_tokens
|
||
|
|
|
||
|
|
def set_input_embeddings(self, value):
|
||
|
|
self.embed_tokens = value
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
input_ids=None,
|
||
|
|
attention_mask=None,
|
||
|
|
encoder_hidden_states=None,
|
||
|
|
encoder_attention_mask=None,
|
||
|
|
head_mask=None,
|
||
|
|
cross_attn_head_mask=None,
|
||
|
|
past_key_values=None,
|
||
|
|
inputs_embeds=None,
|
||
|
|
use_cache=None,
|
||
|
|
output_attentions=None,
|
||
|
|
output_hidden_states=None,
|
||
|
|
return_dict=None,
|
||
|
|
):
|
||
|
|
r"""
|
||
|
|
Args:
|
||
|
|
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
|
||
|
|
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
|
||
|
|
provide it.
|
||
|
|
|
||
|
|
Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
|
||
|
|
[`PreTrainedTokenizer.__call__`] for details.
|
||
|
|
|
||
|
|
[What are input IDs?](../glossary#input-ids)
|
||
|
|
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||
|
|
Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
|
||
|
|
|
||
|
|
- 1 for tokens that are **not masked**,
|
||
|
|
- 0 for tokens that are **masked**.
|
||
|
|
|
||
|
|
[What are attention masks?](../glossary#attention-mask)
|
||
|
|
encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
|
||
|
|
Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
|
||
|
|
of the decoder.
|
||
|
|
encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
|
||
|
|
Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values
|
||
|
|
selected in `[0, 1]`:
|
||
|
|
|
||
|
|
- 1 for tokens that are **not masked**,
|
||
|
|
- 0 for tokens that are **masked**.
|
||
|
|
|
||
|
|
[What are attention masks?](../glossary#attention-mask)
|
||
|
|
head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
|
||
|
|
Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
|
||
|
|
|
||
|
|
- 1 indicates the head is **not masked**,
|
||
|
|
- 0 indicates the head is **masked**.
|
||
|
|
|
||
|
|
cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
|
||
|
|
Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention
|
||
|
|
on hidden heads. Mask values selected in `[0, 1]`:
|
||
|
|
|
||
|
|
- 1 indicates the head is **not masked**,
|
||
|
|
- 0 indicates the head is **masked**.
|
||
|
|
|
||
|
|
past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
|
||
|
|
Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
|
||
|
|
shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
|
||
|
|
shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
|
||
|
|
|
||
|
|
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
|
||
|
|
cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
|
||
|
|
|
||
|
|
If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
|
||
|
|
that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
|
||
|
|
all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
|
||
|
|
inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
|
||
|
|
Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
|
||
|
|
This is useful if you want more control over how to convert `input_ids` indices into associated vectors
|
||
|
|
than the model's internal embedding lookup matrix.
|
||
|
|
output_attentions (`bool`, *optional*):
|
||
|
|
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
|
||
|
|
returned tensors for more detail.
|
||
|
|
output_hidden_states (`bool`, *optional*):
|
||
|
|
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
|
||
|
|
for more detail.
|
||
|
|
return_dict (`bool`, *optional*):
|
||
|
|
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
||
|
|
"""
|
||
|
|
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
|
||
|
|
)
|
||
|
|
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||
|
|
|
||
|
|
# retrieve input_ids and inputs_embeds
|
||
|
|
if input_ids is not None and inputs_embeds is not None:
|
||
|
|
raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
|
||
|
|
elif input_ids is not None:
|
||
|
|
input = input_ids
|
||
|
|
input_ids = input_ids.view(-1, input.shape[-1])
|
||
|
|
elif inputs_embeds is not None:
|
||
|
|
input_shape = inputs_embeds.size()[:-1]
|
||
|
|
input = inputs_embeds[:, :, -1]
|
||
|
|
else:
|
||
|
|
raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
|
||
|
|
|
||
|
|
# past_key_values_length
|
||
|
|
past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
|
||
|
|
|
||
|
|
if inputs_embeds is None:
|
||
|
|
inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
|
||
|
|
|
||
|
|
if self.config.use_learned_position_embeddings:
|
||
|
|
embed_pos = self.embed_positions(input, past_key_values_length=past_key_values_length)
|
||
|
|
else:
|
||
|
|
embed_pos = self.embed_positions(input_ids, past_key_values_length=past_key_values_length)
|
||
|
|
|
||
|
|
hidden_states = inputs_embeds + embed_pos
|
||
|
|
|
||
|
|
if self.layernorm_embedding is not None:
|
||
|
|
hidden_states = self.layernorm_embedding(hidden_states)
|
||
|
|
|
||
|
|
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
|
||
|
|
|
||
|
|
input_shape = input.shape
|
||
|
|
|
||
|
|
attention_mask = _prepare_4d_causal_attention_mask(
|
||
|
|
attention_mask, input_shape, inputs_embeds, past_key_values_length
|
||
|
|
)
|
||
|
|
|
||
|
|
# expand encoder attention mask
|
||
|
|
if encoder_hidden_states is not None and encoder_attention_mask is not None:
|
||
|
|
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||
|
|
encoder_attention_mask = _prepare_4d_attention_mask(
|
||
|
|
encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
|
||
|
|
)
|
||
|
|
|
||
|
|
if self.gradient_checkpointing and self.training:
|
||
|
|
if use_cache:
|
||
|
|
logger.warning_once(
|
||
|
|
"`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`..."
|
||
|
|
)
|
||
|
|
use_cache = False
|
||
|
|
|
||
|
|
# decoder layers
|
||
|
|
all_hidden_states = () if output_hidden_states else None
|
||
|
|
all_self_attns = () if output_attentions else None
|
||
|
|
all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
|
||
|
|
next_decoder_cache = () if use_cache else None
|
||
|
|
|
||
|
|
# check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
|
||
|
|
for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
|
||
|
|
if attn_mask is not None:
|
||
|
|
if attn_mask.size()[0] != (len(self.layers)):
|
||
|
|
raise ValueError(
|
||
|
|
f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
|
||
|
|
f" {head_mask.size()[0]}."
|
||
|
|
)
|
||
|
|
for idx, decoder_layer in enumerate(self.layers):
|
||
|
|
# add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
|
||
|
|
if output_hidden_states:
|
||
|
|
all_hidden_states += (hidden_states,)
|
||
|
|
if self.training:
|
||
|
|
dropout_probability = torch.rand([])
|
||
|
|
if dropout_probability < self.layerdrop:
|
||
|
|
continue
|
||
|
|
|
||
|
|
past_key_value = past_key_values[idx] if past_key_values is not None else None
|
||
|
|
|
||
|
|
if self.gradient_checkpointing and self.training:
|
||
|
|
layer_outputs = self._gradient_checkpointing_func(
|
||
|
|
decoder_layer.__call__,
|
||
|
|
hidden_states,
|
||
|
|
attention_mask,
|
||
|
|
encoder_hidden_states,
|
||
|
|
encoder_attention_mask,
|
||
|
|
head_mask[idx] if head_mask is not None else None,
|
||
|
|
cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
|
||
|
|
None,
|
||
|
|
output_attentions,
|
||
|
|
use_cache,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
layer_outputs = decoder_layer(
|
||
|
|
hidden_states,
|
||
|
|
attention_mask=attention_mask,
|
||
|
|
encoder_hidden_states=encoder_hidden_states,
|
||
|
|
encoder_attention_mask=encoder_attention_mask,
|
||
|
|
layer_head_mask=(head_mask[idx] if head_mask is not None else None),
|
||
|
|
cross_attn_layer_head_mask=(
|
||
|
|
cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
|
||
|
|
),
|
||
|
|
past_key_value=past_key_value,
|
||
|
|
output_attentions=output_attentions,
|
||
|
|
use_cache=use_cache,
|
||
|
|
)
|
||
|
|
hidden_states = layer_outputs[0]
|
||
|
|
|
||
|
|
if use_cache:
|
||
|
|
next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)
|
||
|
|
|
||
|
|
if output_attentions:
|
||
|
|
all_self_attns += (layer_outputs[1],)
|
||
|
|
|
||
|
|
if encoder_hidden_states is not None:
|
||
|
|
all_cross_attentions += (layer_outputs[2],)
|
||
|
|
|
||
|
|
# add hidden states from the last decoder layer
|
||
|
|
if output_hidden_states:
|
||
|
|
all_hidden_states += (hidden_states,)
|
||
|
|
|
||
|
|
next_cache = next_decoder_cache if use_cache else None
|
||
|
|
if not return_dict:
|
||
|
|
return tuple(
|
||
|
|
v
|
||
|
|
for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions]
|
||
|
|
if v is not None
|
||
|
|
)
|
||
|
|
return BaseModelOutputWithPastAndCrossAttentions(
|
||
|
|
last_hidden_state=hidden_states,
|
||
|
|
past_key_values=next_cache,
|
||
|
|
hidden_states=all_hidden_states,
|
||
|
|
attentions=all_self_attns,
|
||
|
|
cross_attentions=all_cross_attentions,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@add_start_docstrings(
|
||
|
|
"The TrOCR Model with a language modeling head. Can be used for summarization.",
|
||
|
|
TROCR_START_DOCSTRING,
|
||
|
|
)
|
||
|
|
class TrOCRDecoderWrapper(TrOCRPreTrainedModel):
|
||
|
|
"""
|
||
|
|
This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is
|
||
|
|
used in combination with the [`EncoderDecoderModel`] framework.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, config):
|
||
|
|
super().__init__(config)
|
||
|
|
self.decoder = TrOCRDecoder(config)
|
||
|
|
|
||
|
|
def forward(self, *args, **kwargs):
|
||
|
|
return self.decoder(*args, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
@add_start_docstrings(
|
||
|
|
"The TrOCR Decoder with a language modeling head. Can be used as the decoder part of [`EncoderDecoderModel`] and"
|
||
|
|
" [`VisionEncoderDecoder`].",
|
||
|
|
TROCR_START_DOCSTRING,
|
||
|
|
)
|
||
|
|
class TrOCRForCausalLM(TrOCRPreTrainedModel):
|
||
|
|
_tied_weights_keys = ["output_projection.weight"]
|
||
|
|
|
||
|
|
def __init__(self, config):
|
||
|
|
config = copy.deepcopy(config)
|
||
|
|
config.is_decoder = True
|
||
|
|
config.is_encoder_decoder = False
|
||
|
|
super().__init__(config)
|
||
|
|
self.model = TrOCRDecoderWrapper(config)
|
||
|
|
|
||
|
|
self.output_projection = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||
|
|
|
||
|
|
# Initialize weights and apply final processing
|
||
|
|
self.post_init()
|
||
|
|
|
||
|
|
def get_input_embeddings(self):
|
||
|
|
return self.model.decoder.embed_tokens
|
||
|
|
|
||
|
|
def set_input_embeddings(self, value):
|
||
|
|
self.model.decoder.embed_tokens = value
|
||
|
|
|
||
|
|
def get_output_embeddings(self):
|
||
|
|
return self.output_projection
|
||
|
|
|
||
|
|
def set_output_embeddings(self, new_embeddings):
|
||
|
|
self.output_projection = new_embeddings
|
||
|
|
|
||
|
|
def set_decoder(self, decoder):
|
||
|
|
self.model.decoder = decoder
|
||
|
|
|
||
|
|
def get_decoder(self):
|
||
|
|
return self.model.decoder
|
||
|
|
|
||
|
|
@replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
input_ids: Optional[torch.LongTensor] = None,
|
||
|
|
attention_mask: Optional[torch.Tensor] = None,
|
||
|
|
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||
|
|
encoder_attention_mask: Optional[torch.LongTensor] = None,
|
||
|
|
head_mask: Optional[torch.Tensor] = None,
|
||
|
|
cross_attn_head_mask: Optional[torch.Tensor] = None,
|
||
|
|
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
|
||
|
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||
|
|
labels: Optional[torch.LongTensor] = None,
|
||
|
|
use_cache: Optional[bool] = None,
|
||
|
|
output_attentions: Optional[bool] = None,
|
||
|
|
output_hidden_states: Optional[bool] = None,
|
||
|
|
return_dict: Optional[bool] = None,
|
||
|
|
) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
|
||
|
|
r"""
|
||
|
|
Args:
|
||
|
|
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
|
||
|
|
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
|
||
|
|
provide it.
|
||
|
|
|
||
|
|
Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
|
||
|
|
[`PreTrainedTokenizer.__call__`] for details.
|
||
|
|
|
||
|
|
[What are input IDs?](../glossary#input-ids)
|
||
|
|
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||
|
|
Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
|
||
|
|
|
||
|
|
- 1 for tokens that are **not masked**,
|
||
|
|
- 0 for tokens that are **masked**.
|
||
|
|
|
||
|
|
[What are attention masks?](../glossary#attention-mask)
|
||
|
|
encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
|
||
|
|
Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
|
||
|
|
if the model is configured as a decoder.
|
||
|
|
encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||
|
|
Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used
|
||
|
|
in the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
|
||
|
|
head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
|
||
|
|
Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
|
||
|
|
|
||
|
|
- 1 indicates the head is **not masked**,
|
||
|
|
- 0 indicates the head is **masked**.
|
||
|
|
|
||
|
|
cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
|
||
|
|
Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`:
|
||
|
|
|
||
|
|
- 1 indicates the head is **not masked**,
|
||
|
|
- 0 indicates the head is **masked**.
|
||
|
|
|
||
|
|
past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
|
||
|
|
Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
|
||
|
|
shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
|
||
|
|
shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. The two additional
|
||
|
|
tensors are only required when the model is used as a decoder in a Sequence to Sequence model.
|
||
|
|
|
||
|
|
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
|
||
|
|
cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
|
||
|
|
|
||
|
|
If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
|
||
|
|
that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
|
||
|
|
all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
|
||
|
|
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||
|
|
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
|
||
|
|
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
|
||
|
|
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
|
||
|
|
use_cache (`bool`, *optional*):
|
||
|
|
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
|
||
|
|
(see `past_key_values`).
|
||
|
|
|
||
|
|
- 1 for tokens that are **not masked**,
|
||
|
|
- 0 for tokens that are **masked**.
|
||
|
|
output_attentions (`bool`, *optional*):
|
||
|
|
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
|
||
|
|
returned tensors for more detail.
|
||
|
|
output_hidden_states (`bool`, *optional*):
|
||
|
|
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
|
||
|
|
for more detail.
|
||
|
|
return_dict (`bool`, *optional*):
|
||
|
|
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
|
||
|
|
Example:
|
||
|
|
|
||
|
|
```python
|
||
|
|
>>> from transformers import (
|
||
|
|
... TrOCRConfig,
|
||
|
|
... TrOCRProcessor,
|
||
|
|
... TrOCRForCausalLM,
|
||
|
|
... ViTConfig,
|
||
|
|
... ViTModel,
|
||
|
|
... VisionEncoderDecoderModel,
|
||
|
|
... )
|
||
|
|
>>> import requests
|
||
|
|
>>> from PIL import Image
|
||
|
|
|
||
|
|
>>> # TrOCR is a decoder model and should be used within a VisionEncoderDecoderModel
|
||
|
|
>>> # init vision2text model with random weights
|
||
|
|
>>> encoder = ViTModel(ViTConfig())
|
||
|
|
>>> decoder = TrOCRForCausalLM(TrOCRConfig())
|
||
|
|
>>> model = VisionEncoderDecoderModel(encoder=encoder, decoder=decoder)
|
||
|
|
|
||
|
|
>>> # If you want to start from the pretrained model, load the checkpoint with `VisionEncoderDecoderModel`
|
||
|
|
>>> processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
|
||
|
|
>>> model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
|
||
|
|
|
||
|
|
>>> # load image from the IAM dataset
|
||
|
|
>>> url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"
|
||
|
|
>>> image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
|
||
|
|
>>> pixel_values = processor(image, return_tensors="pt").pixel_values
|
||
|
|
>>> text = "industry, ' Mr. Brown commented icily. ' Let us have a"
|
||
|
|
|
||
|
|
>>> # training
|
||
|
|
>>> model.config.decoder_start_token_id = processor.tokenizer.eos_token_id
|
||
|
|
>>> model.config.pad_token_id = processor.tokenizer.pad_token_id
|
||
|
|
>>> model.config.vocab_size = model.config.decoder.vocab_size
|
||
|
|
|
||
|
|
>>> labels = processor.tokenizer(text, return_tensors="pt").input_ids
|
||
|
|
>>> outputs = model(pixel_values, labels=labels)
|
||
|
|
>>> loss = outputs.loss
|
||
|
|
>>> round(loss.item(), 2)
|
||
|
|
5.30
|
||
|
|
|
||
|
|
>>> # inference
|
||
|
|
>>> generated_ids = model.generate(pixel_values)
|
||
|
|
>>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
||
|
|
>>> generated_text
|
||
|
|
'industry, " Mr. Brown commented icily. " Let us have a'
|
||
|
|
```"""
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||
|
|
outputs = self.model.decoder(
|
||
|
|
input_ids=input_ids,
|
||
|
|
attention_mask=attention_mask,
|
||
|
|
encoder_hidden_states=encoder_hidden_states,
|
||
|
|
encoder_attention_mask=encoder_attention_mask,
|
||
|
|
head_mask=head_mask,
|
||
|
|
cross_attn_head_mask=cross_attn_head_mask,
|
||
|
|
past_key_values=past_key_values,
|
||
|
|
inputs_embeds=inputs_embeds,
|
||
|
|
use_cache=use_cache,
|
||
|
|
output_attentions=output_attentions,
|
||
|
|
output_hidden_states=output_hidden_states,
|
||
|
|
return_dict=return_dict,
|
||
|
|
)
|
||
|
|
|
||
|
|
logits = self.output_projection(outputs[0])
|
||
|
|
|
||
|
|
loss = None
|
||
|
|
if labels is not None:
|
||
|
|
loss_fct = CrossEntropyLoss()
|
||
|
|
loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
|
||
|
|
|
||
|
|
if not return_dict:
|
||
|
|
output = (logits,) + outputs[1:]
|
||
|
|
return (loss,) + output if loss is not None else output
|
||
|
|
|
||
|
|
return CausalLMOutputWithCrossAttentions(
|
||
|
|
loss=loss,
|
||
|
|
logits=logits,
|
||
|
|
past_key_values=outputs.past_key_values,
|
||
|
|
hidden_states=outputs.hidden_states,
|
||
|
|
attentions=outputs.attentions,
|
||
|
|
cross_attentions=outputs.cross_attentions,
|
||
|
|
)
|
||
|
|
|
||
|
|
def prepare_inputs_for_generation(
|
||
|
|
self, input_ids, past_key_values=None, attention_mask=None, use_cache=None, **kwargs
|
||
|
|
):
|
||
|
|
# if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
|
||
|
|
if attention_mask is None:
|
||
|
|
attention_mask = input_ids.new_ones(input_ids.shape)
|
||
|
|
|
||
|
|
if past_key_values:
|
||
|
|
past_length = past_key_values[0][0].shape[2]
|
||
|
|
|
||
|
|
# Some generation methods already pass only the last input ID
|
||
|
|
if input_ids.shape[1] > past_length:
|
||
|
|
remove_prefix_length = past_length
|
||
|
|
else:
|
||
|
|
# Default to old behavior: keep only final ID
|
||
|
|
remove_prefix_length = input_ids.shape[1] - 1
|
||
|
|
|
||
|
|
input_ids = input_ids[:, remove_prefix_length:]
|
||
|
|
# first step, decoder_cached_states are empty
|
||
|
|
return {
|
||
|
|
"input_ids": input_ids, # encoder_outputs is defined. input_ids not needed
|
||
|
|
"attention_mask": attention_mask,
|
||
|
|
"past_key_values": past_key_values,
|
||
|
|
"use_cache": use_cache,
|
||
|
|
}
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _reorder_cache(past_key_values, beam_idx):
|
||
|
|
reordered_past = ()
|
||
|
|
for layer_past in past_key_values:
|
||
|
|
reordered_past += (
|
||
|
|
tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
|
||
|
|
)
|
||
|
|
return reordered_past
|