openocr-python 0.0.9__py3-none-any.whl → 0.1.0.dev0__py3-none-any.whl

openocr/openrec/modeling/decoders/smtr_decoder.py

@@ -71,6 +71,7 @@ class SSMatchLayer(nn.Module):
         drop_path=0.0,
         act_layer=nn.GELU,
         epsilon=1e-6,
+        is_last_layer=False,
     ):
         super().__init__()
         self.dim = dim
@@ -96,17 +97,18 @@ class SSMatchLayer(nn.Module):
             proj_drop=drop)
         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
         self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
+        self.is_last_layer = is_last_layer
 
     def forward(self, question_f, prompt_f, visual_f, mask=None):
-
         question_f = question_f + self.drop_path(
             self.images_to_question_cross_attn(self.normq1(question_f),
                                                self.normkv1(prompt_f), mask))
         question_f = question_f.reshape(visual_f.shape[0], -1, self.dim)
         question_f = self.question_to_images_cross_attn(
             self.normq2(question_f), self.normkv2(visual_f))
-
-        return question_f
+        if self.is_last_layer:
+            return question_f
+        return question_f.flatten(0, 1).unsqueeze(1)
 
 
 class SMTRDecoder(nn.Module):
@@ -152,7 +154,9 @@ class SMTRDecoder(nn.Module):
                          dynq2img_heads=dynq2img_heads,
                          mlp_ratio=4.0,
                          qkv_bias=True,
-                         drop_path=dpr[i]) for i in range(num_layer)
+                         drop_path=dpr[i],
+                         is_last_layer=i == num_layer - 1)
+            for i in range(num_layer)
         ])
 
         self.ds = ds
@@ -351,7 +355,6 @@ class SMTRDecoder(nn.Module):
             else:
                 return torch.concat(next_logits_all + pre_logits_all[::-1], 1)
 
-
     def forward_test_bi_attn(self, x):
         self.attn_maps = []
         if not self.ds:
@@ -366,81 +369,122 @@ class SMTRDecoder(nn.Module):
             next = self.next_token
             pre = self.pre_token
             next_pre = torch.concat([next, pre], 0)
-            next_pre = next_pre.squeeze(1) #2, 1, dim
+            next_pre = next_pre.squeeze(1)  #2, 1, dim
 
             prompt_next_embed = self.prompt_next_embed.squeeze(1)
             prompt_pre_embed = self.prompt_pre_embed.squeeze(1)
 
-            next_id = torch.full([1, self.sub_str_len], self.bos_next, dtype=torch.long, device=x.device)
-            pre_id = torch.full([1, self.sub_str_len], self.bos_pre, dtype=torch.long, device=x.device)
+            next_id = torch.full([1, self.sub_str_len],
+                                 self.bos_next,
+                                 dtype=torch.long,
+                                 device=x.device)
+            pre_id = torch.full([1, self.sub_str_len],
+                                self.bos_pre,
+                                dtype=torch.long,
+                                device=x.device)
             # prompt_next_bos = self.char_embed(prompt_id)
             # pred_prob_list = torch.full([bs, self.sub_str_len], self.ignore_index, dtype=torch.long, device=x.device)
-            next_pred_id_list = torch.full([1, self.max_len], self.ignore_index, dtype=torch.long, device=x.device)
-            pre_pred_id_list = torch.full([1, self.max_len], self.ignore_index, dtype=torch.long, device=x.device)
+            next_pred_id_list = torch.full([1, self.max_len],
+                                           self.ignore_index,
+                                           dtype=torch.long,
+                                           device=x.device)
+            pre_pred_id_list = torch.full([1, self.max_len],
+                                          self.ignore_index,
+                                          dtype=torch.long,
+                                          device=x.device)
             next_logits_all = []
             pre_logits_all = []
             attn_map_next = []
             attn_map_pre = []
