Update app.py

app.py

@@ -1,55 +1,64 @@
-
 import gradio as gr
-from transformers import pipeline
-import re
-
-# 번역기 초기화 함수 (언어 선택 시 로딩)
-def load_translator(lang_code):
-    model_map = {
-        "en": "Helsinki-NLP/opus-mt-ko-en",
-        "zh": "Helsinki-NLP/opus-mt-ko-zh",
-        "vi": "Helsinki-NLP/opus-mt-ko-vi"
-    }
-    if lang_code == "ko":
-        return None
-    return pipeline("translation", model=model_map[lang_code])
-
-# 규칙 기반 요약
-def rule_based_summary(text):
-    key_phrases = re.findall(r"(폭염|호우|지진|감염병|불|화재|대피|주의|통제|착용|자제|신고|피해).*?(바랍니다|하세요|주의)", text)
-    flat = [" ".join(kp) for kp in key_phrases]
-    if not flat:
-        # 중요 문장이 없다면 명사 위주 추출
-        fallback = re.findall(r"[가-힣]{2,}", text)
-        return ". ".join(fallback[:3]) if fallback else "중요 문장 없음"
-    return ". ".join(flat[:3])
-
-# 전체 처리 함수
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+from googletrans import Translator
+
+# 모델 로드
+model_name = "Aoi785/kobart-disaster-summary-v1"
+tokenizer = AutoTokenizer.from_pretrained("digit82/kobart-summarization")
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
+
+# 번역기
+translator = Translator()
+
+# 요약 + 번역 함수
 def summarize_and_translate(text, lang):
     if not text.strip():
-        return "입력된 내용이 없습니다.", ""
-    summary = rule_based_summary(text)
-    if lang == "ko":
-        return summary, summary
+        return "입력된 텍스트가 없습니다.", ""
+
+    # 요약
     try:
-        translator = load_translator(lang)
-        result = translator(summary, max_length=100)
-        return summary, result[0]["translation_text"]
+        summary = summarizer(text, max_length=50, min_length=10, do_sample=False)[0]['summary_text']
     except Exception as e:
-        return summary, f"[번역 오류] {str(e)}"
+        summary = f"요약 실패: {str(e)}"
+        return summary, ""
+
+    # 번역
+    try:
+        if lang == "ko":
+            translated = summary
+        else:
+            translated = translator.translate(summary, dest=lang).text
+    except Exception as e:
+        translated = f"번역 실패: {str(e)}"
+
+    return summary, translated
 
 # Gradio UI
-iface = gr.Interface(
-    fn=summarize_and_translate,
-    inputs=[
-        gr.Textbox(lines=5, label="재난 문자 입력", placeholder="재난안전문자를 여기에 붙여넣으세요."),
-        gr.Dropdown(["ko", "en", "zh", "vi"], value="ko", label="번역 언어 선택")
-    ],
-    outputs=[
-        gr.Textbox(label="📌 핵심 요약"),
-        gr.Textbox(label="🌍 번역 요약")
-    ],
-    title="📢 DISASTER SNAP v5 Lite"
-)
-
-iface.launch()
+with gr.Blocks() as demo:
+    gr.Markdown("## 🆘 재난 문자 요약 및 번역기")
+    
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(label="재난 문자 입력", lines=5, placeholder="예: 내일까지 장마전선 영향으로 많은 비 예상...")
+            lang = gr.Dropdown(
+                choices=["ko(한국어)", "en(English)", "zh(중국어)", "vi(베트남어)"],
+                value="ko(한국어)",
+                label="번역 언어 선택"
+            )
+            run_button = gr.Button("요약 및 번역 실행")
+
+        with gr.Column():
+            output_summary = gr.Textbox(label="✅ 요약 결과")
+            output_trans = gr.Textbox(label="🌐 번역 결과")
+
+    def process_lang_code(lang_dropdown):
+        return lang_dropdown.split("(")[0]
+
+    run_button.click(
+        fn=lambda txt, ln: summarize_and_translate(txt, process_lang_code(ln)),
+        inputs=[input_text, lang],
+        outputs=[output_summary, output_trans]
+    )
+
+demo.launch()