|
|
import gradio as gr |
|
|
from transformers import pipeline |
|
|
|
|
|
|
|
|
translation_pipeline = pipeline( |
|
|
"translation", model="facebook/mbart-large-50-many-to-many-mmt" |
|
|
) |
|
|
|
|
|
|
|
|
generation_pipeline = pipeline( |
|
|
"text-generation", model="gpt2" |
|
|
) |
|
|
|
|
|
LANG_CODES = { |
|
|
"English": "en_XX", |
|
|
"Spanish": "es_XX", |
|
|
"Portuguese": "pt_XX", |
|
|
"Italian": "it_IT", |
|
|
"Korean": "ko_KR", |
|
|
} |
|
|
|
|
|
def translate(text, src, tgt): |
|
|
if src == tgt: |
|
|
return text |
|
|
out = translation_pipeline(text, src_lang=LANG_CODES[src], tgt_lang=LANG_CODES[tgt]) |
|
|
return out[0]["translation_text"] |
|
|
|
|
|
def generate(text, max_length=100, temperature=1.0, top_p=0.95): |
|
|
out = generation_pipeline( |
|
|
text, max_length=int(max_length), temperature=float(temperature), top_p=float(top_p) |
|
|
) |
|
|
return out[0]["generated_text"] |
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
gr.Markdown("## 🌎 Translation & Text Generation Demo") |
|
|
with gr.Tab("Translation"): |
|
|
src_lang = gr.Dropdown(list(LANG_CODES.keys()), label="Source Language", value="English") |
|
|
tgt_lang = gr.Dropdown(list(LANG_CODES.keys()), label="Target Language", value="Spanish") |
|
|
inp = gr.Textbox(label="Text to Translate") |
|
|
out = gr.Textbox(label="Translation") |
|
|
btn = gr.Button("Translate") |
|
|
btn.click(fn=translate, inputs=[inp, src_lang, tgt_lang], outputs=out) |
|
|
with gr.Tab("Text Generation"): |
|
|
prompt = gr.Textbox(label="Prompt") |
|
|
gen_out = gr.Textbox(label="Generated Text") |
|
|
max_len = gr.Slider(10, 256, value=100, label="Max Length") |
|
|
temp = gr.Slider(0.5, 1.5, value=1.0, label="Temperature") |
|
|
topp = gr.Slider(0.5, 1.0, value=0.95, label="Top-p") |
|
|
gen_btn = gr.Button("Generate") |
|
|
gen_btn.click( |
|
|
fn=generate, |
|
|
inputs=[prompt, max_len, temp, topp], |
|
|
outputs=gen_out, |
|
|
) |
|
|
|
|
|
demo.launch() |
|
|
|
