Update app.py

app.py

@@ -1,36 +1,34 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-import os
+from transformers import MarianMTModel, MarianTokenizer
 
-# Cohere Command R+ 모델 ID 정의
-COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"
+# 한국어 -> 영어 모델 로드
+model_name_ko_en = "Helsinki-NLP/opus-mt-ko-en"
+tokenizer_ko_en = MarianTokenizer.from_pretrained(model_name_ko_en)
+model_ko_en = MarianMTModel.from_pretrained(model_name_ko_en)
 
-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)
+# 영어 -> 한국어 모델 로드
+model_name_en_ko = "Helsinki-NLP/opus-mt-en-ko"
+tokenizer_en_ko = MarianTokenizer.from_pretrained(model_name_en_ko)
+model_en_ko = MarianMTModel.from_pretrained(model_name_en_ko)
 
 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=50,  # 불필요한 텍스트 생성을 방지하기 위해 토큰 수 제한
-            temperature=0.3,    # 창의성을 낮춰 정확한 번역 유도
-            top_p=0.9           # 높은 확률의 단어만 선택하도록 설정
-        )
-        # 번역 결과에서 불필요한 부분 제거
-        return response.strip().split("\n")[0]  # 첫 번째 줄만 반환
+        if source_lang == "한국어" and target_lang == "영어":
+            # 한국어 -> 영어 번역
+            tokenized = tokenizer_ko_en(text, return_tensors="pt", padding=True)
+            translated = model_ko_en.generate(**tokenized)
+            result = tokenizer_ko_en.decode(translated[0], skip_special_tokens=True)
+        elif source_lang == "영어" and target_lang == "한국어":
+            # 영어 -> 한국어 번역
+            tokenized = tokenizer_en_ko(text, return_tensors="pt", padding=True)
+            translated = model_en_ko.generate(**tokenized)
+            result = tokenizer_en_ko.decode(translated[0], skip_special_tokens=True)
+        else:
+            return "지원하지 않는 언어 조합입니다."
+        return result
     except Exception as e:
         return f"오류가 발생했습니다: {str(e)}"