import gradio as gr
from translate import translator_fn


def predict(text):
    result = translator_fn(text)
    return {
        "input_text": result.input_text,
        "input_tokens": result.input_tokens,
        "n_input": result.n_input,
        "output_text": result.output_text,
        "output_tokens": result.output_tokens,
        "n_output": result.n_output,
        "output_scores": result.output_scores,
        "cross_attention": result.cross_attention.tolist(),
    }


gradio_app = gr.Interface(
    predict,
    inputs=gr.Text(placeholder="Enter a sentence to translate...", label="Input text"),
    outputs=[gr.Json(label="Model output")],
    title="En2Ru Scientific Translator",
    description="Translate scientific texts from English to Russian",
)

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