code changed

app.py

@@ -1,5 +1,6 @@
 import torch
 from torchvision import transforms as T
+from transformers import AutoTokenizer
 import gradio as gr
 
 class App:
@@ -14,6 +15,7 @@ class App:
             T.ToTensor(),
             T.Normalize(0.5, 0.5)
         ])
+        self._tokenizer_cache = {}
 
     def _get_model(self, name):
         if name in self._model_cache:
@@ -22,21 +24,33 @@ class App:
         self._model_cache[name] = model
         return model
 
+    def _get_tokenizer(self, name):
+        if name in self._tokenizer_cache:
+            return self._tokenizer_cache[name]
+        tokenizer = AutoTokenizer.from_pretrained(name)
+        self._tokenizer_cache[name] = tokenizer
+        return tokenizer
+
     @torch.inference_mode()
     def __call__(self, model_name, image):
         if image is None:
             return '', []
         model = self._get_model(model_name)
+        tokenizer = self._get_tokenizer(model_name)
+
         image = self._preprocess(image.convert('RGB')).unsqueeze(0)
         # Greedy decoding
         pred = model(image).softmax(-1)
-        label, _ = model.tokenizer.decode(pred)
-        raw_label, raw_confidence = model.tokenizer.decode(pred, raw=True)
+        
+        # Tokenize input data
+        label = tokenizer.decode(pred.argmax(-1)[0].tolist(), skip_special_tokens=True)
+        raw_label, raw_confidence = tokenizer.decode(pred.argmax(-1)[0].tolist(), raw=True)
+        
         # Format confidence values
-        max_len = 25 if model_name == 'crnn' else len(label[0]) + 1
-        conf = list(map('{:0.1f}'.format, raw_confidence[0][:max_len].tolist()))
-        return label[0], [raw_label[0][:max_len], conf]
-
+        max_len = 25 if model_name == 'crnn' else len(label) + 1
+        conf = list(map('{:0.1f}'.format, pred[0, :, :max_len].tolist()))
+        
+        return label, [raw_label[:max_len], conf]
 
 def main():
     app = App()
@@ -49,14 +63,12 @@ def main():
                 read_upload = gr.Button('Read Text')
 
         output = gr.Textbox(max_lines=1, label='Model output')
-        #adv_output = gr.Checkbox(label='Show detailed output')
         raw_output = gr.Dataframe(row_count=2, col_count=0, label='Raw output with confidence values ([0, 1] interval; [B] - BLANK token; [E] - EOS token)')
 
         read_upload.click(app, inputs=[model_name, image_upload], outputs=[output, raw_output])
-        #adv_output.change(lambda x: gr.update(visible=x), inputs=adv_output, outputs=raw_output)
+
     demo.queue(max_size=20)
     demo.launch()
 
-
 if __name__ == '__main__':
-    main()
+    main()