We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Models] Add SAR teacher forçing during training (decoder)
Teacher forçing (feeding the LSTM decoder of SAR with ground-truth characters during training) is not implemented, and it should really improve performances

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `doctr/models/recognition/sar.py`
Content:
```
1 # Copyright (C) 2021, Mindee.
2 
3 # This program is licensed under the Apache License version 2.
4 # See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
5 
6 from copy import deepcopy
7 import tensorflow as tf
8 from tensorflow.keras import Sequential, layers
9 from typing import Tuple, Dict, List, Any, Optional
10 
11 from .. import vgg
12 from ..utils import load_pretrained_params
13 from .core import RecognitionModel
14 from .core import RecognitionPostProcessor
15 from doctr.utils.repr import NestedObject
16 
17 __all__ = ['SAR', 'SARPostProcessor', 'sar_vgg16_bn']
18 
19 default_cfgs: Dict[str, Dict[str, Any]] = {
20     'sar_vgg16_bn': {
21         'backbone': 'vgg16_bn', 'rnn_units': 512, 'max_length': 40, 'num_decoders': 2,
22         'input_shape': (64, 256, 3),
23         'post_processor': 'SARPostProcessor',
24         'vocab': ('3K}7eé;5àÎYho]QwV6qU~W"XnbBvcADfËmy.9ÔpÛ*{CôïE%M4#ÈR:g@T$x?0î£|za1ù8,OG€P-'
25                   'kçHëÀÂ2É/ûIJ\'j(LNÙFut[)èZs+&°Sd=Ï!<â_Ç>rêi`l'),
26         'url': None,
27     },
28 }
29 
30 
31 class AttentionModule(layers.Layer, NestedObject):
32     """Implements attention module of the SAR model
33 
34     Args:
35         attention_units: number of hidden attention units
36 
37     """
38     def __init__(
39         self,
40         attention_units: int
41     ) -> None:
42 
43         super().__init__()
44         self.hidden_state_projector = layers.Conv2D(
45             attention_units, 1, strides=1, use_bias=False, padding='same', kernel_initializer='he_normal',
46         )
47         self.features_projector = layers.Conv2D(
48             attention_units, 3, strides=1, use_bias=True, padding='same', kernel_initializer='he_normal',
49         )
50         self.attention_projector = layers.Conv2D(
51             1, 1, strides=1, use_bias=False, padding="same", kernel_initializer='he_normal',
52         )
53         self.flatten = layers.Flatten()
54 
55     def call(
56         self,
57         features: tf.Tensor,
58         hidden_state: tf.Tensor,
59         **kwargs: Any,
60     ) -> tf.Tensor:
61 
62         [H, W] = features.get_shape().as_list()[1:3]
63         # shape (N, 1, 1, rnn_units) -> (N, 1, 1, attention_units)
64         hidden_state_projection = self.hidden_state_projector(hidden_state, **kwargs)
65         # shape (N, H, W, vgg_units) -> (N, H, W, attention_units)
66         features_projection = self.features_projector(features, **kwargs)
67         projection = tf.math.tanh(hidden_state_projection + features_projection)
68         # shape (N, H, W, attention_units) -> (N, H, W, 1)
69         attention = self.attention_projector(projection, **kwargs)
70         # shape (N, H, W, 1) -> (N, H * W)
71         attention = self.flatten(attention)
72         attention = tf.nn.softmax(attention)
73         # shape (N, H * W) -> (N, H, W, 1)
74         attention_map = tf.reshape(attention, [-1, H, W, 1])
75         glimpse = tf.math.multiply(features, attention_map)
76         # shape (N, H * W) -> (N, 1)
77         glimpse = tf.reduce_sum(glimpse, axis=[1, 2])
78         return glimpse
79 
80 
81 class SARDecoder(layers.Layer, NestedObject):
82     """Implements decoder module of the SAR model
83 
84     Args:
85         rnn_units: number of hidden units in recurrent cells
86         max_length: maximum length of a sequence
87         vocab_size: number of classes in the model alphabet
88         embedding_units: number of hidden embedding units
89         attention_units: number of hidden attention units
90         num_decoder_layers: number of LSTM layers to stack
91 
92 
93     """
94     def __init__(
95         self,
96         rnn_units: int,
97         max_length: int,
98         vocab_size: int,
99         embedding_units: int,
100         attention_units: int,
101         num_decoder_layers: int = 2
102     ) -> None:
103 
104         super().__init__()
105         self.vocab_size = vocab_size
106         self.embed = layers.Dense(embedding_units, use_bias=False)
107         self.attention_module = AttentionModule(attention_units)
108         self.output_dense = layers.Dense(vocab_size + 1, use_bias=True)
109         self.max_length = max_length
110         self.lstm_decoder = layers.StackedRNNCells(
111             [layers.LSTMCell(rnn_units, dtype=tf.float32, implementation=1) for _ in range(num_decoder_layers)]
112         )
113 
114     def call(
115         self,
116         features: tf.Tensor,
117         holistic: tf.Tensor,
118         **kwargs: Any,
119     ) -> tf.Tensor:
120 
121         # initialize states (each of shape (N, rnn_units))
122         states = self.lstm_decoder.get_initial_state(
123             inputs=None, batch_size=features.shape[0], dtype=tf.float32
124         )
125         # run first step of lstm
126         # holistic: shape (N, rnn_units)
127         _, states = self.lstm_decoder(holistic, states, **kwargs)
128         # Initialize with the index of virtual START symbol (placed after <eos>)
129         symbol = tf.fill(features.shape[0], self.vocab_size + 1)
130         logits_list = []
131         for _ in range(self.max_length + 1):  # keep 1 step for <eos>
132             # one-hot symbol with depth vocab_size + 1
133             # embeded_symbol: shape (N, embedding_units)
134             embeded_symbol = self.embed(tf.one_hot(symbol, depth=self.vocab_size + 1), **kwargs)
135             logits, states = self.lstm_decoder(embeded_symbol, states, **kwargs)
136             glimpse = self.attention_module(
137                 features, tf.expand_dims(tf.expand_dims(logits, axis=1), axis=1), **kwargs,
138             )
139             # logits: shape (N, rnn_units), glimpse: shape (N, 1)
140             logits = tf.concat([logits, glimpse], axis=-1)
141             # shape (N, rnn_units + 1) -> (N, vocab_size + 1)
142             logits = self.output_dense(logits, **kwargs)
143             # update symbol with predicted logits for t+1 step
144             symbol = tf.argmax(logits, axis=-1)
145             logits_list.append(logits)
146         outputs = tf.stack(logits_list, axis=1)  # shape (N, max_length + 1, vocab_size + 1)
147 
148         return outputs
149 
150 
151 class SAR(RecognitionModel):
152     """Implements a SAR architecture as described in `"Show, Attend and Read:A Simple and Strong Baseline for
153     Irregular Text Recognition" <https://arxiv.org/pdf/1811.00751.pdf>`_.
154 
155     Args:
156         feature_extractor: the backbone serving as feature extractor
157         vocab_size: size of the alphabet
158         rnn_units: number of hidden units in both encoder and decoder LSTM
159         embedding_units: number of embedding units
160         attention_units: number of hidden units in attention module
161         max_length: maximum word length handled by the model
162         num_decoders: number of LSTM to stack in decoder layer
163 
164     """
165 
166     _children_names: List[str] = ['feat_extractor', 'encoder', 'decoder']
167 
168     def __init__(
169         self,
170         feature_extractor,
171         vocab_size: int = 110,
172         rnn_units: int = 512,
173         embedding_units: int = 512,
174         attention_units: int = 512,
175         max_length: int = 30,
176         num_decoders: int = 2,
177         cfg: Optional[Dict[str, Any]] = None,
178     ) -> None:
179 
180         super().__init__(cfg=cfg)
181 
182         self.feat_extractor = feature_extractor
183 
184         self.encoder = Sequential(
185             [
186                 layers.LSTM(units=rnn_units, return_sequences=True),
187                 layers.LSTM(units=rnn_units, return_sequences=False)
188             ]
189         )
190 
191         self.decoder = SARDecoder(
192             rnn_units, max_length, vocab_size, embedding_units, attention_units, num_decoders,
193 
194         )
195 
196     def call(
197         self,
198         x: tf.Tensor,
199         **kwargs: Any,
200     ) -> tf.Tensor:
201 
202         features = self.feat_extractor(x, **kwargs)
203         pooled_features = tf.reduce_max(features, axis=1)  # vertical max pooling
204         encoded = self.encoder(pooled_features, **kwargs)
205         decoded = self.decoder(features, encoded, **kwargs)
206 
207         return decoded
208 
209 
210 class SARPostProcessor(RecognitionPostProcessor):
211     """Post processor for SAR architectures
212 
213     Args:
214         vocab: string containing the ordered sequence of supported characters
215         ignore_case: if True, ignore case of letters
216         ignore_accents: if True, ignore accents of letters
217     """
218 
219     def __call__(
220         self,
221         logits: tf.Tensor,
222     ) -> List[str]:
223         # compute pred with argmax for attention models
224         pred = tf.math.argmax(logits, axis=2)
225 
226         # decode raw output of the model with tf_label_to_idx
227         pred = tf.cast(pred, dtype='int32')
228         decoded_strings_pred = tf.strings.reduce_join(inputs=tf.nn.embedding_lookup(self._embedding, pred), axis=-1)
229         decoded_strings_pred = tf.strings.split(decoded_strings_pred, "<eos>")
230         decoded_strings_pred = tf.sparse.to_dense(decoded_strings_pred.to_sparse(), default_value='not valid')[:, 0]
231         words_list = [word.decode() for word in list(decoded_strings_pred.numpy())]
232 
233         if self.ignore_case:
234             words_list = [word.lower() for word in words_list]
235 
236         if self.ignore_accents:
237             raise NotImplementedError
238 
239         return words_list
240 
241 
242 def _sar_vgg(arch: str, pretrained: bool, input_shape: Tuple[int, int, int] = None, **kwargs: Any) -> SAR:
243 
244     # Patch the config
245     _cfg = deepcopy(default_cfgs[arch])
246     _cfg['input_shape'] = input_shape or _cfg['input_shape']
247     _cfg['vocab_size'] = kwargs.get('vocab_size', len(_cfg['vocab']))
248     _cfg['rnn_units'] = kwargs.get('rnn_units', _cfg['rnn_units'])
249     _cfg['embedding_units'] = kwargs.get('embedding_units', _cfg['rnn_units'])
250     _cfg['attention_units'] = kwargs.get('attention_units', _cfg['rnn_units'])
251     _cfg['max_length'] = kwargs.get('max_length', _cfg['max_length'])
252     _cfg['num_decoders'] = kwargs.get('num_decoders', _cfg['num_decoders'])
253 
254     # Feature extractor
255     feat_extractor = vgg.__dict__[default_cfgs[arch]['backbone']](
256         input_shape=_cfg['input_shape'],
257         include_top=False,
258     )
259 
260     kwargs['vocab_size'] = _cfg['vocab_size']
261     kwargs['rnn_units'] = _cfg['rnn_units']
262     kwargs['embedding_units'] = _cfg['embedding_units']
263     kwargs['attention_units'] = _cfg['attention_units']
264     kwargs['max_length'] = _cfg['max_length']
265     kwargs['num_decoders'] = _cfg['num_decoders']
266 
267     # Build the model
268     model = SAR(feat_extractor, cfg=_cfg, **kwargs)
269     # Load pretrained parameters
270     if pretrained:
271         load_pretrained_params(model, default_cfgs[arch]['url'])
272 
273     return model
274 
275 
276 def sar_vgg16_bn(pretrained: bool = False, **kwargs: Any) -> SAR:
277     """SAR with a VGG16 feature extractor as described in `"Show, Attend and Read:A Simple and Strong
278     Baseline for Irregular Text Recognition" <https://arxiv.org/pdf/1811.00751.pdf>`_.
279 
280     Example::
281         >>> import tensorflow as tf
282         >>> from doctr.models import sar_vgg16_bn
283         >>> model = sar_vgg16_bn(pretrained=False)
284         >>> input_tensor = tf.random.uniform(shape=[1, 64, 256, 3], maxval=1, dtype=tf.float32)
285         >>> out = model(input_tensor)
286 
287     Args:
288         pretrained (bool): If True, returns a model pre-trained on ImageNet
289 
290     Returns:
291         text recognition architecture
292     """
293 
294     return _sar_vgg('sar_vgg16_bn', pretrained, **kwargs)
295
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/doctr/models/recognition/sar.py b/doctr/models/recognition/sar.py
--- a/doctr/models/recognition/sar.py
+++ b/doctr/models/recognition/sar.py
@@ -89,7 +89,6 @@
         attention_units: number of hidden attention units
         num_decoder_layers: number of LSTM layers to stack
 
