import gradio as gr
from dicta_onnx import Dicta

dicta = Dicta('./dicta-1.0.int8.onnx')

theme = gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto")])

def add_diacritics(text: str, mark_niqqud_male: bool) -> str:
    mark_matres_lectionis = "|" if mark_niqqud_male else None
    return dicta.add_diacritics(text, mark_matres_lectionis=mark_matres_lectionis)

with gr.Blocks(theme=theme) as demo:
    gr.Markdown("# Hebrew Text Diacritization")
    gr.Markdown("Enter Hebrew text without diacritics, and the model will add them.")

    text_input = gr.Textbox(
        placeholder="Type text without diacritics...",
        label="Input Text",
        lines=5,
        rtl=True,
        elem_classes=["input"],
        value="מה קורה? יאלה תתחילו לשחק עם הנקדן!",
    )

    mark_niqqud_male = gr.Checkbox(label="Mark Niqqud Male (Matres Lectionis)")

    diacritized_output = gr.Textbox(
        label="Text with Diacritics",
        lines=7,
        rtl=True,
        elem_classes=["output"]
    )
    
    submit_button = gr.Button("Add Diacritics")

    submit_button.click(
        fn=add_diacritics,
        inputs=[text_input, mark_niqqud_male],
        outputs=[diacritized_output]
    )

if __name__ == "__main__":
    demo.launch()