import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# 모델과 토크나이저는 한 번만 로드되도록 함수 외부에 선언
model = AutoModelForSeq2SeqLM.from_pretrained("Kyudan/opus-mt-en-ro-finetuned-en-to-ro")
tokenizer = AutoTokenizer.from_pretrained("Kyudan/opus-mt-en-ro-finetuned-en-to-ro")

def respond(text):
    # 입력 텍스트를 토큰화
    inputs = tokenizer.encode(text, return_tensors="pt")

    # 모델을 사용하여 번역
    outputs = model.generate(inputs)

    # 번역된 텍스트를 디코딩
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return translated_text

# Gradio 인터페이스 설정
demo = gr.Interface(
    fn=respond,
    inputs="text",
    outputs="text",
    title="Translate English to Romanian"
)

if __name__ == "__main__":
    demo.launch(share=True)