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

with gr.Row():
    gr.Column()
    gr.Column(gr.Image("translation-1/logo_transparent_small.png", alt="DSFSI Logo", elem_id="logo", label=None))
    gr.Column()

description = """
<p style='text-align: center;'>
    Siswati to English Translation
</p>
<p>
    This space provides a bidirectional translation service from Siswati to English.
</p>
"""

article = """
<div style='text-align: center;'>
    <a href='https://github.com/dsfsi/en-ss-m2m100-combo' target='_blank'>En-Ss GitHub</a> |
    <a href='https://github.com/dsfsi/ss-en-m2m100-combo' target='_blank'>Ss-En GitHub</a> |
    <a href='https://docs.google.com/forms/d/e/1FAIpQLSf7S36dyAUPx2egmXbFpnTBuzoRulhL5Elu-N1eoMhaO7v10w/viewform' target='_blank'>Feedback Form</a> 
</div>
<br/>
<p style='text-align: center;'>
    <h2>Translate | Siswati to English</h2>
</p>
"""

authors = """
<div style='text-align: center;'>
    Authors: Vukosi Marivate, Richard Lastrucci
</div>
"""

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.Textbox(lines=5, placeholder="Enter text (maximum 5 lines)", label="Input"),
        gr.Radio(
            choices=['en->ss', 'ss->en'],
            default='en->ss',
            label='Direction'),
    ],
    outputs="text"
)

iface.launch(enable_queue=True)

gr.markdown(authors, unsafe_allow_html=True)
gr.markdown(citation, unsafe_allow_html=True)
gr.markdown(doi, unsafe_allow_html=True)