Update app.py

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
 import torch
+import fitz  # PyMuPDF
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # モデルID
 model_id = "tencent/HY-MT1.5-1.8B"
 
 # 環境に合わせてデバイスと精度を自動選択
-# Freeスペース(CPU)の場合はfloat32、GPUがある場合はfloat16を使用
 if torch.cuda.is_available():
     device = "cuda"
     dtype = torch.float16
@@ -20,17 +20,41 @@ print(f"Loading model on {device} with {dtype}...")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map=device, # autoではなく明示的に指定
+    device_map=device,
     torch_dtype=dtype
 )
 
+def extract_text_from_pdf(pdf_file):
+    """PDF에서 텍스트 추출"""
+    if pdf_file is None:
+        return ""
+    
+    try:
+        doc = fitz.open(pdf_file.name)
+        full_text = ""
+        
+        for page_num, page in enumerate(doc, 1):
+            text = page.get_text("text")
+            if text.strip():
+                full_text += f"\n--- Page {page_num} ---\n{text.strip()}\n"
+        
+        doc.close()
+        return full_text.strip()
+    
+    except Exception as e:
+        return f"❌ PDF 추출 오류: {str(e)}"
+
 def translate_text(source_text, target_lang):
+    """텍스트 번역"""
+    if not source_text or not source_text.strip():
+        return "입력 텍스트가 없습니다."
+    
     # プロンプトの切り替えロジック
     if "Chinese" in target_lang or "中文" in target_lang:
         prompt = f"将以下文本翻译为{target_lang}，注意只需要输出翻译后的结果，不要额外解释：\n{source_text}"
     else:
         prompt = f"Translate the following segment into {target_lang}, without additional explanation.\n{source_text}"
-
+    
     messages = [{"role": "user", "content": prompt}]
     
     # 入力処理
@@ -40,7 +64,7 @@ def translate_text(source_text, target_lang):
         add_generation_prompt=False, 
         return_tensors="pt"
     ).to(device)
-
+    
     # 生成実行
     with torch.no_grad():
         generated_ids = model.generate(
@@ -50,7 +74,7 @@ def translate_text(source_text, target_lang):
             top_p=0.6,
             repetition_penalty=1.05
         )
-
+    
     # 出力処理
     input_length = text_input.shape[1]
     response = generated_ids[0][input_length:]
@@ -58,27 +82,137 @@ def translate_text(source_text, target_lang):
     
     return decoded_output
 
+def translate_long_text(source_text, target_lang, chunk_size=1500):
+    """긴 텍스트를 청크로 나눠서 번역"""
+    if not source_text or not source_text.strip():
+        return "입력 텍스트가 없습니다."
+    
+    # 짧은 텍스트는 바로 번역
+    if len(source_text) <= chunk_size:
+        return translate_text(source_text, target_lang)
+    
+    # 긴 텍스트는 문단 단위로 분할
+    paragraphs = source_text.split('\n\n')
+    chunks = []
+    current_chunk = ""
+    
+    for para in paragraphs:
+        if len(current_chunk) + len(para) < chunk_size:
+            current_chunk += para + "\n\n"
+        else:
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+            current_chunk = para + "\n\n"
+    
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+    
+    # 각 청크 번역
+    translated_chunks = []
+    for i, chunk in enumerate(chunks):
+        print(f"Translating chunk {i+1}/{len(chunks)}...")
+        translated = translate_text(chunk, target_lang)
+        translated_chunks.append(translated)
+    
+    return "\n\n".join(translated_chunks)
+
+def process_pdf_and_translate(pdf_file, target_lang):
+    """PDF 업로드 → 텍스트 추출 → 번역"""
+    if pdf_file is None:
+        return "", "PDF 파일을 업로드해주세요."
+    
+    # 텍스트 추출
+    extracted_text = extract_text_from_pdf(pdf_file)
+    
+    if extracted_text.startswith("❌"):
+        return "", extracted_text
+    
+    if not extracted_text.strip():
+        return "", "PDF에서 텍스트를 추출할 수 없습니다."
+    
+    # 번역
+    translated_text = translate_long_text(extracted_text, target_lang)
+    
+    return extracted_text, translated_text
+
+def translate_input_text(source_text, target_lang):
+    """입력 텍스트 번역"""
+    return translate_long_text(source_text, target_lang)
+
 # UIの構築
