import os

import gradio as gr

import subprocess

from transformers import AutoModel

from huggingface_hub import snapshot_download

MODELS_PATH = "./models"
HF_CACHE_DIR = "./hf_cache"

def download_model(repo_id, revision="main"):
    return snapshot_download(repo_id=repo_id, revision=revision, local_dir=os.path.join(MODELS_PATH, repo_id), cache_dir=HF_CACHE_DIR)

model_dir_gl_en = download_model("proxectonos/Nos_MT-OpenNMT-gl-en", revision="main")
model_dir_en_gl = download_model("proxectonos/Nos_MT-OpenNMT-en-gl", revision="main")

def translate(input_text):
    print input_text
    command = f"onmt_translate -src {input_text} -model {os.path.join(MODELS_PATH, model_dir_en_gl)}/NOS-MT-OpenNMT-en-gl --output ./output_file.txt --replace_unk"
    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate()
    if process.returncode != 0:
        raise Exception(f"Error occurred: {stderr.decode().strip()}")
    return stdout.decode().strip()

demo = gr.Interface(fn=translate, inputs="textbox", outputs="textbox")
demo.launch()