-
     """
     def __init__(
         self,
@@ -98,7 +97,7 @@
         vocab_size: int,
         embedding_units: int,
         attention_units: int,
-        num_decoder_layers: int = 2
+        num_decoder_layers: int = 2,
     ) -> None:
 
         super().__init__()
@@ -115,6 +114,7 @@
         self,
         features: tf.Tensor,
         holistic: tf.Tensor,
+        labels: Optional[tf.sparse.SparseTensor] = None,
         **kwargs: Any,
     ) -> tf.Tensor:
 
@@ -128,7 +128,7 @@
         # Initialize with the index of virtual START symbol (placed after <eos>)
         symbol = tf.fill(features.shape[0], self.vocab_size + 1)
         logits_list = []
-        for _ in range(self.max_length + 1):  # keep 1 step for <eos>
+        for t in range(self.max_length + 1):  # keep 1 step for <eos>
             # one-hot symbol with depth vocab_size + 1
             # embeded_symbol: shape (N, embedding_units)
             embeded_symbol = self.embed(tf.one_hot(symbol, depth=self.vocab_size + 1), **kwargs)
@@ -141,7 +141,13 @@
             # shape (N, rnn_units + 1) -> (N, vocab_size + 1)
             logits = self.output_dense(logits, **kwargs)
             # update symbol with predicted logits for t+1 step
-            symbol = tf.argmax(logits, axis=-1)
+            if kwargs.get('training'):
+                dense_labels = tf.sparse.to_dense(
+                    labels, default_value=self.vocab_size
+                )
+                symbol = dense_labels[:, t]
+            else:
+                symbol = tf.argmax(logits, axis=-1)
             logits_list.append(logits)
         outputs = tf.stack(logits_list, axis=1)  # shape (N, max_length + 1, vocab_size + 1)
 
@@ -196,13 +202,19 @@
     def call(
         self,
         x: tf.Tensor,
+        labels: Optional[tf.sparse.SparseTensor] = None,
         **kwargs: Any,
     ) -> tf.Tensor:
 
         features = self.feat_extractor(x, **kwargs)
         pooled_features = tf.reduce_max(features, axis=1)  # vertical max pooling
         encoded = self.encoder(pooled_features, **kwargs)
-        decoded = self.decoder(features, encoded, **kwargs)
+        if kwargs.get('training'):
+            if labels is None:
+                raise ValueError('Need to provide labels during training for teacher forcing')
+            decoded = self.decoder(features, encoded, labels, **kwargs)
+        else:
+            decoded = self.decoder(features, encoded, **kwargs)
 
         return decoded

{"golden_diff": "diff --git a/doctr/models/recognition/sar.py b/doctr/models/recognition/sar.py\n--- a/doctr/models/recognition/sar.py\n+++ b/doctr/models/recognition/sar.py\n@@ -89,7 +89,6 @@\n         attention_units: number of hidden attention units\n         num_decoder_layers: number of LSTM layers to stack\n \n-\n     \"\"\"\n     def __init__(\n         self,\n@@ -98,7 +97,7 @@\n         vocab_size: int,\n         embedding_units: int,\n         attention_units: int,\n-        num_decoder_layers: int = 2\n+        num_decoder_layers: int = 2,\n     ) -> None:\n \n         super().__init__()\n@@ -115,6 +114,7 @@\n         self,\n         features: tf.Tensor,\n         holistic: tf.Tensor,\n+        labels: Optional[tf.sparse.SparseTensor] = None,\n         **kwargs: Any,\n     ) -> tf.Tensor:\n \n@@ -128,7 +128,7 @@\n         # Initialize with the index of virtual START symbol (placed after <eos>)\n         symbol = tf.fill(features.shape[0], self.vocab_size + 1)\n         logits_list = []\n-        for _ in range(self.max_length + 1):  # keep 1 step for <eos>\n+        for t in range(self.max_length + 1):  # keep 1 step for <eos>\n             # one-hot symbol with depth vocab_size + 1\n             # embeded_symbol: shape (N, embedding_units)\n             embeded_symbol = self.embed(tf.one_hot(symbol, depth=self.vocab_size + 1), **kwargs)\n@@ -141,7 +141,13 @@\n             # shape (N, rnn_units + 1) -> (N, vocab_size + 1)\n             logits = self.output_dense(logits, **kwargs)\n             # update symbol with predicted logits for t+1 step\n-            symbol = tf.argmax(logits, axis=-1)\n+            if kwargs.get('training'):\n+                dense_labels = tf.sparse.to_dense(\n+                    labels, default_value=self.vocab_size\n+                )\n+                symbol = dense_labels[:, t]\n+            else:\n+                symbol = tf.argmax(logits, axis=-1)\n             logits_list.append(logits)\n         outputs = tf.stack(logits_list, axis=1)  # shape (N, max_length + 1, vocab_size + 1)\n \n@@ -196,13 +202,19 @@\n     def call(\n         self,\n         x: tf.Tensor,\n+        labels: Optional[tf.sparse.SparseTensor] = None,\n         **kwargs: Any,\n     ) -> tf.Tensor:\n \n         features = self.feat_extractor(x, **kwargs)\n         pooled_features = tf.reduce_max(features, axis=1)  # vertical max pooling\n         encoded = self.encoder(pooled_features, **kwargs)\n-        decoded = self.decoder(features, encoded, **kwargs)\n+        if kwargs.get('training'):\n+            if labels is None:\n+                raise ValueError('Need to provide labels during training for teacher forcing')\n+            decoded = self.decoder(features, encoded, labels, **kwargs)\n+        else:\n+            decoded = self.decoder(features, encoded, **kwargs)\n \n         return decoded\n", "issue": "[Models] Add SAR teacher for\u00e7ing during training (decoder)\nTeacher for\u00e7ing (feeding the LSTM decoder of SAR with ground-truth characters during training) is not implemented, and it should really improve performances\n", "before_files": [{"content": "# Copyright (C) 2021, Mindee.\n\n# This program is licensed under the Apache License version 2.\n# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.\n\nfrom copy import deepcopy\nimport tensorflow as tf\nfrom tensorflow.keras import Sequential, layers\nfrom typing import Tuple, Dict, List, Any, Optional\n\nfrom .. import vgg\nfrom ..utils import load_pretrained_params\nfrom .core import RecognitionModel\nfrom .core import RecognitionPostProcessor\nfrom doctr.utils.repr import NestedObject\n\n__all__ = ['SAR', 'SARPostProcessor', 'sar_vgg16_bn']\n\ndefault_cfgs: Dict[str, Dict[str, Any]] = {\n    'sar_vgg16_bn': {\n        'backbone': 'vgg16_bn', 'rnn_units': 512, 'max_length': 40, 'num_decoders': 2,\n        'input_shape': (64, 256, 3),\n        'post_processor': 'SARPostProcessor',\n        'vocab': ('3K}7e\u00e9;5\u00e0\u00ceYho]QwV6qU~W\"XnbBvcADf\u00cbmy.9\u00d4p\u00db*{C\u00f4\u00efE%M4#\u00c8R:g@T$x?0\u00ee\u00a3|za1\u00f98,OG\u20acP-'\n                  'k\u00e7H\u00eb\u00c0\u00c22\u00c9/\u00fbIJ\\'j(LN\u00d9Fut[)\u00e8Zs+&\u00b0Sd=\u00cf!<\u00e2_\u00c7>r\u00eai`l'),\n        'url': None,\n    },\n}\n\n\nclass AttentionModule(layers.Layer, NestedObject):\n    \"\"\"Implements attention module of the SAR model\n\n    Args:\n        attention_units: number of hidden attention units\n\n    \"\"\"\n    def __init__(\n        self,\n        attention_units: int\n    ) -> None:\n\n        super().__init__()\n        self.hidden_state_projector = layers.Conv2D(\n            attention_units, 1, strides=1, use_bias=False, padding='same', kernel_initializer='he_normal',\n        )\n        self.features_projector = layers.Conv2D(\n            attention_units, 3, strides=1, use_bias=True, padding='same', kernel_initializer='he_normal',\n        )\n        self.attention_projector = layers.Conv2D(\n            1, 1, strides=1, use_bias=False, padding=\"same\", kernel_initializer='he_normal',\n        )\n        self.flatten = layers.Flatten()\n\n    def call(\n        self,\n        features: tf.Tensor,\n        hidden_state: tf.Tensor,\n        **kwargs: Any,\n    ) -> tf.Tensor:\n\n        [H, W] = features.get_shape().as_list()[1:3]\n        # shape (N, 1, 1, rnn_units) -> (N, 1, 1, attention_units)\n        hidden_state_projection = self.hidden_state_projector(hidden_state, **kwargs)\n        # shape (N, H, W, vgg_units) -> (N, H, W, attention_units)\n        features_projection = self.features_projector(features, **kwargs)\n        projection = tf.math.tanh(hidden_state_projection + features_projection)\n        # shape (N, H, W, attention_units) -> (N, H, W, 1)\n        attention = self.attention_projector(projection, **kwargs)\n        # shape (N, H, W, 1) -> (N, H * W)\n        attention = self.flatten(attention)\n        attention = tf.nn.softmax(attention)\n        # shape (N, H * W) -> (N, H, W, 1)\n        attention_map = tf.reshape(attention, [-1, H, W, 1])\n        glimpse = tf.math.multiply(features, attention_map)\n        # shape (N, H * W) -> (N, 1)\n        glimpse = tf.reduce_sum(glimpse, axis=[1, 2])\n        return glimpse\n\n\nclass SARDecoder(layers.Layer, NestedObject):\n    \"\"\"Implements decoder module of the SAR model\n\n    Args:\n        rnn_units: number of hidden units in recurrent cells\n        max_length: maximum length of a sequence\n        vocab_size: number of classes in the model alphabet\n        embedding_units: number of hidden embedding units\n        attention_units: number of hidden attention units\n        num_decoder_layers: number of LSTM layers to stack\n\n\n    \"\"\"\n    def __init__(\n        self,\n        rnn_units: int,\n        max_length: int,\n        vocab_size: int,\n        embedding_units: int,\n        attention_units: int,\n        num_decoder_layers: int = 2\n    ) -> None:\n\n        super().__init__()\n        self.vocab_size = vocab_size\n        self.embed = layers.Dense(embedding_units, use_bias=False)\n        self.attention_module = AttentionModule(attention_units)\n        self.output_dense = layers.Dense(vocab_size + 1, use_bias=True)\n        self.max_length = max_length\n        self.lstm_decoder = layers.StackedRNNCells(\n            [layers.LSTMCell(rnn_units, dtype=tf.float32, implementation=1) for _ in range(num_decoder_layers)]\n        )\n\n    def call(\n        self,\n        features: tf.Tensor,\n        holistic: tf.Tensor,\n        **kwargs: Any,\n    ) -> tf.Tensor:\n\n        # initialize states (each of shape (N, rnn_units))\n        states = self.lstm_decoder.get_initial_state(\n            inputs=None, batch_size=features.shape[0], dtype=tf.float32\n        )\n        # run first step of lstm\n        # holistic: shape (N, rnn_units)\n        _, states = self.lstm_decoder(holistic, states, **kwargs)\n        # Initialize with the index of virtual START symbol (placed after <eos>)\n        symbol = tf.fill(features.shape[0], self.vocab_size + 1)\n        logits_list = []\n        for _ in range(self.max_length + 1):  # keep 1 step for <eos>\n            # one-hot symbol with depth vocab_size + 1\n            # embeded_symbol: shape (N, embedding_units)\n            embeded_symbol = self.embed(tf.one_hot(symbol, depth=self.vocab_size + 1), **kwargs)\n            logits, states = self.lstm_decoder(embeded_symbol, states, **kwargs)\n            glimpse = self.attention_module(\n                features, tf.expand_dims(tf.expand_dims(logits, axis=1), axis=1), **kwargs,\n            )\n            # logits: shape (N, rnn_units), glimpse: shape (N, 1)\n            logits = tf.concat([logits, glimpse], axis=-1)\n            # shape (N, rnn_units + 1) -> (N, vocab_size + 1)\n            logits = self.output_dense(logits, **kwargs)\n            # update symbol with predicted logits for t+1 step\n            symbol = tf.argmax(logits, axis=-1)\n            logits_list.append(logits)\n        outputs = tf.stack(logits_list, axis=1)  # shape (N, max_length + 1, vocab_size + 1)\n\n        return outputs\n\n\nclass SAR(RecognitionModel):\n    \"\"\"Implements a SAR architecture as described in `\"Show, Attend and Read:A Simple and Strong Baseline for\n    Irregular Text Recognition\" <https://arxiv.org/pdf/1811.00751.pdf>`_.\n\n    Args:\n        feature_extractor: the backbone serving as feature extractor\n        vocab_size: size of the alphabet\n        rnn_units: number of hidden units in both encoder and decoder LSTM\n        embedding_units: number of embedding units\n        attention_units: number of hidden units in attention module\n        max_length: maximum word length handled by the model\n        num_decoders: number of LSTM to stack in decoder layer\n\n    \"\"\"\n\n    _children_names: List[str] = ['feat_extractor', 'encoder', 'decoder']\n\n    def __init__(\n        self,\n        feature_extractor,\n        vocab_size: int = 110,\n        rnn_units: int = 512,\n        embedding_units: int = 512,\n        attention_units: int = 512,\n        max_length: int = 30,\n        num_decoders: int = 2,\n        cfg: Optional[Dict[str, Any]] = None,\n    ) -> None:\n\n        super().__init__(cfg=cfg)\n\n        self.feat_extractor = feature_extractor\n\n        self.encoder = Sequential(\n            [\n                layers.LSTM(units=rnn_units, return_sequences=True),\n                layers.LSTM(units=rnn_units, return_sequences=False)\n            ]\n        )\n\n        self.decoder = SARDecoder(\n            rnn_units, max_length, vocab_size, embedding_units, attention_units, num_decoders,\n\n        )\n\n    def call(\n        self,\n        x: tf.Tensor,\n        **kwargs: Any,\n    ) -> tf.Tensor:\n\n        features = self.feat_extractor(x, **kwargs)\n        pooled_features = tf.reduce_max(features, axis=1)  # vertical max pooling\n        encoded = self.encoder(pooled_features, **kwargs)\n        decoded = self.decoder(features, encoded, **kwargs)\n\n        return decoded\n\n\nclass SARPostProcessor(RecognitionPostProcessor):\n    \"\"\"Post processor for SAR architectures\n\n    Args:\n        vocab: string containing the ordered sequence of supported characters\n        ignore_case: if True, ignore case of letters\n        ignore_accents: if True, ignore accents of letters\n    \"\"\"\n\n    def __call__(\n        self,\n        logits: tf.Tensor,\n    ) -> List[str]:\n        # compute pred with argmax for attention models\n        pred = tf.math.argmax(logits, axis=2)\n\n        # decode raw output of the model with tf_label_to_idx\n        pred = tf.cast(pred, dtype='int32')\n        decoded_strings_pred = tf.strings.reduce_join(inputs=tf.nn.embedding_lookup(self._embedding, pred), axis=-1)\n        decoded_strings_pred = tf.strings.split(decoded_strings_pred, \"<eos>\")\n        decoded_strings_pred = tf.sparse.to_dense(decoded_strings_pred.to_sparse(), default_value='not valid')[:, 0]\n        words_list = [word.decode() for word in list(decoded_strings_pred.numpy())]\n\n        if self.ignore_case:\n            words_list = [word.lower() for word in words_list]\n\n        if self.ignore_accents:\n            raise NotImplementedError\n\n        return words_list\n\n\ndef _sar_vgg(arch: str, pretrained: bool, input_shape: Tuple[int, int, int] = None, **kwargs: Any) -> SAR:\n\n    # Patch the config\n    _cfg = deepcopy(default_cfgs[arch])\n    _cfg['input_shape'] = input_shape or _cfg['input_shape']\n    _cfg['vocab_size'] = kwargs.get('vocab_size', len(_cfg['vocab']))\n    _cfg['rnn_units'] = kwargs.get('rnn_units', _cfg['rnn_units'])\n    _cfg['embedding_units'] = kwargs.get('embedding_units', _cfg['rnn_units'])\n    _cfg['attention_units'] = kwargs.get('attention_units', _cfg['rnn_units'])\n    _cfg['max_length'] = kwargs.get('max_length', _cfg['max_length'])\n    _cfg['num_decoders'] = kwargs.get('num_decoders', _cfg['num_decoders'])\n\n    # Feature extractor\n    feat_extractor = vgg.__dict__[default_cfgs[arch]['backbone']](\n        input_shape=_cfg['input_shape'],\n        include_top=False,\n    )\n\n    kwargs['vocab_size'] = _cfg['vocab_size']\n    kwargs['rnn_units'] = _cfg['rnn_units']\n    kwargs['embedding_units'] = _cfg['embedding_units']\n    kwargs['attention_units'] = _cfg['attention_units']\n    kwargs['max_length'] = _cfg['max_length']\n    kwargs['num_decoders'] = _cfg['num_decoders']\n\n    # Build the model\n    model = SAR(feat_extractor, cfg=_cfg, **kwargs)\n    # Load pretrained parameters\n    if pretrained:\n        load_pretrained_params(model, default_cfgs[arch]['url'])\n\n    return model\n\n\ndef sar_vgg16_bn(pretrained: bool = False, **kwargs: Any) -> SAR:\n    \"\"\"SAR with a VGG16 feature extractor as described in `\"Show, Attend and Read:A Simple and Strong\n    Baseline for Irregular Text Recognition\" <https://arxiv.org/pdf/1811.00751.pdf>`_.\n\n    Example::\n        >>> import tensorflow as tf\n        >>> from doctr.models import sar_vgg16_bn\n        >>> model = sar_vgg16_bn(pretrained=False)\n        >>> input_tensor = tf.random.uniform(shape=[1, 64, 256, 3], maxval=1, dtype=tf.float32)\n        >>> out = model(input_tensor)\n\n    Args:\n        pretrained (bool): If True, returns a model pre-trained on ImageNet\n\n    Returns:\n        text recognition architecture\n    \"\"\"\n\n    return _sar_vgg('sar_vgg16_bn', pretrained, **kwargs)\n", "path": "doctr/models/recognition/sar.py"}], "after_files": [{"content": "# Copyright (C) 2021, Mindee.\n\n# This program is licensed under the Apache License version 2.\n# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.\n\nfrom copy import deepcopy\nimport tensorflow as tf\nfrom tensorflow.keras import Sequential, layers\nfrom typing import Tuple, Dict, List, Any, Optional\n\nfrom .. import vgg\nfrom ..utils import load_pretrained_params\nfrom .core import RecognitionModel\nfrom .core import RecognitionPostProcessor\nfrom doctr.utils.repr import NestedObject\n\n__all__ = ['SAR', 'SARPostProcessor', 'sar_vgg16_bn']\n\ndefault_cfgs: Dict[str, Dict[str, Any]] = {\n    'sar_vgg16_bn': {\n        'backbone': 'vgg16_bn', 'rnn_units': 512, 'max_length': 40, 'num_decoders': 2,\n        'input_shape': (64, 256, 3),\n        'post_processor': 'SARPostProcessor',\n        'vocab': ('3K}7e\u00e9;5\u00e0\u00ceYho]QwV6qU~W\"XnbBvcADf\u00cbmy.9\u00d4p\u00db*{C\u00f4\u00efE%M4#\u00c8R:g@T$x?0\u00ee\u00a3|za1\u00f98,OG\u20acP-'\n                  'k\u00e7H\u00eb\u00c0\u00c22\u00c9/\u00fbIJ\\'j(LN\u00d9Fut[)\u00e8Zs+&\u00b0Sd=\u00cf!<\u00e2_\u00c7>r\u00eai`l'),\n        'url': None,\n    },\n}\n\n\nclass AttentionModule(layers.Layer, NestedObject):\n    \"\"\"Implements attention module of the SAR model\n\n    Args:\n        attention_units: number of hidden attention units\n\n    \"\"\"\n    def __init__(\n        self,\n        attention_units: int\n    ) -> None:\n\n        super().__init__()\n        self.hidden_state_projector = layers.Conv2D(\n            attention_units, 1, strides=1, use_bias=False, padding='same', kernel_initializer='he_normal',\n        )\n        self.features_projector = layers.Conv2D(\n            attention_units, 3, strides=1, use_bias=True, padding='same', kernel_initializer='he_normal',\n        )\n        self.attention_projector = layers.Conv2D(\n            1, 1, strides=1, use_bias=False, padding=\"same\", kernel_initializer='he_normal',\n        )\n        self.flatten = layers.Flatten()\n\n    def call(\n        self,\n        features: tf.Tensor,\n        hidden_state: tf.Tensor,\n        **kwargs: Any,\n    ) -> tf.Tensor:\n\n        [H, W] = features.get_shape().as_list()[1:3]\n        # shape (N, 1, 1, rnn_units) -> (N, 1, 1, attention_units)\n        hidden_state_projection = self.hidden_state_projector(hidden_state, **kwargs)\n        # shape (N, H, W, vgg_units) -> (N, H, W, attention_units)\n        features_projection = self.features_projector(features, **kwargs)\n        projection = tf.math.tanh(hidden_state_projection + features_projection)\n        # shape (N, H, W, attention_units) -> (N, H, W, 1)\n        attention = self.attention_projector(projection, **kwargs)\n        # shape (N, H, W, 1) -> (N, H * W)\n        attention = self.flatten(attention)\n        attention = tf.nn.softmax(attention)\n        # shape (N, H * W) -> (N, H, W, 1)\n        attention_map = tf.reshape(attention, [-1, H, W, 1])\n        glimpse = tf.math.multiply(features, attention_map)\n        # shape (N, H * W) -> (N, 1)\n        glimpse = tf.reduce_sum(glimpse, axis=[1, 2])\n        return glimpse\n\n\nclass SARDecoder(layers.Layer, NestedObject):\n    \"\"\"Implements decoder module of the SAR model\n\n    Args:\n        rnn_units: number of hidden units in recurrent cells\n        max_length: maximum length of a sequence\n        vocab_size: number of classes in the model alphabet\n        embedding_units: number of hidden embedding units\n        attention_units: number of hidden attention units\n        num_decoder_layers: number of LSTM layers to stack\n\n    \"\"\"\n    def __init__(\n        self,\n        rnn_units: int,\n        max_length: int,\n        vocab_size: int,\n        embedding_units: int,\n        attention_units: int,\n        num_decoder_layers: int = 2,\n    ) -> None:\n\n        super().__init__()\n        self.vocab_size = vocab_size\n        self.embed = layers.Dense(embedding_units, use_bias=False)\n        self.attention_module = AttentionModule(attention_units)\n        self.output_dense = layers.Dense(vocab_size + 1, use_bias=True)\n        self.max_length = max_length\n        self.lstm_decoder = layers.StackedRNNCells(\n            [layers.LSTMCell(rnn_units, dtype=tf.float32, implementation=1) for _ in range(num_decoder_layers)]\n        )\n\n    def call(\n        self,\n        features: tf.Tensor,\n        holistic: tf.Tensor,\n        labels: Optional[tf.sparse.SparseTensor] = None,\n        **kwargs: Any,\n    ) -> tf.Tensor:\n\n        # initialize states (each of shape (N, rnn_units))\n        states = self.lstm_decoder.get_initial_state(\n            inputs=None, batch_size=features.shape[0], dtype=tf.float32\n        )\n        # run first step of lstm\n        # holistic: shape (N, rnn_units)\n        _, states = self.lstm_decoder(holistic, states, **kwargs)\n        # Initialize with the index of virtual START symbol (placed after <eos>)\n        symbol = tf.fill(features.shape[0], self.vocab_size + 1)\n        logits_list = []\n        for t in range(self.max_length + 1):  # keep 1 step for <eos>\n            # one-hot symbol with depth vocab_size + 1\n            # embeded_symbol: shape (N, embedding_units)\n            embeded_symbol = self.embed(tf.one_hot(symbol, depth=self.vocab_size + 1), **kwargs)\n            logits, states = self.lstm_decoder(embeded_symbol, states, **kwargs)\n            glimpse = self.attention_module(\n                features, tf.expand_dims(tf.expand_dims(logits, axis=1), axis=1), **kwargs,\n            )\n            # logits: shape (N, rnn_units), glimpse: shape (N, 1)\n            logits = tf.concat([logits, glimpse], axis=-1)\n            # shape (N, rnn_units + 1) -> (N, vocab_size + 1)\n            logits = self.output_dense(logits, **kwargs)\n            # update symbol with predicted logits for t+1 step\n            if kwargs.get('training'):\n                dense_labels = tf.sparse.to_dense(\n                    labels, default_value=self.vocab_size\n                )\n                symbol = dense_labels[:, t]\n            else:\n                symbol = tf.argmax(logits, axis=-1)\n            logits_list.append(logits)\n        outputs = tf.stack(logits_list, axis=1)  # shape (N, max_length + 1, vocab_size + 1)\n\n        return outputs\n\n\nclass SAR(RecognitionModel):\n    \"\"\"Implements a SAR architecture as described in `\"Show, Attend and Read:A Simple and Strong Baseline for\n    Irregular Text Recognition\" <https://arxiv.org/pdf/1811.00751.pdf>`_.\n\n    Args:\n        feature_extractor: the backbone serving as feature extractor\n        vocab_size: size of the alphabet\n        rnn_units: number of hidden units in both encoder and decoder LSTM\n        embedding_units: number of embedding units\n        attention_units: number of hidden units in attention module\n        max_length: maximum word length handled by the model\n        num_decoders: number of LSTM to stack in decoder layer\n\n    \"\"\"\n\n    _children_names: List[str] = ['feat_extractor', 'encoder', 'decoder']\n\n    def __init__(\n        self,\n        feature_extractor,\n        vocab_size: int = 110,\n        rnn_units: int = 512,\n        embedding_units: int = 512,\n        attention_units: int = 512,\n        max_length: int = 30,\n        num_decoders: int = 2,\n        cfg: Optional[Dict[str, Any]] = None,\n    ) -> None:\n\n        super().__init__(cfg=cfg)\n\n        self.feat_extractor = feature_extractor\n\n        self.encoder = Sequential(\n            [\n                layers.LSTM(units=rnn_units, return_sequences=True),\n                layers.LSTM(units=rnn_units, return_sequences=False)\n            ]\n        )\n\n        self.decoder = SARDecoder(\n            rnn_units, max_length, vocab_size, embedding_units, attention_units, num_decoders,\n\n        )\n\n    def call(\n        self,\n        x: tf.Tensor,\n        labels: Optional[tf.sparse.SparseTensor] = None,\n        **kwargs: Any,\n    ) -> tf.Tensor:\n\n        features = self.feat_extractor(x, **kwargs)\n        pooled_features = tf.reduce_max(features, axis=1)  # vertical max pooling\n        encoded = self.encoder(pooled_features, **kwargs)\n        if kwargs.get('training'):\n            if labels is None:\n                raise ValueError('Need to provide labels during training for teacher forcing')\n            decoded = self.decoder(features, encoded, labels, **kwargs)\n        else:\n            decoded = self.decoder(features, encoded, **kwargs)\n\n        return decoded\n\n\nclass SARPostProcessor(RecognitionPostProcessor):\n    \"\"\"Post processor for SAR architectures\n\n    Args:\n        vocab: string containing the ordered sequence of supported characters\n        ignore_case: if True, ignore case of letters\n        ignore_accents: if True, ignore accents of letters\n    \"\"\"\n\n    def __call__(\n        self,\n        logits: tf.Tensor,\n    ) -> List[str]:\n        # compute pred with argmax for attention models\n        pred = tf.math.argmax(logits, axis=2)\n\n        # decode raw output of the model with tf_label_to_idx\n        pred = tf.cast(pred, dtype='int32')\n        decoded_strings_pred = tf.strings.reduce_join(inputs=tf.nn.embedding_lookup(self._embedding, pred), axis=-1)\n        decoded_strings_pred = tf.strings.split(decoded_strings_pred, \"<eos>\")\n        decoded_strings_pred = tf.sparse.to_dense(decoded_strings_pred.to_sparse(), default_value='not valid')[:, 0]\n        words_list = [word.decode() for word in list(decoded_strings_pred.numpy())]\n\n        if self.ignore_case:\n            words_list = [word.lower() for word in words_list]\n\n        if self.ignore_accents:\n            raise NotImplementedError\n\n        return words_list\n\n\ndef _sar_vgg(arch: str, pretrained: bool, input_shape: Tuple[int, int, int] = None, **kwargs: Any) -> SAR:\n\n    # Patch the config\n    _cfg = deepcopy(default_cfgs[arch])\n    _cfg['input_shape'] = input_shape or _cfg['input_shape']\n    _cfg['vocab_size'] = kwargs.get('vocab_size', len(_cfg['vocab']))\n    _cfg['rnn_units'] = kwargs.get('rnn_units', _cfg['rnn_units'])\n    _cfg['embedding_units'] = kwargs.get('embedding_units', _cfg['rnn_units'])\n    _cfg['attention_units'] = kwargs.get('attention_units', _cfg['rnn_units'])\n    _cfg['max_length'] = kwargs.get('max_length', _cfg['max_length'])\n    _cfg['num_decoders'] = kwargs.get('num_decoders', _cfg['num_decoders'])\n\n    # Feature extractor\n    feat_extractor = vgg.__dict__[default_cfgs[arch]['backbone']](\n        input_shape=_cfg['input_shape'],\n        include_top=False,\n    )\n\n    kwargs['vocab_size'] = _cfg['vocab_size']\n    kwargs['rnn_units'] = _cfg['rnn_units']\n    kwargs['embedding_units'] = _cfg['embedding_units']\n    kwargs['attention_units'] = _cfg['attention_units']\n    kwargs['max_length'] = _cfg['max_length']\n    kwargs['num_decoders'] = _cfg['num_decoders']\n\n    # Build the model\n    model = SAR(feat_extractor, cfg=_cfg, **kwargs)\n    # Load pretrained parameters\n    if pretrained:\n        load_pretrained_params(model, default_cfgs[arch]['url'])\n\n    return model\n\n\ndef sar_vgg16_bn(pretrained: bool = False, **kwargs: Any) -> SAR:\n    \"\"\"SAR with a VGG16 feature extractor as described in `\"Show, Attend and Read:A Simple and Strong\n    Baseline for Irregular Text Recognition\" <https://arxiv.org/pdf/1811.00751.pdf>`_.\n\n    Example::\n        >>> import tensorflow as tf\n        >>> from doctr.models import sar_vgg16_bn\n        >>> model = sar_vgg16_bn(pretrained=False)\n        >>> input_tensor = tf.random.uniform(shape=[1, 64, 256, 3], maxval=1, dtype=tf.float32)\n        >>> out = model(input_tensor)\n\n    Args:\n        pretrained (bool): If True, returns a model pre-trained on ImageNet\n\n    Returns:\n        text recognition architecture\n    \"\"\"\n\n    return _sar_vgg('sar_vgg16_bn', pretrained, **kwargs)\n", "path": "doctr/models/recognition/sar.py"}]}