import requests
import os
import gradio as gr

title = "Translate Text"
description = """"""
article = "Check out [the original repo](https://huggingface.co/language-tools/language-translation) that this demo is based off of."


TRANSLATION_API_URL = "https://api-inference.huggingface.co/models/t5-base"
LANG_ID_API_URL = "https://noe30ht5sav83xm1.us-east-1.aws.endpoints.huggingface.cloud"
ACCESS_TOKEN = os.environ.get("ACCESS_TOKEN")
# ACCESS_TOKEN = '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'
headers = {"Authorization": f"Bearer {ACCESS_TOKEN}"}


def query(payload):
  translation_response = requests.post(TRANSLATION_API_URL, headers=headers, json={
      "inputs": payload, "wait_for_model": True})
  translation = translation_response.json()[0]['translation_text']

  lang_id_response = requests.post(LANG_ID_API_URL, headers=headers, json={
      "inputs": payload, "wait_for_model": True})
  lang_id = lang_id_response.json()[0][0]

  return [lang_id, translation]


gr.Interface(
    query,
    gr.Textbox(lines=2),
    outputs=[
        gr.Textbox(lines=3, label="Detected Language"),
        gr.Textbox(lines=3, label="Translation")
    ],
    title=title,
    description=description,
    article=article
).launch()