app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import sys

# 1. 모델 및 토크나이저 설정
# Hugging Face Spaces의 무료 CPU 환경(16GB RAM)에 맞춰 600M 모델 사용
model_name = "facebook/nllb-200-distilled-600M"

print(f"모델({model_name})을 로드하는 중입니다... 잠시만 기다려주세요.")

# 전역 변수로 선언
tokenizer = None
model = None

try:
    # 토크나이저와 모델 로드
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    print("모델 로드 완료!")
except Exception as e:
    # 모델 로드 실패 시 앱을 강제로 종료하여 Logs 탭에서 정확한 원인을 볼 수 있게 함
    print(f"❌ 모델 로드 중 치명적인 오류 발생: {e}")
    sys.exit(1)

# 2. 언어 코드 매핑
LANG_CODES = {
    "영어 (English)": "eng_Latn",
    "일본어 (Japanese)": "jpn_Jpan",
    "중국어 (Chinese Simplified)": "zho_Hans"
}

TARGET_LANG_CODE = "kor_Hang"  # 한국어

def translate_text(text, source_lang_name):
    """
    입력 텍스트를 한국어로 번역
    """
    if not text:
        return "번역할 내용을 입력해주세요."

    if model is None or tokenizer is None:
        return "모델이 로드되지 않았습니다. 서버 로그를 확인해주세요."

    try:
        # 입력 언어 코드 가져오기
        src_code = LANG_CODES.get(source_lang_name)
        
        # 번역 옵션 설정: 입력 언어 지정
        tokenizer.src_lang = src_code
        
        # 입력 텍스트 토큰화
        inputs = tokenizer(text, return_tensors="pt")
        
        # [중요] 메모리 최적화를 위해 no_grad() 사용
        with torch.no_grad():

            target_token_id = tokenizer.convert_tokens_to_ids(TARGET_LANG_CODE)
            
            generated_tokens = model.generate(
                **inputs,
                forced_bos_token_id=target_token_id,
                max_length=500,
                # [중요] CPU 환경 안정성을 위해 Beam Search 대신 Greedy Search 사용
                num_beams=1
            )
        
        # 결과 디코딩
        result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
        return result
        
    except Exception as e:
        return f"번역 에러: {str(e)}"

# 3. Gradio 인터페이스
with gr.Blocks(title="한글로 (Hangullo) - 다국어 번역기") as demo:
    gr.Markdown(
        """
        # 🇰🇷 한글로 (Hangullo)
        **영어, 일본어, 중국어**를 입력하면 자연스러운 **한국어**로 번역해 드립니다.
        *(Powered by Meta NLLB-200)*
        """
    )
    
    with gr.Row():
        with gr.Column():
            src_lang = gr.Dropdown(
                choices=list(LANG_CODES.keys()), 
                value="영어 (English)", 
                label="입력 언어"
            )
            input_text = gr.Textbox(
                lines=5, 
                placeholder="번역할 문장을 입력하세요...", 
                label="입력 (Source)"
            )
            translate_btn = gr.Button("한국어로 변환", variant="primary")
            
        with gr.Column():
            output_text = gr.Textbox(
                lines=5, 
                label="한국어 결과 (Korean)", 
                interactive=False
            )
    
    # 예제 데이터
    gr.Examples(
        examples=[
            ["The quick brown fox jumps over the lazy dog.", "영어 (English)"],
            ["AIの発展によって、私たちの生活は大きく変化しています。", "일본어 (Japanese)"],
            ["今天天气真好，我们去公园散步吧。", "중국어 (Chinese Simplified)"]
        ],
        inputs=[input_text, src_lang]
    )

    translate_btn.click(
        fn=translate_text, 
        inputs=[input_text, src_lang], 
        outputs=output_text
    )

# 4. 앱 실행
if __name__ == "__main__":
    # [중요] 큐(Queue)를 활성화하여 요청 충돌 방지
    demo.queue().launch()