import gradio as gr
from transformers import pipeline

# 3. 번역 파이프라인 생성
print("한국어->영어 번역 모델(NHNDQ/nllb-finetuned-ko2en)을 로딩합니다...")
ko_to_en_translator = pipeline('translation', model='NHNDQ/nllb-finetuned-ko2en')

print("영어->한국어 번역 모델(NHNDQ/nllb-finetuned-en2ko)을 로딩합니다...")
en_to_ko_translator = pipeline('translation', model='NHNDQ/nllb-finetuned-en2ko')

print("\n모든 모델 로딩 완료!")

# 4. 번역을 수행하는 메인 함수 정의
def translate_text(text, direction):
    """
    선택된 번역 방향에 따라 적절한 NLLB 파이프라인을 사용하여 텍스트를 번역합니다.
    """
    try:
        # 입력 텍스트가 비어있는 경우, 빈 문자열 반환
        if not text.strip():
            return ""

        if direction == "한국어 → 영어":
            # [수정] src_lang과 tgt_lang을 명시적으로 지정합니다.
            result = ko_to_en_translator(text, src_lang="kor_Hang", tgt_lang="eng_Latn")
        else: # "영어 → 한국어"
            # [수정] src_lang과 tgt_lang을 명시적으로 지정합니다.
            result = en_to_ko_translator(text, src_lang="eng_Latn", tgt_lang="kor_Hang")

        # 파이프라인 결과는 리스트 안의 딕셔너리 형태 [{'translation_text': '...'}] 입니다.
        return result[0]['translation_text']
    except Exception as e:
        return f"번역 중 오류가 발생했습니다: {e}"

# 5. Gradio 인터페이스 생성 및 실행
demo = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(lines=7, label="입력 텍스트", placeholder="번역할 내용을 입력하세요..."),
        gr.Radio(
            choices=["한국어 → 영어", "영어 → 한국어"],
            value="한국어 → 영어",
            label="번역 방향 선택"
        )
    ],
    outputs=gr.Textbox(lines=7, label="번역 결과"),
    title="NHN NLLB 기반 한국어-영어 번역기 🇰🇷↔️🇬🇧",
    description="NHN의 NLLB 미세 조정 모델을 사용합니다. 번역할 내용과 방향을 선택해 주세요.",
    allow_flagging="never"
)

# Colab에서 외부 공유 링크를 생성하기 위해 share=True 옵션 사용
demo.launch(share=True)