-            mask_pad = torch.zeros([bs, 1], dtype=torch.float32, device=x.device)
-            for j in range(0, min(70, self.max_len-1)):
-                
-                prompt_char_next = torch.concat([prompt_next_embed[:, :1, :], prompt_next_embed[:, 1:, :] + self.char_embed(next_id)], 1) # b, sub_l, dim
-                prompt_char_pre = torch.concat([prompt_pre_embed[:, :1, :], prompt_pre_embed[:, 1:, :] + self.char_embed(pre_id)], 1) # b, sub_l, dim
-                prompt_char = torch.concat([prompt_char_next, prompt_char_pre], 0) #2, 6, dim
+            mask_pad = torch.zeros([bs, 1],
+                                   dtype=torch.float32,
+                                   device=x.device)
+            for j in range(0, min(70, self.max_len - 1)):
+
+                prompt_char_next = torch.concat([
+                    prompt_next_embed[:, :1, :],
+                    prompt_next_embed[:, 1:, :] + self.char_embed(next_id)
+                ], 1)  # b, sub_l, dim
+                prompt_char_pre = torch.concat([
+                    prompt_pre_embed[:, :1, :],
+                    prompt_pre_embed[:, 1:, :] + self.char_embed(pre_id)
+                ], 1)  # b, sub_l, dim
+                prompt_char = torch.concat([prompt_char_next, prompt_char_pre],
+                                           0)  #2, 6, dim
                 # prompt_char = prompt_char.flatten(0, 1)
 
-                mask_next = torch.where(next_id == self.bos_next, float('-inf'), 0) # b, subs_l
-                mask_pre = torch.where(pre_id == self.bos_pre, float('-inf'), 0) # b, subs_l
-                mask = torch.concat([mask_next, mask_pre], 0) #2, 5
-                mask = torch.concat([mask_pad, mask], 1) # 2, 6
+                mask_next = torch.where(next_id == self.bos_next,
+                                        float('-inf'), 0)  # b, subs_l
+                mask_pre = torch.where(pre_id == self.bos_pre, float('-inf'),
+                                       0)  # b, subs_l
+                mask = torch.concat([mask_next, mask_pre], 0)  #2, 5
+                mask = torch.concat([mask_pad, mask], 1)  # 2, 6
                 pred_token = next_pre
-                visual_f_i = visual_f[:2] # 2 l dim
+                visual_f_i = visual_f[:2]  # 2 l dim
                 for layer in self.cmff_decoder:
-                    pred_token = layer(pred_token, prompt_char, visual_f_i, mask.unsqueeze(1))
-                
-                
+                    pred_token = layer(pred_token, prompt_char, visual_f_i,
+                                       mask.unsqueeze(1))
+
                 logits_next_i = self.ques1_head(self.norm_pred(pred_token))
                 logits = F.softmax(logits_next_i, -1)
-                pred_id_i = logits.argmax(-1) #2, 1
+                pred_id_i = logits.argmax(-1)  #2, 1
                 # print(pred_id_i.shape)
-                
-                next_pred_id_list[:, j:j+1] = pred_id_i[:1]
-                pre_pred_id_list[:, j:j+1] = pred_id_i[1:2]
+
+                next_pred_id_list[:, j:j + 1] = pred_id_i[:1]
+                pre_pred_id_list[:, j:j + 1] = pred_id_i[1:2]
                 if not (next_pred_id_list == self.eos).any(dim=-1).all():
                     next_logits_all.append(logits[:1])
-                    attn_map_next.append(self.cmff_decoder[-1].question_to_images_cross_attn.attn_map[0])
+                    attn_map_next.append(
+                        self.cmff_decoder[-1].question_to_images_cross_attn.
+                        attn_map[0])
                     next_id = torch.concat([next_id[:, 1:], pred_id_i[:1]], 1)
                 if not (pre_pred_id_list == self.eos).any(dim=-1).all():
                     pre_logits_all.append(logits[1:2])
-                    attn_map_pre.append(self.cmff_decoder[-1].question_to_images_cross_attn.attn_map[1])
+                    attn_map_pre.append(
+                        self.cmff_decoder[-1].question_to_images_cross_attn.
+                        attn_map[1])
                     pre_id = torch.concat([pred_id_i[1:2], pre_id[:, :-1]], 1)
-                
-                if (next_pred_id_list == self.eos).any(dim=-1).all() and (pre_pred_id_list == self.eos).any(dim=-1).all():
+
+                if (next_pred_id_list == self.eos).any(dim=-1).all() and (
+                        pre_pred_id_list == self.eos).any(dim=-1).all():
                     break
                 # print(next_id, pre_id)
             # exit(0)
