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.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):  
            pass
        with gr.Column(scale=4, min_width=1000): 
            gr.Image("logo_transparent_small.png", elem_id="logo", show_label=False, width=500)
            gr.Markdown(
                """
                <h1 style='text-align: center;'>Siswati-English Translation</h1>
                <p style='text-align: center;'>This space provides a bidirectional translation service from Siswati to English.</p>
                """
            )
        with gr.Column(scale=1):  
            pass
    
    with gr.Row():
        with gr.Column(scale=1):  
            pass
        with gr.Column(scale=4, min_width=1000): 
            inp_text = gr.Textbox(lines=5, placeholder="Enter text (maximum 5 lines)", label="Input")
            direction = gr.Radio(choices=['en->ss', 'ss->en'], label='Direction')
            translate_button = gr.Button("Translate")
            output_text = gr.Textbox(label="Output")
            translate_button.click(translate, inputs=[inp_text, direction], outputs=output_text)
        with gr.Column(scale=1):  
            pass

    with gr.Row():
        with gr.Column(scale=1):  
            pass
        with gr.Column(scale=4, min_width=1000): 
            gr.Markdown(
                """
                <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>
                """
            )
        with gr.Column(scale=1):  
            pass
    

    with gr.Accordion("More Information", open=False):
        gr.Markdown("""
        <h4 style="text-align: center;">Authors</h4>
        <div style='text-align: center;'>Vukosi Marivate, Richard Lastrucci</div>
        """)

demo.launch()