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, AutoModel, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
 
 from bertviz import model_view, head_view
 from bertviz_gradio import head_view_mod
@@ -23,44 +23,44 @@ from bertviz_gradio import head_view_mod
 
 
 model_es = "Helsinki-NLP/opus-mt-en-es"
-# model_fr = "Helsinki-NLP/opus-mt-en-fr"
-# model_zh = "Helsinki-NLP/opus-mt-en-zh"
-# model_sw = "Helsinki-NLP/opus-mt-en-sw"
+model_fr = "Helsinki-NLP/opus-mt-en-fr"
+model_zh = "Helsinki-NLP/opus-mt-en-zh"
+model_sw = "Helsinki-NLP/opus-mt-en-sw"
 
 tokenizer_es = AutoTokenizer.from_pretrained(model_es)
-# tokenizer_fr = AutoTokenizer.from_pretrained(model_fr)
-# tokenizer_zh = AutoTokenizer.from_pretrained(model_zh)
-# tokenizer_sw = AutoTokenizer.from_pretrained(model_sw)
+tokenizer_fr = AutoTokenizer.from_pretrained(model_fr)
+tokenizer_zh = AutoTokenizer.from_pretrained(model_zh)
+tokenizer_sw = AutoTokenizer.from_pretrained(model_sw)
 
-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)
+model_tr_es =AutoModelForSeq2SeqLM.from_pretrained(model_es)
+model_tr_fr = AutoModelForSeq2SeqLM.from_pretrained(model_fr)
+model_tr_zh =AutoModelForSeq2SeqLM.from_pretrained(model_zh)
+model_tr_sw = AutoModelForSeq2SeqLM.from_pretrained(model_sw)
 
 model_es = inseq.load_model("Helsinki-NLP/opus-mt-en-es", "input_x_gradient")
-# model_fr = inseq.load_model("Helsinki-NLP/opus-mt-en-fr", "input_x_gradient")
-# model_zh = inseq.load_model("Helsinki-NLP/opus-mt-en-zh", "input_x_gradient")
-# model_sw = inseq.load_model("Helsinki-NLP/opus-mt-en-sw", "input_x_gradient")
+model_fr = inseq.load_model("Helsinki-NLP/opus-mt-en-fr", "input_x_gradient")
+model_zh = inseq.load_model("Helsinki-NLP/opus-mt-en-zh", "input_x_gradient")
+model_sw = inseq.load_model("Helsinki-NLP/opus-mt-en-sw", "input_x_gradient")
 
 dict_models = {
 	'en-es': model_es,
-	# 'en-fr': model_fr,
-	# 'en-zh': model_zh,
-	# 'en-sw': model_sw,
+	'en-fr': model_fr,
+	'en-zh': model_zh,
+	'en-sw': model_sw,
 }
 
 dict_models_tr = {
 	'en-es': model_tr_es,
-	# 'en-fr': model_tr_fr,
-	# 'en-zh': model_tr_zh,
-	# 'en-sw': model_tr_sw,
+	'en-fr': model_tr_fr,
+	'en-zh': model_tr_zh,
+	'en-sw': model_tr_sw,
 }
 
 dict_tokenizer_tr = {
 	'en-es': tokenizer_es,
-	# 'en-fr': tokenizer_fr,
-	# 'en-zh': tokenizer_zh,
-	# 'en-sw': tokenizer_sw,
+	'en-fr': tokenizer_fr,
+	'en-zh': tokenizer_zh,
+	'en-sw': tokenizer_sw,
 }
 
 saliency_examples = [
@@ -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(share=True)