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()