import gradio as gr
from transformers import pipeline

# prefix = "<2id> "
# madlad = pipeline("translation", model="google/madlad400-3b-mt")
lulu = pipeline("translation", model="tirtohadi/lulu-v1")


def translate(text):
    # Split input text into paragraphs
    paragraphs = text.split("\n\n")  # Assuming paragraphs are separated by two newline characters
    
    # Translate each paragraph
    translated_paragraphs_lulu = []
    #translated_paragraphs_madlad = []
    for paragraph in paragraphs:
        # Call your custom model here to translate each paragraph
        # translated_paragraph_madlad = madlad(prefix + paragraph, max_length=400)[0]["translation_text"]   
        # translated_paragraphs_madlad.append(translated_paragraph_madlad)
    
        translated_paragraph_lulu = lulu(paragraph, max_length=400)[0]["translation_text"]   
        translated_paragraphs_lulu.append(translated_paragraph_lulu)

    
    # Join translated paragraphs back into text
    translated_text_lulu = "\n\n".join(translated_paragraphs_lulu)
    # translated_text_madlad = "\n\n".join(translated_paragraphs_madlad)
    
    return translated_text_lulu #,translated_text_madlad
    
with gr.Blocks() as demo:
    gr.HTML("<h2>Lulu Translate</h2>")
    gr.Markdown("Lulu is a Christian domain specific machine translation")
    with gr.Row():        
        input_text1 = gr.Textbox(label="English Text",lines=5)
        output_lulu = gr.Textbox(label="Indonesian Translation",lines=5)
        
    with gr.Row():
        with gr.Column(scale=1):        
            btn = gr.Button("Translate")
            btn.click(fn=translate, inputs=input_text1, outputs=[output_lulu], api_name="translate")

demo.launch()