balaramas's picture
Update app.py
a5d0aa6 verified
raw
history blame
778 Bytes
import gradio as gr
import subprocess
subprocess.check_call(["pip", "install", "transformers"])
subprocess.check_call(["pip", "install", "torch"])
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
tokenizer = AutoTokenizer.from_pretrained("balaramas/mbart-enhiriser")
model = AutoModelForSeq2SeqLM.from_pretrained("balaramas/mbart-enhiriser")
def greet(ar_en):
tokenizer.src_lang = "en_XX"
encoded_ar = tokenizer(ar_en, return_tensors="pt")
generated_tokens = model.generate(
**encoded_ar,
forced_bos_token_id=tokenizer.lang_code_to_id["hi_IN"]
)
output = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
return output
iface = gr.Interface(fn=greet, inputs="text", outputs="text")
iface.launch()