import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# 모델과 토크나이저 로드
model_name = "Helsinki-NLP/opus-mt-ko-en"  # 한국어 to 영어 모델
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

# 번역 함수 정의
def translate(text):
    # 텍스트를 토큰화하고 모델에 입력
    tokenized_text = tokenizer.prepare_seq2seq_batch([text], return_tensors="pt")
    translated = model.generate(**tokenized_text)
    # 번역된 텍스트를 디코딩
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Gradio 인터페이스 설정
iface = gr.Interface(
    fn=translate,
    inputs=gr.inputs.Textbox(lines=2, placeholder="번역할 텍스트를 입력하세요..."),
    outputs="text",
    title="한국어 to 영어 번역기",
    description="한국어 텍스트를 영어로 번역해주는 간단한 번역기입니다."
)

# 앱 실행
iface.launch()