from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import gradio as grad

model_name = "Helsinki-NLP/opus-mt-ko-en"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

#opus_translator = pipeline("translation", model=model_name)

def translate(text):
  inputs = tokenizer(text, return_tensors="pt")
  translation_output = model.generate(**input)
  response = tokenizer.decode(translation_output[0], skip_special_tokens=True)

  #response = opus_translator(text)
  return response

# Web UI
grad.Interface(translate, inputs=["text"], outputs=["text"]).launch()