File size: 1,931 Bytes
b85212b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | import gradio as gr
from transformers import pipeline
# Translation pipeline (mBART)
translation_pipeline = pipeline(
"translation", model="facebook/mbart-large-50-many-to-many-mmt"
)
# Text generation pipeline (GPT-2)
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()
|