revert 2

app.py

@@ -1,6 +1,5 @@
 import torch
 from torchvision import transforms as T
-from transformers import AutoTokenizer
 import gradio as gr
 
 class App:
@@ -15,7 +14,6 @@ class App:
             T.ToTensor(),
             T.Normalize(0.5, 0.5)
         ])
-        self._tokenizer_cache = {}
 
     def _get_model(self, name):
         if name in self._model_cache:
@@ -24,33 +22,21 @@ 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)
-        
-        # 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)
-        
+        label, _ = model.tokenizer.decode(pred)
+        raw_label, raw_confidence = model.tokenizer.decode(pred, raw=True)
         # Format confidence values
-        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]
+        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]
+
 
 def main():
     app = App()
@@ -63,12 +49,14 @@ 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()