Update app.py

app.py

@@ -1,45 +1,81 @@
 import spaces
 import gradio as gr
 import torch
-
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from optimum.bettertransformer import BetterTransformer
 
-tokenizer = AutoTokenizer.from_pretrained(
-    "google/madlad400-3b-mt",
-    use_fast=True
-)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model_hf = AutoModelForSeq2SeqLM.from_pretrained(
-    "google/madlad400-3b-mt",
-    torch_dtype=torch.bfloat16
-)
+tokenizer_3b_mt = AutoTokenizer.from_pretrained("google/madlad400-3b-mt", use_fast=True)
+language_codes = [token for token in tokenizer_3b_mt.get_vocab().keys() if token.startswith("<2")]
+remove_codes = ['<2>', '<2en_xx_simple>', '<2translate>', '<2back_translated>', '<2zxx_xx_dtynoise>', '<2transliterate>']
+language_codes = [token for token in language_codes if token not in remove_codes]
+
+model_choices = [
+    "google/madlad400-3b-mt", 
+    "google/madlad400-7b-mt", 
+    "google/madlad400-10b-mt", 
+    "google/madlad400-7b-mt-bt"
+]
 
-model = BetterTransformer.transform(model_hf, keep_original=True)
+model_resources = {}
+
+def load_tokenizer_model(model_name):
+    """
+    Load tokenizer and model for a chosen model name.
+    """
+    if model_name not in model_resources:
+        # Load tokenizer and model for first time
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.float16)
+        model.to_bettertransformer()
+        model.to(device)
+        model_resources[model_name] = (tokenizer, model)
+    return model_resources[model_name]
 
 @spaces.GPU
-def translate(text):
+def translate(text, target_language, model_name):
     """
-    Translates the input text from English to Hawaiian.
+    Translate the input text from English to another language.
     """
-    text = "<2haw> " + text
+    # Load tokenizer and model if not already loaded
+    tokenizer, model = load_tokenizer_model(model_name)
     
-    inputs = tokenizer(
-        text,
-        return_tensors="pt"
-    )
+    text = target_language + text
+    input_ids = tokenizer(text, return_tensors="pt").input_ids.to(device)
     
-    outputs = model.generate(**inputs, max_new_tokens=1000)
+    outputs = model.generate(input_ids=input_ids, max_new_tokens=128000)
     text_translated = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    
+
     return text_translated[0]
 
+title = "MADLAD-400 Translation"
+description = """
+Translation from English to over 400 languages based on [research](https://arxiv.org/pdf/2309.04662) by Google DeepMind and Google Research
+"""
+
+input_text = gr.Textbox(
+    label="Text",
+    placeholder="Enter text here"
+)
+target_language = gr.Dropdown(
+    choices=language_codes,
+    value="<2haw>",
+    label="Target language"
+)
+model_choice = gr.Dropdown(
+    choices=model_choices, 
+    value="google/madlad400-3b-mt", 
+    label="Model"
+)
+output_text = gr.Textbox(label="Translation")
+
 demo = gr.Interface(
     fn=translate,
-    inputs=[gr.Textbox(label="English")],
-    outputs=[gr.Textbox(label="Hawaiian")],
-    title="MADLAD-400-3B-MT English-to-Hawaiian Translation",
-    description="[Code](https://github.com/darylalim/madlad-400-3b-mt-eng-to-haw-translation)")
+    inputs=[input_text, target_language, model_choice],
+    outputs=output_text,
+    title=title,
+    description=description
+)
 
 demo.queue()