-            if len(next_logits_all) > self.sub_str_len and len(pre_logits_all) > self.sub_str_len:
-                next_logits_all_ = torch.concat(next_logits_all[:-1], 1) # 1, l
-                pre_logits_all_ = torch.concat(pre_logits_all[:-1][::-1], 1) #1, l
+            if len(next_logits_all) > self.sub_str_len and len(
+                    pre_logits_all) > self.sub_str_len:
+                next_logits_all_ = torch.concat(next_logits_all[:-1],
+                                                1)  # 1, l
+                pre_logits_all_ = torch.concat(pre_logits_all[:-1][::-1],
+                                               1)  #1, l
 
                 next_id = next_logits_all_.argmax(-1)[:, -self.sub_str_len:]
                 pre_id = pre_logits_all_.argmax(-1)[:, :self.sub_str_len]
                 next_logits_all_mid = []
                 attn_map_next_mid = []
                 ques_next = self.next_token.tile([1, 1, 1, 1]).squeeze(1)
-                mask_pad = torch.zeros([1, 1], dtype=torch.float32, device=x.device)
-                for j in range(0, min(70, self.max_len-1)):
-                    
-                    prompt_next = torch.concat([prompt_next_embed[:, :1, :], prompt_next_embed[:, 1:, :] + self.char_embed(next_id)], 1) # b, sub_l, dim
-                    mask_next = torch.where(next_id == self.bos_next, float('-inf'), 0) # b, subs_l
+                mask_pad = torch.zeros([1, 1],
+                                       dtype=torch.float32,
+                                       device=x.device)
+                for j in range(0, min(70, self.max_len - 1)):
+
+                    prompt_next = torch.concat([
+                        prompt_next_embed[:, :1, :],
+                        prompt_next_embed[:, 1:, :] + self.char_embed(next_id)
+                    ], 1)  # b, sub_l, dim
+                    mask_next = torch.where(next_id == self.bos_next,
+                                            float('-inf'), 0)  # b, subs_l
                     mask = torch.concat([mask_pad, mask_next], 1)
                     # prompt_next = self.char_embed(prompt_id)
                     ques_next_i = ques_next
                     visual_f_i = visual_f[2:3]
                     for layer in self.cmff_decoder:
-                        ques_next_i = layer(ques_next_i, prompt_next, visual_f_i, mask.unsqueeze(1))
-                    logits_next_i = self.ques1_head(self.norm_pred(ques_next_i))
-                    attn_map_next_mid.append(self.cmff_decoder[-1].question_to_images_cross_attn.attn_map[0])
+                        ques_next_i = layer(ques_next_i, prompt_next,
+                                            visual_f_i, mask.unsqueeze(1))
+                    logits_next_i = self.ques1_head(
+                        self.norm_pred(ques_next_i))
+                    attn_map_next_mid.append(
+                        self.cmff_decoder[-1].question_to_images_cross_attn.
+                        attn_map[0])
                     logits = F.softmax(logits_next_i, -1)
                     pred_id_i = logits.argmax(-1)
                     next_logits_all_mid.append(logits)
@@ -449,12 +493,19 @@ class SMTRDecoder(nn.Module):
                         break
                 next_logits_all_mid = torch.concat(next_logits_all_mid, 1)
                 # next_logits_all_ = torch.concat([next_logits_all_, next_logits_all], 1)
-                self.attn_maps = [attn_map_next, attn_map_next_mid, attn_map_pre[::-1]]
-                return [torch.concat(next_logits_all, 1), next_logits_all_mid, torch.concat(pre_logits_all[::-1], 1)]
+                self.attn_maps = [
+                    attn_map_next, attn_map_next_mid, attn_map_pre[::-1]
+                ]
+                return [
+                    torch.concat(next_logits_all, 1), next_logits_all_mid,
+                    torch.concat(pre_logits_all[::-1], 1)
+                ]
             else:
                 self.attn_maps = [attn_map_next, attn_map_pre[::-1]]
-                return [torch.concat(next_logits_all, 1), torch.concat(pre_logits_all[::-1], 1)]
-
+                return [
+                    torch.concat(next_logits_all, 1),
+                    torch.concat(pre_logits_all[::-1], 1)
+                ]
 
     def forward_test(self, x):
         self.attn_maps = []
@@ -579,7 +630,7 @@ class SMTRDecoder(nn.Module):
         prompt_char_next = torch.concat([
             prompt_next_embed[:, :, :1, :],
             prompt_next_embed[:, :, 1:, :] + self.char_embed(subs)
-        ], 2)  # b, n, sub_l, dim
+        ], 2)  # b, n, subs_l, dim
         next = self.next_token.tile([bs, max_len_curr, 1, 1])
 
         max_len_curr_pre = targets[6].max()

