Improve demo

app.py

@@ -1,50 +1,64 @@
+import time
+
 from transformers import T5ForConditionalGeneration, T5Tokenizer, GenerationConfig
 import gradio as gr
 
 MODEL_NAME = "jbochi/madlad400-3b-mt"
 
-
-default_max_length = 200
-
-print("Using `{}`.".format(MODEL_NAME))
-
+print(f"Loading {MODEL_NAME} tokenizer...")
 tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
-print("T5Tokenizer loaded from pretrained.")
-
+print(f"Loading {MODEL_NAME} model...")
 model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, device_map="auto")
-print("T5ForConditionalGeneration loaded from pretrained.")
 
 
-def inference(max_length, input_text, history=[]):
-    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+def inference(input_text, target_language, max_length):
+    global model, tokenizer
+    start_time = time.time()
+    input_ids = tokenizer(
+        f"<2{target_language}> {input_text}", return_tensors="pt"
+    ).input_ids
     outputs = model.generate(
-        input_ids=input_ids,
-        generation_config=GenerationConfig(max_length=max_length, decoder_start_token_id=2),
+        input_ids=input_ids.to(model.device),
+        generation_config=GenerationConfig(max_length=max_length),
     )
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    history.append((input_text, result))
-    return history, history
-
+    end_time = time.time()
+    result = {
+        'result': result,
+        'inference_time': end_time - start_time,
+        'input_token_ids': input_ids[0].tolist(),
+        'output_token_ids': outputs[0].tolist(),
+    }
+    return result
 
-with gr.Blocks() as demo:
-    with gr.Row():
-        gr.Markdown(
-            "<h1>Demo of {}</h1><p>See more at Hugging Face: <a href='https://huggingface.co/{}'>{}</a>.</p>".format(
-                MODEL_NAME, MODEL_NAME, MODEL_NAME
-            )
-        )
-        max_length = gr.Number(
-            value=default_max_length, label="maximum length of response"
-        )
 
-    chatbot = gr.Chatbot(label=MODEL_NAME)
-    state = gr.State([])
-
-    with gr.Row():
-        txt = gr.Textbox(
-            show_label=False, placeholder="<2es> text to translate"
-        )
+def run():
+    tokens = [tokenizer.decode(i) for i in range(500)]
+    lang_codes = [token[2:-1] for token in tokens if token.startswith("<2")]
+    inputs = [
+        gr.components.Textbox(lines=5, label="Input text"),
+        gr.components.Dropdown(lang_codes, value="en", label="Target Language"),
+        gr.components.Slider(
+            minimum=5,
+            maximum=500,
+            value=200,
+            label="Max length",
+        ),
+    ]
+    outputs = gr.components.JSON()
+    title = f"{MODEL_NAME} demo"
+    demo_status = "Demo is running on CPU"
+    description = (
+        f"Details: https://huggingface.co/{MODEL_NAME}. {demo_status}"
+    )
+    gr.Interface(
+        inference,
+        inputs,
+        outputs,
+        title=title,
+        description=description,
+    ).launch()
 
-    txt.submit(fn=inference, inputs=[max_length, txt, state], outputs=[chatbot, state])
 
-demo.launch()
+if __name__ == "__main__":
+    run()