import gradio as gr
from transformers import pipeline

translater_en_ss = pipeline("translation", model="dsfsi/en-ss-m2m100-combo", src_lang="en", tgt_lang="ss")
translater_ss_en = pipeline("translation", model="dsfsi/ss-en-m2m100-combo", src_lang="ss", tgt_lang="en")

def translate(inp, direction):
    if direction == 'en->ss':
        res = translater_en_ss(inp, max_length=512, early_stopping=True)[0]['translation_text']
    else:
        res = translater_ss_en(inp, max_length=512, early_stopping=True)[0]['translation_text']
    return res

description = """
<p>
<center>
Multi-domain Translation Between Siswati and English
</center>
</p>
"""
article = "<p style='text-align: center'><a href='https://huggingface.co/dsfsi/en-ss-m2m100-combo' target='_blank'>by dsfsi</a></p></center></p>"

examples = [
    ["Thank you for your help", "en->ss"],
    ["Ngiyabonga ngesiciniseko sakho", "ss->en"]
]

iface = gr.Interface(
    fn=translate,
    title="Siswati-English Translation",
    description=description,
    article=article,
    examples=examples,
    inputs=[
        gr.components.Textbox(lines=5, placeholder="Enter text (maximum 5 lines)", label="Input"),
        gr.components.Radio(
            choices=['en->ss', 'ss->en'],
            default='en->ss',
            label='Direction'),
    ],
    outputs="text"
)

iface.launch(enable_queue=True)