import os
import gradio as gr
import tensorflow as tf
import tensorflow_text as tf_text
from huggingface_hub import Repository

repo = Repository(
	local_dir="nmt-bahdanau-attention",
	clone_from="pyimagesearch/nmt-bahdanau-attention",
	use_auth_token=os.environ.get("token")
)
reloaded = tf.saved_model.load("nmt-bahdanau-attention/translator")

title="Neural Machine Translation with Bahdanau's Attention"
description="The model used here is a POC and not SOTA on NMT."

examples=["how are you?", "good morning.", "how is your health?"]

def get_translation(sentence):
	result = reloaded.translate(
		sourceText=tf.constant([sentence])
	)["text"].numpy()[0].decode()
	return result

nmt_space = gr.Interface(
	fn=get_translation,
	inputs=gr.Textbox(label="English Sentence"),
	outputs=gr.Textbox(label="French Sentence"),
	title=title,
	description=description,
	examples=examples,
)

nmt_space.launch()