Update app.py

gradio_samples/bertviz/app.py

@@ -15,7 +15,7 @@ from tqdm.notebook import tqdm
 from torch.utils.data import DataLoader
 from functools import partial
 
-from transformers import AutoTokenizer, MarianTokenizer, AutoModel, AutoModelForSeq2SeqLM, MarianMTModel
+from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM
 
 from bertviz import model_view, head_view
 from bertviz_gradio import head_view_mod
@@ -32,7 +32,7 @@ tokenizer_es = AutoTokenizer.from_pretrained(model_es)
 # tokenizer_zh = AutoTokenizer.from_pretrained(model_zh)
 # tokenizer_sw = AutoTokenizer.from_pretrained(model_sw)
 
-model_tr_es = MarianMTModel.from_pretrained(model_es)
+model_tr_es = AutoModel.from_pretrained(model_es)
 # model_tr_fr = MarianMTModel.from_pretrained(model_fr)
 # model_tr_zh = MarianMTModel.from_pretrained(model_zh)
 # model_tr_sw = MarianMTModel.from_pretrained(model_sw)
@@ -196,4 +196,4 @@ with gr.Blocks(js="plotsjs_bertviz.js") as demo:
     # demo.load(None,None,None,js="plotsjs.js")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()