openocr/openrec/modeling/unirec_modeling/configuration_unirec.py

@@ -0,0 +1,166 @@
+# 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.
+"""UniRecConfig model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class UniRecConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`M2M100Model`]. It is used to instantiate an
+    M2M100 model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the M2M100
+    [facebook/m2m100_418M](https://huggingface.co/facebook/m2m100_418M) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 50265):
+            Vocabulary size of the M2M100 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`M2M100Model`] or
+        d_model (`int`, *optional*, defaults to 1024):
+            Dimensionality of the layers and the pooler layer.
+        encoder_layers (`int`, *optional*, defaults to 12):
+            Number of encoder layers.
+        decoder_layers (`int`, *optional*, defaults to 12):
+            Number of decoder layers.
+        encoder_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        decoder_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        decoder_ffn_dim (`int`, *optional*, defaults to 4096):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
+        encoder_ffn_dim (`int`, *optional*, defaults to 4096):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
+        activation_function (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"silu"` and `"gelu_new"` are supported.
+        dropout (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        activation_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for activations inside the fully connected layer.
+        classifier_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for classifier.
+        max_position_embeddings (`int`, *optional*, defaults to 1024):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        init_std (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        encoder_layerdrop (`float`, *optional*, defaults to 0.0):
+            The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
+            for more details.
+        decoder_layerdrop (`float`, *optional*, defaults to 0.0):
+            The LayerDrop probability for the decoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
+            for more details.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models).
+
+    Example:
+
+    ```python
+    >>> from transformers import M2M100Config, M2M100Model
+
+    >>> # Initializing a M2M100 facebook/m2m100_418M style configuration
+    >>> configuration = M2M100Config()
+
+    >>> # Initializing a model (with random weights) from the facebook/m2m100_418M style configuration
+    >>> model = M2M100Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = 'm2m_100'
+    keys_to_ignore_at_inference = ['past_key_values']
+    attribute_map = {
+        'num_attention_heads': 'encoder_attention_heads',
+        'hidden_size': 'd_model'
+    }
+
+    def __init__(
+        self,
+        vocab_size=50000,
+        max_position_embeddings=3072,
+        decoder_layers=6,
+        decoder_ffn_dim=1536,
+        decoder_attention_heads=6,
+        encoder_layerdrop=0.0,
+        decoder_layerdrop=0.0,
+        use_cache=True,
+        is_encoder_decoder=True,
+        activation_function='relu',
+        d_model=384,
+        dropout=0.1,
+        attention_dropout=0.1,
+        activation_dropout=0.0,
+        init_std=0.02,
+        decoder_start_token_id=0,
+        scale_embedding=True,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        depths=[2, 2, 9, 2],
+        dims=[64, 128, 256, 384],
+        mixer=[['Conv'] * 2, ['Conv'] * 2,
+               ['Conv'] * 6 + ['FGlobal', 'Global', 'Global'], ['Global'] * 2],
+        num_heads=[2, 4, 4, 6],
+        sub_k=[[2, 2], [2, 2], [2, 2], [2, 2]],
+        mlp_ratio=4,
+        kernel_size=[3, 3],
+        drop_path_rate=0.1,
+        label_smoothing=0.1,
+        torch_dtype='bfloat16',
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.d_model = d_model
+        self.decoder_ffn_dim = decoder_ffn_dim
+        self.decoder_layers = decoder_layers
+        self.decoder_attention_heads = decoder_attention_heads
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.activation_function = activation_function
+        self.init_std = init_std
+        self.encoder_layerdrop = encoder_layerdrop
+        self.decoder_layerdrop = decoder_layerdrop
+        self.use_cache = use_cache
+        self.scale_embedding = scale_embedding  # scale factor will be sqrt(d_model) if True
+        self.depths = depths
+        self.dims = dims
+        self.mixer = mixer
+        self.num_heads = num_heads
+        self.sub_k = sub_k
+        self.mlp_ratio = mlp_ratio
+        self.kernel_size = kernel_size
+        self.drop_path_rate = drop_path_rate
+        self.label_smoothing = label_smoothing
+        self.torch_dtype = torch_dtype
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            is_encoder_decoder=is_encoder_decoder,
+            decoder_start_token_id=decoder_start_token_id,
+            **kwargs,
+        )
+
+
+__all__ = ['UniRecConfig']