ZVN00's picture
Update app.py
c1c3ddc
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
tokenizervin = AutoTokenizer.from_pretrained("vinai/vinai-translate-en2vi-v2", src_lang="en_XX")
modelvin = AutoModelForSeq2SeqLM.from_pretrained("vinai/vinai-translate-en2vi-v2")
def translate(en_text: str) -> str:
input_ids = tokenizervin(en_text, return_tensors="pt").input_ids
output_ids = modelvin.generate(
input_ids,
decoder_start_token_id=tokenizervin.lang_code_to_id["vi_VN"],
num_return_sequences=1,
max_length=4096,
num_beams=5,
early_stopping=True
)
vi_text = tokenizervin.batch_decode(output_ids, skip_special_tokens=True)
vi_text = '\n'.join(vi_text)
return vi_text
tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-1.3B", src_lang="eng_Latn")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-1.3B")
def translation(text):
translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang="eng_Latn", tgt_lang="vie_Latn")
output = translator(text, max_length=4096)
output = output[0]['translation_text']
return output
description_string = """
<div align="center">
<img src="https://i.spread.name/83a58f6a-475f-4b57-9c0f-b937a7919126_BIMBIP.png" width=200px>
<span style="font-size: 24px; font-weight: bold;">BIPSOUP DRIVE-THRU</span><br>
<span style="font-size: 18px; font-weight: bold;">You are almost ready to be entrusted with the secret ingredient of my Biptopia Soup!</span>
</div>
"""
iface = gr.Interface(
fn=translate,
description=description_string,
inputs=[gr.Textbox(lines=7, max_lines=100, placeholder="From BimBip with love", label="Ingredient"),],
outputs=[gr.Textbox(lines=7, max_lines=100, placeholder="Nhanh, Gọn, Mướt", label="Fastsoup", show_label=True, show_copy_button=True),])
iface2 = gr.Interface(
fn=translation,
description=description_string,
inputs=[gr.Textbox(lines=7, max_lines=100, placeholder="From BimBip with love", label="Ingredient"),],
outputs=[gr.Textbox(lines=7, max_lines=100, placeholder="Nhanh, Gọn, Mướt", label="Fastsoup", show_label=True, show_copy_button=True),])
demo = gr.TabbedInterface([iface, iface2], ["VinAI", "NLLB"], theme=gr.themes.Soft())
if __name__ == "__main__":
demo.launch()