| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| checkpoint = "projecte-aina/m2m100-418M-ft-de-ca" | |
| model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint) | |
| tokenizer = AutoTokenizer.from_pretrained(checkpoint) | |
| tokenizer.src_lang = "de" | |
| tokenizer.tgt_lang = "ca" | |
| def predict(text): | |
| encoded_ref = tokenizer(text, return_tensors="pt") | |
| generated_tokens = model.generate(**encoded_ref, forced_bos_token_id=tokenizer.get_lang_id(hyp_lang)) | |
| hyp = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True) | |
| return hyp[0] | |
| gr.Interface( | |
| predict, | |
| inputs="text", | |
| outputs="text", | |
| title="Translation", | |
| ).launch() |