import gradio as gr

from typo_check import css, process_input,MODEL_OPTIONS_TYPO
from title_gen import generate_title, MODEL_OPTIONS_TITLE
from content_gen import generate_content, MODEL_OPTIONS_CONTENT, get_default_prompt
from instruct_dv import generate_response, MODEL_OPTIONS_INSTRUCT
from instruct_dv_tuned import generate_response_tuned, MODEL_OPTIONS_INSTRUCT_TUNED
from en_dv_latin import translate, MODEL_OPTIONS_TRANSLATE

def update_textbox_direction(direction):
    # Enable RTL only if the source language is Dhivehi (dv2*)
    is_rtl = direction.startswith("dv2")
    return gr.Textbox(rtl=is_rtl)

# Create Gradio interface using the latest syntax
with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
    with gr.Tabs():
        with gr.Tab("Typo Correction"):
            gr.Markdown("# <center>Dhivehi Typo Correction</center>")
            gr.Markdown("This app uses a fine-tuned T5 model to correct typos in Dhivehi text. Enter text with typos and the model will attempt to fix them.")
            
            with gr.Row():
                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS_TYPO.keys()), value="A3 Model", label="Model")

            with gr.Row():
                input_text = gr.Textbox(
                    lines=1,
                    placeholder="ދިވެހި ބަހުން ލިޔުމެއް މިތާ ލިޔެބަލާ",
                    label="Input Text",
                    rtl=True,
                    elem_classes="textbox1"
                )


            with gr.Row():
                corrected_text = gr.Textbox(
                    lines=1,
                    label="Corrected Text",
                    rtl=True,
                    elem_classes="textbox1"
                )
            
            with gr.Row():
                highlighted_diff = gr.HTML(
                    label="Changes Highlighted",
                    elem_classes="dhivehi-text"
                )
            
            submit_btn = gr.Button("ރަނގަޅު ކޮށްލުމަށް",elem_classes="textbox1")  # "Correct" in Dhivehi
            submit_btn.click(
                fn=process_input,
                inputs=[input_text,model_choice],
                outputs=[corrected_text, highlighted_diff]
            )
            
            gr.Examples(
                examples=[
                    ["މައި ނޭމް އިސް އަހްމދް"],  # My name is Ahmed (with typos)
                    ["ވަރަށ ރައްގޅޭ ދއުވސް"],   # Very good day (with typos)
                    ["މިއަދ ސްކޫލަށް ދިއުމަށްޓަކައި ހެނދުނ ވަރށް އަވަހށް ތެދުވިން"],  # I woke up early today to go to school (with typos)
                    ["ރާއްޖެއަކީ ވަރށް ރީތި ގުދުރަތީ މަންޒަރުތައް ހުރި ގައުމެކެވެ"],  # Maldives is a country with beautiful natural scenery (with typos)
                    ["ދިވެހި ބަހަކީ އަޅުގަނޑުމެންގެ މާދަރި ބަހެވ އަދި އެބަސް ދިރުވައި އާލާކުރން ޖެހެއވެ"],  # Dhivehi is our mother tongue and we need to preserve it (with typos)
                    ["ކުޅިވަރަކީ ހަށިގަނޑގެ ދުޅަހެޔޮ ކަމަށް ވަރށް މުހިންމު ކަމެކެވެ"],  # Sports are very important for physical health (with typos)
                ],
                inputs=input_text,
            )

            gr.Markdown("""## Note:

        This application uses a fine-tuned T5 model specifically designed to correct common typos in Dhivehi text. The model focuses on:

        - Fixing common mistakes in diacritics
        - Correcting missing Dhivehi letters
        - Addressing typical typing errors

        ### What it can fix

        The model works best with common typing mistakes. For example, if you type "ދިވެހ ބަސް" (missing a diacritic), the model will correct it properly.

        ### Limitations

        The model is trained only on common errors and may not fix more complex or unusual mistakes. For instance, if you type "ދިވެހިދަ ބަސް" (with an incorrect vowel mark rather than a missing one), the correction might not work as expected.

        Simply enter your Dhivehi text with typos, and the model will attempt to correct the common errors while preserving the meaning of your text.""")

        with gr.Tab("Title Generation"):
            gr.Markdown("# <center>Dhivehi Article Title Generator</center>")
            gr.Markdown("Enter Dhivehi article content to generate a title using a fine-tuned T5 model. Set the random seed for reproducibility. Enable sampling for creative/random results.")
            
            with gr.Row():
                article_content = gr.Textbox(lines=10, label="Article Content", rtl=True, elem_classes="textbox1")
            with gr.Row():
                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS_TITLE.keys()), value="V6 Model", label="Model")
            with gr.Row():
                seed = gr.Slider(0, 10000, value=42, step=1, label="Random Seed")
                use_sampling = gr.Checkbox(label="Use Sampling (Creative/Random)", value=False)
            with gr.Row():
                generated_title = gr.Textbox(label="Generated Title", rtl=True, elem_classes="textbox1" )
            generate_btn = gr.Button("Generate Title")
            generate_btn.click(
                fn=generate_title,
                inputs=[article_content, seed, use_sampling, model_choice],
                outputs=generated_title
            )

            gr.Examples(
                examples=[
                    ["""މިއަދު ފާހަގަކުރާ "އާތު ޑޭ" ނުވަތަ ދުނިޔޭގެ ދުވަހުގެ މުނާސަބަތުގައި ގޫގުލްގެ ޑޫޑުލްގައި ރާއްޖޭގެ ރީތި ރަށެއް ކަމަށްވާ ށ. އެރިޔަދޫ ހިމަނައިފި އެވެ.

ދުނިޔޭގެ ދުވަހާ ގުޅުވައިގެން، ގޫގުލްގެ ހޯމް ޕޭޖުން ފެންނަ ޑޫޑުލް މިއަދު ހަދާފައިވަނީ ދުނިޔޭގެ އަޖައިބެއް ކަހަލަ ރީތި ހަ ތަނެއްގެ މަންޒަރުތައް ބޭނުންކޮށްގެންނެވެ. ޑޫޑުލްގައި ހިމަނާފައި ވަނީ އިނގިރޭސިބަހުން ގޫގުލް ކިޔާއިރު އޮންނަ ހަ އަކުރުގެ ބައްޓަން ސިފަވާ ގޮތަށް އޮންނަ ތަންތަނެވެ.

ޑޫޑުލްގައި، ހިތްގައިމު އެރިޔަދޫ އޮތީ ގޫގުލްގެ "ޖީ" އަކުރުގެ ގޮތުގަ އެވެ.

އެރިޔަދޫއަކީ ހުއިފިލަނޑާ އާލާކުރުމަށް ސަރުކާރުން ވަރަށް ކުރީއްސުރެ އަމިއްލަ މީހަކަށް ކުއްޔަށް ދޫކޮށްފައި އޮންނަ ރަށެކެވެ. އެ ރަށުގެ ހާއްސަ، ރީތި ނޫފެހި ވިލު ބޭނުންކުރަމުން އަންނަނީ ހުއިފިލަނޑާ އަލާކުރުމަށެވެ.


ޑޫޑުލްގައި އެ ރަށް ހިމެނުމަށް ފަހު ގޫގުލް އިން ބުނީ މިފަހަރު "ޖީ" އަކުރު ހާއްސަކުރ ރާއްޖެއަށް ކަމަށެވެ.

"ބިން ތިރިކޮށް ހެދިފައި އޮންނަ ސުވަރުގެއް ފަދަ ރާއްޖެ. މި ރީތި ޖަޒީރާ ގައުމު އެކުލެވިގެންވަނީ ވައް ބައްޓަމަށް އޮންނަ، މުރަކަ ފަރުތަކެއްގެ މަތީގައި އުފެދިފައިވާ ވިހި އަތޮޅުގެ މައްޗަށް. މި އަތޮޅުތަކުގައި ރީތި ފަޅުތައް ހުރޭ. އަދި 'އެޓޯލް' މި ބަސް އައިސްފައި ވަނީ ވެސް ދިވެހިބަހުން،" ގޫގުލް އިން ބުންޏެވެ.


އެރިޔަދޫ

ޑޫޑުލްގައި "އޯ" އަކުރުގައި ހިމަނާފައި ވަނީ ފްރާންސްގެ އެލްޕްސް ފަރުބަދަތަކެވެ. އެއީ ހިލަ ފަރުބަދަތަކާއި ރީތި ފެހި ގަސްތަކުން ޖަރީވެފައި އޮންނަ ރީތި ސަރަހައްދެކެވެ.

ދެވަނަ "އޯ" އަކުރުގައި ހިމަނާފައި ވަނީ ރީތި ކޯރަކާ އެކު އޮންނަ، ގަސްގަހާގެއްސާއި ދިރުންތަކުން މުއްސަނދި ސަރަހައްދެއް ކަމަށްވާ ކެނެޑާގެ ކޯޓް-ނޯޑެވެ.

ދެވަނަ "ޖީ" އަކުރުގައި އާޖެންޓީނާ ދެކުނުގެ ސަރަހައްދެއް ކަމަށްވާ މެންޑޯޒާގެ މަންޒަރެއް ހިމަނާފައިވާއިރު، "އެލް" އަކުރުގައި ހިމަނާފައި ވަނީ އެމެރިކާގެ ޔޫޓާގެ ހަނަފަސް ސަރަހައްދެކެވެ. އަދި "އީ" އަކުރު ހާއްސަކޮށްފައި ވަނީ އޮސްޓްރޭލިއާގެ ނިއު ސައުތު ވޭލްސްގައި އޮންނަ ހިކިފަސް ސަރަހައްދަކަށެވެ.
"""],
        ["""ބިލްބާއޯ (މެއި 1) - ޔޫއެފާ ޔޫރޮޕާ ލީގު ކާމިޔާބު ކުރިޔަސް މެންޗެސްޓާ ޔުނައިޓެޑްގެ މައްސަލަތަކަށް ހައްލެއް ނުލިބޭނެ ކަމަށް ކޯޗު ރުބެން އަމޯރިމް ބުނެފި އެވެ.

ރަށުން ބޭރުގައި މިރޭ އެތުލެޓިކް ބިލްބާއޯއާ ދެކޮޅަށް ކުޅޭ ސެމީ ފައިނަލްގެ ފުރަތަމަ ލެގުގެ ކުރިމް މީޑިއާއާ ވާހަކަ ދައްކަމުން އަމޯރިމް ބުނީ، މުބާރާތް ކާމިޔާބު ކޮށްގެން ޗެމްޕިއަންސް ލީގު ޖާގައެއް ލިބި، ހަރަދު ކުރާނެ ފައިސާ ލިބުނަސް، ކްލަބުގެ ކޮންސިސްޓެންސީ ދަމަހައްޓަން ގިނަ މަސައްކަތްތަކެއް ކުރަން ޖެހިފައިވާ ކަމަށެވެ.

"ޔޫރޮޕާ ލީގު ކާމިޔާބު ކޮށްގެން އަންނަ ސީޒަންގައި ޗެމްޕިއަންސް ލީގު ކުޅެން ލިބުމަކީ، އަހަރެމެންގެ ކްލަބްގެ ބައިވަރު ކަންކަން ބަދަލު ކުރެވިދާނެ ކަމެއް، ކުރިއަށް އޮތް ހޫނު މޫސުމުގެ ޓްރާންސްފާގައި ވެސް،" ޔުނައިޓެޑްގެ ކޯޗު އަމޯރިމް ބުންޏެވެ.

"ޔޫރޮޕާ ލީގުން މޮޅުވެގެން އަހަރެމެންގެ މައްސަލަތަކެއް ހައްލެއް ނުވާނެ. މޮޅުވެގެން ލިބޭނީ ޗެމްޕިއަންސް ލީގުގެ ޖާގައަކާއި ހަރަދު ކުރާނެ ފައިސާ. އެކަމަކު އަސްލު މައްސަލަތައް އަދިވެސް ހުރީ. އަހަރެމެންގެ ކޮންސިސްޓެންޓް ވާން ޖެހޭ. ރަނގަޅު ނިންމުންތައް ނިންމައި، ރަނގަޅު ރެކްރޫޓްމަންޓް ސިސްޓަމަކާއި ރަނގަޅު އެކަޑަމީއެއް ބޭނުންވޭ. މިއ މި ކުލަބްގެ އަސްލު މަގާމަށް ދިޔުމަށްޓަކައި ހައްލު ކުރަން ޖެހޭ ކަންކަން."

އިނގިރޭސި ޕްރިމިއާ ލީގުގެ 14 ވަނައިގައި އޮތް މެންޗެސްޓާ ޔުނައިޓެޑަށް، އަންނަ ސީޒަންގައި ޔޫރަޕްގެ މުބާރާތެއް ކުޅެވެން އޮތް ހަމައެކަނި ގޮތަކ ޔޫރޮޕާ ލީގު ކާމިޔާބު ކުރުމެވެ.

ޔޫރޮޕާ ލީގުގެ ސެމީ ފައިނަލްގައި އެތުލެޓިކް ބިލްބާއޯ އާއި ޔުނައިޓެޑް މިރޭ ވާދަ ކުރާއިރު، ދެ ވަނަ ސެމީގައި ބައްދަލުކުރާނ ޓޮޓެންހަމް ހޮޓްސްޕާ އާއި ނޯވޭގެ ބޯޑޯ/ގްލިމްޓް އެވެ.
"""],
                ],
                inputs=article_content,
            )

            gr.Markdown("""## Notes:

✨ **What does this do?**

Dhivehi article content and, with the help of a fine-tuned T5 model, conjures up a fitting title! Whether you want something straightforward or a bit more creative, just tweak the random seed or enable sampling for a dash of randomness.

- **Best for:** Quickly summarizing the essence of your article in a catchy title.
- **Creative mode:** Turn on sampling for surprising, out-of-the-box suggestions.
- **Consistent mode:** Keep sampling off for reliable, focused titles.

⚠️ **Experimental Status**
This is an experimental model trained on a very small dataset of Dhivehi news articles. Results may vary significantly.

### Limitations

- The model is only good with clear, well-written content. If your article is too short, vague, or off-topic, the title might be less impressive.
- Sometimes, creativity can go a little too wild—if you get a strange title, try a different seed or turn off sampling.
- Due to limited training data, the model may not handle all topics or writing styles equally well.
- Experimental nature means outputs should be reviewed carefully before use.

""")

        with gr.Tab("Content Generation"):
            gr.Markdown("# <center>Dhivehi Content Generator</center>")
            gr.Markdown("Generate Dhivehi news articles or content using a fine-tuned T5 model. Select your model and tweak generation parameters.")
            with gr.Row():
                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS_CONTENT.keys()), value=list(MODEL_OPTIONS_CONTENT.keys())[0], label="Model")
            with gr.Row():
                prompt_input = gr.Textbox(label="Input Prompt", lines=1, elem_classes="textbox1")
            def update_prompt(model_choice):
                return get_default_prompt(model_choice)
            model_choice.change(fn=update_prompt, inputs=model_choice, outputs=prompt_input)
            gr.Examples(
                examples=[
                    ["ޗޮކްލެޓު ކޭކުގެ ރެސިޕީ އަކީ؟"],
                    ["ދިވެހި ވާހަކަ އެއް ލިޔެލަބަލަ"],
                    ["ގުޅިފަޅުގެ ލޭންޑް ޔޫޒް ޕްލޭނަށް ބަދަލު ގެނެސްފައިވާތީ މިއަދު ގުރުއަތުލުން ބާއްވަނީ"],
                    ["ދެމަފިރިންގެ ގުޅުންތަކާއި ޒިންމާތައް ބިނާކުރަން ޖެހޭ އަސާސްތައް"],
                    ["ކެފެއިންގެ ސަބަބުން އިންސާނާގެ ހަށިގަނޑަށް ކުރާ ޖިސްމާނީ އަދި ނަފްސާނީ އަސަރުތައް ބަޔާންކޮށްދެއްވާށެވެ."],
                ],
                inputs=prompt_input,
            )
            with gr.Row():
                max_tokens_slider = gr.Slider(10, 300, value=150, label="Max New Tokens")
                num_beams_slider = gr.Slider(1, 10, value=4, step=1, label="Beam Size (num_beams)")
            with gr.Row():
                rep_penalty_slider = gr.Slider(1.0, 1.9, value=1.2, step=0.1, label="Repetition Penalty")
                ngram_slider = gr.Slider(0, 10, value=3, step=1, label="No Repeat Ngram Size")
            with gr.Row():
                do_sample_checkbox = gr.Checkbox(label="Do Sample", value=True)
            with gr.Row():
                output_text = gr.Textbox(label="Generated Output", lines=10, elem_classes="textbox1")
            generate_button = gr.Button("Generate Content")
            generate_button.click(
                fn=generate_content,
                inputs=[prompt_input, max_tokens_slider, num_beams_slider, rep_penalty_slider, ngram_slider, do_sample_checkbox, model_choice],
                outputs=output_text
            )
            gr.Markdown("""
**Notice:**

All outputs generated are synthetic, created using fine-tuned models for experimental and educational evaluation. Accuracy is not guaranteed, and the content should not be considered a source of truth. Please avoid applying these results to production environments, critical systems, or real-world decision-making without proper validation.
""")

        with gr.Tab("Instruction Following"):
            gr.Markdown("# <center>Dhivehi Instruction Following</center>")
            gr.Markdown("Enter an instruction and (optionally) input text. The model will generate a response following your instruction. Set the random seed for reproducibility. Enable sampling for creative/random results.")
            with gr.Row():
                instruction = gr.Textbox(lines=2, label="Instruction", rtl=True, elem_classes="textbox1")
            with gr.Row():
                input_text = gr.Textbox(lines=2, label="Input Text (optional)", rtl=True, elem_classes="textbox1")
            with gr.Row():
                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS_INSTRUCT.keys()), value=list(MODEL_OPTIONS_INSTRUCT.keys())[0], label="Model")
            with gr.Row():
                max_tokens_slider = gr.Slider(10, 512, value=256, label="Max New Tokens")
                seed = gr.Slider(0, 10000, value=42, step=1, label="Random Seed")
                use_sampling = gr.Checkbox(label="Use Sampling (Creative/Random)", value=False)
            with gr.Row():
                generated_response = gr.Textbox(label="Model Response", rtl=True, elem_classes="textbox1")
            generate_btn = gr.Button("Generate Response")
            generate_btn.click(
                fn=generate_response,
                inputs=[instruction, input_text, seed, use_sampling, model_choice,max_tokens_slider,num_beams_slider],
                outputs=generated_response
            )
            gr.Examples(
                examples=[
                    ["ދީފައިވާ މައުޟޫޢާ ބެހޭގޮތުން ކުރު ޕެރެގްރާފެއް ލިޔެލާށެވެ.","އިއާދަކުރަނިވި ހަކަތަ ބޭނުންކުރުމުގެ މުހިންމުކަން"],
                    ["އާ މޯބައިލް އެޕް ޕްރޮމޯޓް ކުރުމަށް މާކެޓިންގ ސްޓްރެޓެޖީތަކުގެ ލިސްޓެއް އުފެއްދުން.",""],
                    ["ދިގުމިނުގައި 10ސެންޓިމީޓަރު އަދި ފުޅާމިނަކީ 5ސެންޓިމީޓަރު ހުންނަ ރެކްޓަންގްލަރެއްގެ ސަރަހައްދު ހިސާބުކުރުން.",""],
                    ["ތިރީގައިވާ ބަސްފުޅު ތެދެއް ނުވަތަ ދޮގުގެ ގޮތުގައި ގިންތިކުރުން.","ސުޕްރީމް ކޯޓަކީ އެމެރިކާގެ އެންމެ މަތީ ކޯޓެވެ."],
                ],
                inputs=[instruction, input_text],
            )
            gr.Markdown("""\
**Notes:**
- This tab allows you to give instructions to the model, optionally with input text, for general-purpose generation or task following in Dhivehi.
- Try different seeds or enable sampling for more creative outputs.
- The model is experimental and may not always follow instructions perfectly.
""")

        with gr.Tab("Instruction Tuned"):
            gr.Markdown("# <center>Dhivehi Instruction-Tuned Model</center>")
            gr.Markdown("Generate answers by providing a custom instruction and optional input. This instruction-tuned model is designed for better format awareness, task generalization, and stronger alignment with user intent.")
            with gr.Row():
                instruction = gr.Textbox(lines=2, label="Instruction", rtl=True, elem_classes="textbox1")
            with gr.Row():
                input_text = gr.Textbox(lines=2, label="Input Text (optional)", rtl=True, elem_classes="textbox1")
            with gr.Row():
                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS_INSTRUCT_TUNED.keys()), value=list(MODEL_OPTIONS_INSTRUCT_TUNED.keys())[0], label="Model")
            with gr.Row():
                max_tokens = gr.Slider(64, 1024, value=768, step=16, label="Max New Tokens")
                temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature")
                num_beams = gr.Slider(1, 8, value=4, step=1, label="Number of Beams")
            with gr.Row():
                generated_response = gr.Textbox(label="Model Response", rtl=True, elem_classes="textbox1")
            generate_btn = gr.Button("Generate Response")
            generate_btn.click(
                fn=generate_response_tuned,
                inputs=[instruction, input_text, seed, model_choice,max_tokens,temperature,num_beams],
                outputs=generated_response
            )
            gr.Examples(
                examples=[
                    ["ދީފައިވާ މައުޟޫޢާ ބެހޭގޮތުން ކުރު ޕެރެގްރާފެއް ލިޔެލާށެވެ.","އިއާދަކުރަނިވި ހަކަތަ ބޭނުންކުރުމުގެ މުހިންމުކަން"],
                    ["ދިގުމިނުގެ މިންވަރުތައް ބަދަލުކުރުން.","1 ކިލޯމީޓަރ"],
                    ["ދެ މޯބައިލް ފޯނެއްގެ ފީޗާސް އަޅާކިޔާށެވެ.","އައިފޯން 11 ޕްރޯ އަދި ސެމްސަންގް ގެލެކްސީ އެސް20 އަލްޓްރާ"],
                    ["މަސައްކަތްތައް ފައިދާހުރި ގޮތެއްގައި ހަވާލުކުރުމަށް އަޅަންޖެހޭ ފިޔަވަޅުތައް ބަޔާންކުރުން.",""],
                    ["އާ މޯބައިލް އެޕް ޕްރޮމޯޓް ކުރުމަށް މާކެޓިންގ ސްޓްރެޓެޖީތަކުގެ ލިސްޓެއް އުފެއްދުން.",""],
                    ["ދިގުމިނުގައި 10ސެންޓިމީޓަރު އަދި ފުޅާމިނަކީ 5ސެންޓިމީޓަރު ހުންނަ ރެކްޓަންގްލަރެއްގެ ސަރަހައްދު ހިސާބުކުރުން.",""],
                    ["ތިރީގައިވާ ބަސްފުޅު ތެދެއް ނުވަތަ ދޮގުގެ ގޮތުގައި ގިންތިކުރުން.","ސުޕްރީމް ކޯޓަކީ އެމެރިކާގެ އެންމެ މަތީ ކޯޓެވެ."],
                ],
                inputs=[instruction, input_text],
            )
            gr.Markdown("""
### 📝 Notes:
- This model is **instruction-tuned** using Dhivehi data, designed to follow a wide variety of instructions.
- Provide both an instruction and input for best results, but input is optional for tasks like open-ended generation.
- Use **temperature** > 0 for more diverse outputs, or set to **0** for deterministic answers.
- **Number of beams** increases quality at the cost of speed.
- Assuming the model will have **better format awareness**, can handle **multiple task types**, and often **align more closely with your intent**. (Have to test)
- This model is experimental and may not always follow instructions perfectly. 
""")

        with gr.Tab("Translation Tasks"):
            gr.Markdown("# <center>Dhivehi Translation</center>")
            gr.Markdown("Select a translation direction and enter text to translate between Dhivehi, English and Latin script.")
            # Set up initial choices and examples based on default model
            default_model = list(MODEL_OPTIONS_TRANSLATE.keys())[0]
            if default_model.startswith("MT"):
                initial_choices = ["2dv", "2en"]
                initial_value = "2dv"
            else:
                initial_choices = ["en2dv:", "dv2en:", "dv2latin:", "latin2dv:"]
                initial_value = "en2dv:"

            with gr.Row():
                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS_TRANSLATE.keys()), value=default_model, label="Model")            
            with gr.Row():
                instruction = gr.Dropdown(
                    choices=initial_choices,
                    label="Translation Direction",
                    value=initial_value
                )
            with gr.Row():
                input_text = gr.Textbox(lines=2, label="Text to Translate", rtl=True, elem_classes="textbox1")
            with gr.Row():
                generated_response = gr.Textbox(label="Translated Text", rtl=True, elem_classes="textbox1")
            with gr.Row():
                max_tokens_slider = gr.Slider(10, 128, value=128, label="Max New Tokens")
                num_beams_slider = gr.Slider(1, 10, value=4, step=1, label="Beam Size (num_beams)")
            with gr.Row():
                rep_penalty_slider = gr.Slider(1.0, 1.9, value=1.2, step=0.1, label="Repetition Penalty")
                ngram_slider = gr.Slider(0, 10, value=3, step=1, label="No Repeat Ngram Size")
            generate_btn = gr.Button("Translate")

            generate_btn.click(
                fn=translate,
                inputs=[instruction, input_text, model_choice,max_tokens_slider, num_beams_slider, rep_penalty_slider, ngram_slider],
                outputs=generated_response
            )

            with gr.Row():
                gr.Examples(
                    examples=[
                        "ދުނިޔޭގެ އެކި ކަންކޮޅުތަކުން މިލިއަން މީހުން މައްކާއަށް ޖަމާވެފައި",
                        "Concerns over prepayment of GST raised in parliament",
                        "ވައިބާރުވުމުން ކުޅުދުއްފުށީ އެއާޕޯޓަށް ނުޖެއްސިގެން މޯލްޑިވިއަންގެ ބޯޓެއް އެނބުރި މާލެއަށް",
                        "Paakisthaanuge skoolu bahakah dhin hamalaaehgai thin kuhjakaai bodu dhe meehaku maruvehje"
                    ],
                    inputs=input_text,
                )

            def update_choices(model_name):
                if model_name.startswith("MT"):
                    return gr.update(choices=["2dv", "2en"], value="2dv")
                else:
                    return gr.update(choices=["en2dv:", "dv2en:", "dv2latin:", "latin2dv:"], value="en2dv:")

            model_choice.change(
                fn=update_choices,
                inputs=[model_choice],
                outputs=[instruction]
            )

            gr.Markdown("""\
**Notes:**
- Supports translation between Dhivehi, English and Latin script
- Model trained on news articles and common phrases
- Translation quality may vary based on the domain of the text
""")

# Launch the app
if __name__ == "__main__":
    #demo.launch(server_name="0.0.0.0", server_port=7811)
    demo.launch()