-langs = ["Japanese", "English", "Chinese", "Korean", "French", "German", "Spanish"]
+langs = ["Japanese", "English", "Chinese", "Korean", "French", "German", "Spanish", "한국어", "日本語", "中文"]
 
-with gr.Blocks() as demo:
-    gr.Markdown("# 🚀 HY-MT1.5-1.8B Translator (Spaces)")
-    gr.Markdown("Tencentの1.8Bモデルを使用した翻訳デモです。")
-    
-    with gr.Row():
-        with gr.Column():
-            input_text = gr.Textbox(label="原文 (Source Text)", lines=5, placeholder="ここに入力...")
-            target_lang = gr.Dropdown(choices=langs, value="English", label="翻���先 (Target Language)")
-            submit_btn = gr.Button("翻訳 (Translate)", variant="primary")
+with gr.Blocks(title="HY-MT1.5 Translator") as demo:
+    gr.Markdown("# 🚀 HY-MT1.5-1.8B Translator")
+    gr.Markdown("Tencent의 1.8B 번역 모델을 사용한 텍스트/PDF 번역 데모입니다.")
+    
+    with gr.Tabs():
+        # ============ Tab 1: 텍스트 번역 ============
+        with gr.TabItem("📝 Text Translation"):
+            with gr.Row():
+                with gr.Column():
+                    input_text = gr.Textbox(
+                        label="원문 (Source Text)", 
+                        lines=10, 
+                        placeholder="번역할 텍스트를 입력하세요..."
+                    )
+                    target_lang_text = gr.Dropdown(
+                        choices=langs, 
+                        value="English", 
+                        label="번역 언어 (Target Language)"
+                    )
+                    translate_btn = gr.Button("🔄 번역 (Translate)", variant="primary")
+                
+                with gr.Column():
+                    output_text = gr.Textbox(
+                        label="번역 결과 (Result)", 
+                        lines=10, 
+                        interactive=False
+                    )
+            
+            translate_btn.click(
+                fn=translate_input_text,
+                inputs=[input_text, target_lang_text],
+                outputs=output_text
+            )
         
-        with gr.Column():
-            output_text = gr.Textbox(label="結果 (Result)", lines=5, interactive=False)
-
-    submit_btn.click(
-        fn=translate_text,
-        inputs=[input_text, target_lang],
-        outputs=output_text
-    )
+        # ============ Tab 2: PDF 번역 ============
+        with gr.TabItem("📄 PDF Translation"):
+            gr.Markdown("### PDF 파일을 업로드하면 텍스트를 추출하고 번역합니다.")
+            
+            with gr.Row():
+                with gr.Column():
+                    pdf_input = gr.File(
+                        label="📄 PDF 파일 업로드",
+                        file_types=[".pdf"]
+                    )
+                    target_lang_pdf = gr.Dropdown(
+                        choices=langs, 
+                        value="English", 
+                        label="번역 언어 (Target Language)"
+                    )
+                    translate_pdf_btn = gr.Button("🔄 PDF 번역", variant="primary")
+            
+            with gr.Row():
+                with gr.Column():
+                    extracted_text = gr.Textbox(
+                        label="📋 추출된 원문 (Extracted Text)",
+                        lines=15,
+                        interactive=False
+                    )
+                
+                with gr.Column():
+                    translated_pdf_text = gr.Textbox(
+                        label="📋 번역 결과 (Translated Text)",
+                        lines=15,
+                        interactive=False
+                    )
+            
+            translate_pdf_btn.click(
+                fn=process_pdf_and_translate,
+                inputs=[pdf_input, target_lang_pdf],
+                outputs=[extracted_text, translated_pdf_text]
+            )
 
 # 起動
 demo.launch()