import spaces import gradio as gr from transformers import MT5ForConditionalGeneration, MT5Tokenizer,T5ForConditionalGeneration, T5Tokenizer models = {"finetuned mt5-base":"alakxender/mt5-base-dv-en", "madlad400-3b":"google/madlad400-3b-mt"} def tranlate(text:str,model_name:str): if (len(text)>2000): raise gr.Error(f"Try smaller text, yours is {len(text)}. try to fit to 2000 chars.") if (model_name is None): raise gr.Error("huh! not sure what to do without a model. select a model.") if model_name =='finetuned mt5-base': return mt5_translate(text,model_name) else: return t5_tranlaste(text,model_name) @spaces.GPU(duration=120) def t5_tranlaste(text:str,model_name:str): model = T5ForConditionalGeneration.from_pretrained(models[model_name], device_map="auto") tokenizer = T5Tokenizer.from_pretrained(models[model_name]) text = f"<2en> {text}" input_ids = tokenizer(text, return_tensors="pt").input_ids.to(model.device) outputs = model.generate(input_ids=input_ids, max_new_tokens=1024*2) translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) return translated_text def mt5_translate(text:str, model_name:str): model = MT5ForConditionalGeneration.from_pretrained(models[model_name]) tokenizer = MT5Tokenizer.from_pretrained(models[model_name]) inputs = tokenizer(text, return_tensors="pt") result = model.generate(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], max_new_tokens=1024*2) translated_text = tokenizer.decode(result[0], skip_special_tokens=True) return translated_text css = """ .textbox1 textarea { font-size: 18px !important; font-family: 'MV_Faseyha', 'Faruma', 'A_Faruma' !important; line-height: 1.8 !important; } """ demo = gr.Interface( fn=tranlate, inputs= [ gr.Textbox(lines=5, label="Enter Dhivehi Text", rtl=True, elem_classes="textbox1"), gr.Dropdown(choices=list(models.keys()), label="Select a model", value="finetuned mt5-base"), ], css=css, outputs=gr.Textbox(label="English Translation"), title="Dhivehi to English Translation", description="Translate Dhivehi text to English", examples=[["މާލޭގައި ފެންބޮޑުވާ މަގުތައް މަރާމާތު ކުރަން ފަށައިފި"]] ) demo.launch()