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(description="Model output", label="Model output")],
    title="En2Ru Scientific Translator",
    description="Translate scientific texts from English to Russian",
    examples=[
        [
            r"There is no closed form to implement the KL divergence by the definition of (REF ) and (REF ) for "
            r"Gaussian Mixture Models. Instead, we resort to the Monte Carlo simulation method proposed in [1]}. "
            r"Then, the KL divergence can be caculated by: \(D_{KL_{MC}}(p||q) =\frac{1}{n} \sum _{i=1}^{n} log("
            r"\frac{p(x_i)}{q(x_i)})\) \(D_{KL_{MC}}(q||p) =\frac{1}{n} \sum _{i=1}^{n} log(\frac{q(y_i)}{p(y_i)})\)"],
        [
            r"Almost all currently used classifiers are not intrinsically well-calibrated [1]}, which means their "
            r"output scores can't be interpreted as probabilities. This is an issue when the model is used for "
            r"decision making, as a component in a more general probabilistic pipeline, or simply when one needs a "
            r"quantification of the uncertainty in model's predictions, for example in high risk applications."],
        [
            r"First, with the development of the high-torque electric actuators, such as [1]}, [2]} the robots are "
            r"becoming more dynamical. These actuators allow them not only to move at high speeds, but also to "
            r"rapidly create forces and torques to perform dynamic actions, such as running, jumping, etc."],
    ],
)

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