gorkemozkaya's picture
reversed translation direction
54177cb
import gradio as gr
import sys
sys.path = ['/code/nmt-en-tr/datasets', '/code/nmt-en-tr/models'] + sys.path
import tensorflow_text as tftxt
import tensorflow as tf
from official.core import exp_factory, task_factory
from official.nlp.configs import wmt_transformer_experiments as wmt_te
tokenizer= tftxt.SentencepieceTokenizer(
model=tf.io.gfile.GFile("/code/pretrained_v2/sentencepiece_en_tr.model", "rb").read(),
add_eos=True)
# sys.path = ['/root/.local/lib/python3.10/site-packages', '/root/.local/bin'] + sys.path
task_config = exp_factory.get_exp_config('transformer_tr_en_blended/base').task
task_config.sentencepiece_model_path = '/code/pretrained_v2/sentencepiece_en_tr.model'
translation_task = task_factory.get_task(task_config)
# model_en_tr = translation_task.build_model()
model_tr_en = translation_task.build_model() # we can use the same task
def translate(input_text, model):
tokenized = tokenizer.tokenize(input_text)
translated = model({'inputs' : tf.reshape(tokenized, [1, -1])})
return tokenizer.detokenize(translated['outputs']).numpy()[0].decode('utf-8')
# ignore = translate("test", model_en_tr)
ignore = translate("test", model_tr_en)
model_tr_en.load_weights("/code/pretrained_v2/tr_en/tr_en")
def predict(inp):
return translate(inp, model_tr_en)
def run():
demo = gr.Interface(
fn=predict,
inputs=gr.inputs.Textbox(label="Turkish"),
outputs=gr.outputs.Textbox(label="English"),
)
demo.launch(server_name="0.0.0.0", server_port=7860)
if __name__ == "__main__":
run()