app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os

# Cohere Command R+ 모델 ID 정의
COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"

def get_client():
    """
    Cohere Command R+ 모델을 위한 InferenceClient 생성.
    토큰은 환경 변수에서 가져옴.
    """
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise ValueError("HuggingFace API 토큰이 필요합니다.")
    return InferenceClient(COHERE_MODEL, token=hf_token)

def translate_text(text, source_lang, target_lang):
    """
    텍스트를 번역하는 함수.
    """
    try:
        client = get_client()
        # 프롬프트를 단순화하여 번역만 수행하도록 지시
        prompt = f"Translate the following text from {source_lang} to {target_lang}: '{text}'. Do not add any extra explanations."
        response = client.text_generation(
            prompt,
            max_new_tokens=100,  # 불필요한 텍스트 생성을 방지하기 위해 토큰 수 제한
            temperature=0.3,     # 창의성을 낮춰 정확한 번역 유도
            top_p=0.9            # 높은 확률의 단어만 선택하도록 설정
        )
        return response.strip()  # 불필요한 공백 제거
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"

# Gradio UI 구성
with gr.Blocks() as demo:
    gr.Markdown("# 한국어 ↔ 영어 번역기")

    # 언어 선택
    source_lang = gr.Radio(
        choices=["한국어", "영어"],
        label="원본 언어",
        value="한국어"
    )
    target_lang = gr.Radio(
        choices=["한국어", "영어"],
        label="목표 언어",
        value="영어"
    )

    # 입력 텍스트
    input_text = gr.Textbox(label="번역할 텍스트", lines=5)

    # 번역 결과
    translation_output = gr.Textbox(label="번역 결과", lines=5, interactive=False)

    # 번역 버튼
    translate_button = gr.Button("번역")

    # 번역 함수 연결
    translate_button.click(
        fn=translate_text,
        inputs=[input_text, source_lang, target_lang],
        outputs=translation_output
    )

# 메인 실행부
if __name__ == "__main__":
    demo.launch()