app.py

app.py

@@ -1,55 +1,139 @@
+import re
+import math
 import gradio as gr
+import torch
 from transformers import PreTrainedTokenizerFast, BartForConditionalGeneration
 
-# 1. 모델 & 토크나이저 로드
-MODEL_NAME = "gangyeolkim/kobart-korean-summarizer-v2"
+# ✅ 공개 KoBART 모델
+MODEL_NAME = "gogamza/kobart-base-v2"
+
 tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL_NAME)
 model = BartForConditionalGeneration.from_pretrained(MODEL_NAME)
 
-# 2. 요약 함수
-def summarize(text, min_len, max_len):
-    if not text.strip():
-        return "⚠️ 요약할 텍스트를 입력하세요."
-    
-    # 토큰화
-    inputs = tokenizer(
-        [text],
-        max_length=1024,
-        truncation=True,
-        return_tensors="pt"
+# CPU 동적 양자화 적용
+try:
+    model = torch.quantization.quantize_dynamic(
+        model, {torch.nn.Linear}, dtype=torch.qint8
     )
+except Exception:
+    pass
 
-    # 모델 추론
-    summary_ids = model.generate(
-        inputs["input_ids"],
-        num_beams=4,
-        min_length=min_len,
-        max_length=max_len,
-        early_stopping=True
-    )
+model.eval()
+
+# ===== 유틸 함수 =====
+def normalize_text(text: str) -> str:
+    return re.sub(r"\s+", " ", text).strip()
+
+def split_into_sentences(text: str):
+    text = text.replace("\n", " ")
+    parts = re.split(r"(?<=[\.!?])\s+", text)
+    return [p.strip() for p in parts if p.strip()]
+
+def token_length(s: str) -> int:
+    return len(tokenizer.encode(s, add_special_tokens=False))
+
+def chunk_by_tokens(sentences, max_tokens=900):
+    chunks, cur, cur_tokens = [], [], 0
+    for s in sentences:
+        tl = token_length(s)
+        if tl > max_tokens:
+            piece_size = max(200, int(len(s) * (max_tokens / tl)))
+            for i in range(0, len(s), piece_size):
+                sub = s[i:i+piece_size]
+                if sub.strip():
+                    chunks.append(sub.strip())
+            cur, cur_tokens = [], 0
+            continue
+        if cur_tokens + tl <= max_tokens:
+            cur.append(s)
+            cur_tokens += tl
+        else:
+            if cur:
+                chunks.append(" ".join(cur))
+            cur, cur_tokens = [s], tl
+    if cur:
+        chunks.append(" ".join(cur))
+    return chunks
 
-    # 디코딩
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    return summary
+# ===== 요약 함수 =====
+def summarize_raw(text: str, min_len: int, max_len: int) -> str:
+    inputs = tokenizer([text], max_length=1024, truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        summary_ids = model.generate(
+            inputs["input_ids"],
+            num_beams=4,
+            min_length=min_len,
+            max_length=max_len,
+            early_stopping=True,
+            no_repeat_ngram_size=3
+        )
+    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+
+def apply_style_prompt(text: str, mode: str, final: bool=False) -> str:
+    if mode == "concise":
+        inst = "다음 한국어 텍스트를 핵심만 간결하게 요약하세요."
+    elif mode == "explanatory":
+        inst = "다음 한국어 텍스트를 맥락을 보존하며 이해하기 쉽게 요약하세요."
+    else:
+        inst = "다음 한국어 텍스트를 bullet 형태로 핵심만 요약하세요."
+    if final:
+        inst += " 이 요약은 최종본입니다."
+    return f"{inst}\n\n[텍스트]\n{text}"
+
+def postprocess(summary: str, mode: str) -> str:
+    s = summary.strip()
+    s = re.sub(r"\s+", " ", s)
+    if mode == "bullets":
+        bullets = re.split(r"\s*[-•]\s*", s)
+        bullets = [b.strip() for b in bullets if b.strip()]
+        if len(bullets) > 1:
+            s = "\n".join([f"- {b}" for b in bullets])
+        else:
+            parts = re.split(r"(?<=[\.!?])\s+", s)
+            parts = [p.strip() for p in parts if p.strip()]
+            s = "\n".join([f"- {p}" for p in parts])
+    return s
+
+def summarize_long(text: str, target_chars: int, mode: str):
+    text = normalize_text(text)
+    if not text:
+        return "⚠️ 요약할 텍스트를 입력하세요."
+    approx_tokens = token_length(text)
+    if approx_tokens <= 1000:
+        min_len = max(60, int(target_chars * 0.4 / 2))
+        max_len = max(120, int(target_chars * 0.8 / 2))
+        return postprocess(summarize_raw(apply_style_prompt(text, mode), min_len, max_len), mode)
+    sentences = split_into_sentences(text)
+    chunks = chunk_by_tokens(sentences, max_tokens=900)
+    partial_summaries = []
+    budget_total = int(target_chars * 1.5)
+    per_chunk_chars = max(250, budget_total // max(1, len(chunks)))
+    for c in chunks:
+        min_len = max(50, int(per_chunk_chars * 0.4 / 2))
+        max_len = max(100, int(per_chunk_chars * 0.9 / 2))
+        psum = summarize_raw(apply_style_prompt(c, mode), min_len, max_len)
+        partial_summaries.append(psum)
+    merged = normalize_text(" ".join(partial_summaries))
+    final_min = max(80, int(target_chars * 0.45 / 2))
+    final_max = max(160, int(target_chars * 1.05 / 2))
+    return postprocess(summarize_raw(apply_style_prompt(merged, mode, final=True), final_min, final_max), mode)
+
+# ===== Gradio UI =====
+def ui_summarize(text, target_len, style):
+    mode = {"간결형":"concise", "설명형":"explanatory", "핵심 bullet":"bullets"}[style]
+    return summarize_long(text, int(target_len), mode)
 
-# 3. Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 📝 KoBART 한국어 요약기 (CPU 최적화 가능)")
+    gr.Markdown("## 📝 KoBART 한국어 요약기 (공개 모델 gogamza/kobart-base-v2)")
     with gr.Row():
         with gr.Column():
-            input_text = gr.Textbox(
-                label="원문 입력 (최대 2000자)",
-                lines=15,
-                placeholder="여기에 요약할 한국어 텍스트를 입력하세요."
-            )
-            min_len = gr.Slider(50, 500, value=100, step=10, label="최소 요약 길이")
-            max_len = gr.Slider(100, 1500, value=300, step=10, label="최대 요약 길이")
+            input_text = gr.Textbox(label="원문 입력", lines=16)
+            style = gr.Radio(["간결형", "설명형", "핵심 bullet"], value="간결형", label="요약 스타일")
+            target_len = gr.Slider(300, 1500, value=1000, step=50, label="목표 요약 길이(문자)")
             btn = gr.Button("요약 실행")
         with gr.Column():
-            output_text = gr.Textbox(label="요약 결과", lines=15)
-
-    btn.click(summarize, inputs=[input_text, min_len, max_len], outputs=output_text)
+            output_text = gr.Textbox(label="요약 결과", lines=16)
+    btn.click(ui_summarize, inputs=[input_text, target_len, style], outputs=output_text)
 
-# 4. 실행
 if __name__ == "__main__":
     demo.launch()