Update app.py

app.py

@@ -1,95 +1,68 @@
 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"
-
-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_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)
 
+# 영어 -> 한국어 모델 로드
+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()
-
-        # 프롬프트 규칙 설정
         if source_lang == "한국어" and target_lang == "영어":
-            prompt = f"""
-            You are a 30-year veteran English translator.
-            You are renowned for your accurate and meticulous expressions.
-            Translate the following text from Korean to English:
-            '{text}'
-            - Translate accurately according to the context and flow of the content.
-            - Do not add any extra text or explanations, just provide the translation.
-            """
+            # 한국어 -> 영어 번역
+            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 == "한국어":
-            prompt = f"""
-            너는 30년차 한국어 전문 번역가이다.
-            정확하고 섬세한 표현력으로 유명한 한국어 전문 번역가이다.
-            다음 텍스트를 영어에서 한국어로 번역하라:
-            '{text}'
-            - 문맥과 내용의 흐름에 맞춰 정확한 번역을 한다.
-            - 추가적인 설명 없이 번역 결과만 제공하라.
-            """
+            # 영어 -> 한국어 번역
+            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 "지원하지 않는 언어 조합입니다."
-
-        # 번역에 최적화된 설정
-        response = client.text_generation(
-            prompt,
-            max_new_tokens=200,  # 번역 결과의 길이를 적절히 제한
-            temperature=0.3,     # 창의성을 낮춰 정확한 번역을 유도
-            top_p=0.9            # 높은 확률의 단어만 선택하도록 설정
-        )
-        return response.strip()
+        return result
     except Exception as e:
         return f"오류가 발생했습니다: {str(e)}"
 
-
 # Gradio UI 구성
 with gr.Blocks() as demo:
     gr.Markdown("# 번역기")
 
-    # 번역기 탭
-    with gr.Tab("번역기"):
-        # 언어 선택
-        source_lang = gr.Radio(
-            choices=["한국어", "영어"],
-            label="원본 언어",
-            value="한국어"
-        )
-        target_lang = gr.Radio(
-            choices=["한국어", "영어"],
-            label="목표 언어",
-            value="영어"
-        )
+    # 언어 선택
+    source_lang = gr.Radio(
+        choices=["한국어", "영어"],
+        label="원본 언어",
+        value="한국어"
+    )
+    target_lang = gr.Radio(
+        choices=["한국어", "영어"],
+        label="목표 언어",
+        value="영어"
+    )
 
-        # 입력 텍스트
-        input_text = gr.Textbox(label="번역할 텍스트", lines=5)
+    # 입력 텍스트
+    input_text = gr.Textbox(label="번역할 텍스트", lines=5)
 
-        # 번역 결과
-        translation_output = gr.Textbox(label="번역 결과", lines=5, interactive=False)
+    # 번역 결과
+    translation_output = gr.Textbox(label="번역 결과", lines=5, interactive=False)
 
-        # 번역 버튼
-        translate_button = gr.Button("번역")
+    # 번역 버튼
+    translate_button = gr.Button("번역")
 
-        # 번역 함수 연결
-        translate_button.click(
-            fn=translate_text,
-            inputs=[input_text, source_lang, target_lang],
-            outputs=translation_output
-        )
+    # 번역 함수 연결
+    translate_button.click(
+        fn=translate_text,
+        inputs=[input_text, source_lang, target_lang],
+        outputs=translation_output
+    )
 
 # 메인 실행��
 if __name__ == "__main__":