Update app.py

app.py

@@ -1,11 +1,24 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-tokenizer = AutoTokenizer.from_pretrained("TwentyNine/byt5-small-ainu-latinizer-lr-11")
-model = AutoModelForSeq2SeqLM.from_pretrained("TwentyNine/byt5-small-ainu-latinizer-lr-11")
+tokenizer = AutoTokenizer.from_pretrained("TwentyNine/byt5-small-ainu-latinizer-cos_w_restarts")
+model1 = AutoModelForSeq2SeqLM.from_pretrained("TwentyNine/byt5-small-ainu-latinizer-cos_w_restarts")
+model2 = AutoModelForSeq2SeqLM.from_pretrained("TwentyNine/byt5-small-ainu-latinizer-polynomial")
+model3 = AutoModelForSeq2SeqLM.from_pretrained("TwentyNine/byt5-small-ainu-latinizer-linear")
 
-def transcribe(input_str):
+def transcribe(input_str, model_index):
     output_str = ''
+    model = None
+    
+    match model_index:
+        case 1:
+            model = model1
+        case 2:
+            model = model2
+        case 3:
+            model = model3
+        case _:
+            model = model1
 
     for input in input_str.split('\n'):
         input_enc  = tokenizer.encode(input.strip(), return_tensors='pt')
@@ -20,7 +33,7 @@ def transcribe(input_str):
 
 gradio_app = gr.Interface(
     transcribe,
-    inputs=gr.Textbox(label='Input (kana)', value='トゥイマ　ヒ　ワ　エエㇰ　ワ　ヒオーイオイ。ピㇼカノ　ヌカㇻ　ヤン！', placeholder='トゥイマ　ヒ　ワ　エエㇰ　ワ　ヒオーイオイ。ピㇼカノ　ヌカㇻ　ヤン！', info='Ainu text written in Japanese katakana (input).', interactive=True, autofocus=True),
+    inputs=(gr.Textbox(label='Input (kana)', value='トゥイマ　ヒ　ワ　エエㇰ　ワ　ヒオーイオイ。ピㇼカノ　ヌカㇻ　ヤン！', placeholder='トゥイマ　ヒ　ワ　エエㇰ　ワ　ヒオーイオイ。ピㇼカノ　ヌカㇻ　ヤン！', info='Ainu text written in Japanese katakana (input).', interactive=True, autofocus=True)), gr.radio(label="Training scheduler type", choices=[("Cosine with Restarts", 1), ("Polynomial", 2), ("Linear", 3)]),
     outputs=gr.Textbox(label='Output (alphabet)', info='Ainu text written in the Latin alphabet (output).'),
     title='BYT5 Ainu Kana-Latin Converter (V1)',
     article='<p>Example sentence borrowed from <a href="https://www.hakusuisha.co.jp/book/b584600.html">New Express Ainu-go</a> by <a href="https://researchmap.jp/read0064265/?lang=english">Professor NAKAGAWA Hiroshi</a> of Chiba University.</p>'