remove bigscience/bloom

app.py

@@ -2,25 +2,21 @@ import os
 import torch
 import gradio as gr
 import time
-from transformers import AutoTokenizer, AutoModel, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from flores200_codes import flores_codes
 
 
 def load_models():
     # build model and tokenizer
     model_name_dict = {
-        "bloom": "bigscience/bloom",
-        #"nllb-distilled-600M": "facebook/nllb-200-distilled-600M",
-        #'nllb-1.3B': 'facebook/nllb-200-1.3B',
-        #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
-        #'nllb-3.3B': 'facebook/nllb-200-3.3B',
+        "nllb-distilled-600M": "facebook/nllb-200-distilled-600M",
     }
 
     model_dict = {}
 
     for call_name, real_name in model_name_dict.items():
         print("\tLoading model: %s" % call_name)
-        model = AutoModel.from_pretrained(real_name)
+        model = AutoModelForSeq2SeqLMl.from_pretrained(real_name)
         tokenizer = AutoTokenizer.from_pretrained(real_name)
         model_dict[call_name + "_model"] = model
         model_dict[call_name + "_tokenizer"] = tokenizer
@@ -69,7 +65,6 @@ if __name__ == "__main__":
 
     # define gradio demo
     lang_codes = list(flores_codes.keys())
-    # inputs = [gr.inputs.Radio(['nllb-distilled-600M', 'nllb-1.3B', 'nllb-distilled-1.3B'], label='NLLB Model'),
     inputs = [
         gr.inputs.Dropdown(lang_codes, default="English", label="Source"),
         gr.inputs.Dropdown(lang_codes, default="Nepali", label="Target"),