remove big model

app.py

@@ -7,15 +7,11 @@ import unicodedata
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the big model
-big_tokenizer = NllbTokenizer.from_pretrained("hunterschep/amis-zh-3.3B")
-big_model = AutoModelForSeq2SeqLM.from_pretrained("hunterschep/amis-zh-3.3B").to(device)
-
 # Load the small model
 small_tokenizer = NllbTokenizer.from_pretrained("hunterschep/amis-zh-600M")
 small_model = AutoModelForSeq2SeqLM.from_pretrained("hunterschep/amis-zh-600M").to(device)
 
-# Fix tokenizers
+# Fix tokenizer
 def fix_tokenizer(tokenizer, new_lang='ami_Latn'):
     old_len = len(tokenizer) - int(new_lang in tokenizer.added_tokens_encoder)
     tokenizer.lang_code_to_id[new_lang] = old_len - 1
@@ -28,12 +24,11 @@ def fix_tokenizer(tokenizer, new_lang='ami_Latn'):
     tokenizer.added_tokens_encoder = {}
     tokenizer.added_tokens_decoder = {}
 
-fix_tokenizer(big_tokenizer)
 fix_tokenizer(small_tokenizer)
 
 # Translation function
-def translate(text, model_type, src_lang, tgt_lang):
-    tokenizer, model = (big_tokenizer, big_model) if model_type == "Large" else (small_tokenizer, small_model)
+def translate(text, src_lang, tgt_lang):
+    tokenizer, model = small_tokenizer, small_model
     if src_lang == "zho_Hant":
         text = preproc_chinese(text)
     tokenizer.src_lang = src_lang
@@ -71,7 +66,6 @@ def switch_direction(src_lang, tgt_lang):
 
 with gr.Blocks() as demo:
     gr.Markdown("# AMIS - Chinese Translation Tool")
-    model_type = gr.Radio(choices=["Small", "Large"], value="Small", label="Model Type")
     src_lang = gr.Radio(choices=["zho_Hant", "ami_Latn"], value="zho_Hant", label="Source Language")
     tgt_lang = gr.Radio(choices=["ami_Latn", "zho_Hant"], value="ami_Latn", label="Target Language")
     input_text = gr.Textbox(label="Input Text", placeholder="Enter text here...")
@@ -79,7 +73,7 @@ with gr.Blocks() as demo:
     translate_btn = gr.Button("Translate")
     switch_btn = gr.Button("Switch Direction")
 
-    translate_btn.click(translate, inputs=[input_text, model_type, src_lang, tgt_lang], outputs=output_text)
+    translate_btn.click(translate, inputs=[input_text, src_lang, tgt_lang], outputs=output_text)
     switch_btn.click(switch_direction, inputs=[src_lang, tgt_lang], outputs=[src_lang, tgt_lang])
 
 if __name__ == "__main__":