Update app.py

app.py

@@ -1,127 +1,127 @@
-import torch
-import gradio as gr
-import requests
-from bs4 import BeautifulSoup
-from transformers import T5Tokenizer, T5ForConditionalGeneration, T5TokenizerFast, AutoTokenizer, AutoModelForSeq2SeqLM
-import nltk
-
-# nltk의 'punkt' 데이터 다운로드 (문장 토큰화를 위해)
-nltk.download('punkt')
-
-# 뉴스 기사 본문 또는 URL을 입력받아,
-# 만약 입력값이 "http"로 시작하면 네이버 뉴스 URL로 간주하여 본문을 추출하고,
-# 그렇지 않으면 그대로 텍스트를 반환하는 함수
-def extract_naver_news_article(text_or_url):
-    try:
-        text_or_url = text_or_url.strip()
-        if text_or_url.startswith("http"):
-            headers = {
-                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
-            }
-            response = requests.get(text_or_url, headers=headers)
-            if response.status_code != 200:
-                return f"페이지 요청 실패 (HTTP {response.status_code})"
-            soup = BeautifulSoup(response.text, "html.parser")
-            article_content = soup.select_one("#dic_area")
-            if not article_content:
-                article_content = soup.select_one(".go_trans._article_content")
-            if not article_content:
-                return "기사 본문을 찾을 수 없습니다."
-            paragraphs = article_content.find_all("p")
-            article_text = "\n".join([p.get_text(strip=True) for p in paragraphs]) if paragraphs else article_content.get_text(strip=True)
-            return article_text
-        else:
-            return text_or_url
-    except Exception as e:
-        return f"오류 발생: {e}"
-
-# 제목 생성 모델 로드 함수 (파인튜닝된 T5 모델)
-def load_title_model(checkpoint_path):
-    try:
-        title_tokenizer = T5TokenizerFast.from_pretrained(checkpoint_path)
-        title_model = T5ForConditionalGeneration.from_pretrained(checkpoint_path)
-        return title_model, title_tokenizer
-    except Exception as e:
-        print(f"Error loading title model: {e}")
-        return None, None
-
-# 뉴스 기사 요약 모델 로드 함수 
-def load_summarization_model(model_dir):
-    try:
-        summarizer_tokenizer = AutoTokenizer.from_pretrained(model_dir)
-        summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
-        return summarizer_model, summarizer_tokenizer
-    except Exception as e:
-        print(f"Error loading summarization model: {e}")
-        return None, None
-
-# 모델 경로 설정 (사용자 환경에 맞게 변경)
-checkpoint_path = r"C:\Users\onebe\_ABH\prac\NLP\final\title_gen"  # 파인튜닝된 T5 제목 생성 모델 경로
-model_dir = r"C:\Users\onebe\_ABH\prac\NLP\final\abs_gen"                   # 뉴스 기사 요약 모델 경로
-
-title_model, title_tokenizer = load_title_model(checkpoint_path)
-summarizer_model, summarizer_tokenizer = load_summarization_model(model_dir)
-
-# 뉴스 기사 제목 생성 함수
-def generate_title(article_text, max_length=20, num_beams=10, early_stopping=True):
-    try:
-        input_ids = title_tokenizer.encode(f"summarize: {article_text}", return_tensors="pt", truncation=True)
-        outputs = title_model.generate(
-            input_ids,
-            max_length=max_length,
-            num_beams=num_beams,
-            early_stopping=early_stopping
-        )
-        title = title_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return title
-    except Exception as e:
-        return f"Error generating title: {e}"
-
-# 뉴스 기사 요약 생성 함수
-def generate_summary(article_text, max_length=40, num_beams=10, early_stopping=True):
-    try:
-        input_ids = summarizer_tokenizer.encode(article_text, return_tensors="pt", truncation=True)
-        outputs = summarizer_model.generate(
-            input_ids,
-            max_length=max_length,
-            num_beams=num_beams,
-            early_stopping=early_stopping
-        )
-        summary = summarizer_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return summary
-    except Exception as e:
-        return f"Error generating summary: {e}"
-
-# 뉴스 기사 본문 또는 URL을 입력받아 본문 추출 후 제목과 요약 생성하는 함수
-def process_news(input_text, title_max_length, title_num_beams, title_early_stopping,
-                 summary_max_length, summary_num_beams, summary_early_stopping):
-    article_text = extract_naver_news_article(input_text)
-    title = generate_title(article_text, max_length=title_max_length, num_beams=title_num_beams, early_stopping=title_early_stopping)
-    summary = generate_summary(article_text, max_length=summary_max_length, num_beams=summary_num_beams, early_stopping=summary_early_stopping)
-    return article_text, title, summary
-
-# Gradio 인터페이스 실행 함수
-def launch_gradio_interface():
-    interface = gr.Interface(
-        fn=process_news,
-        inputs=[
-            gr.Textbox(label="뉴스 기사 본문 또는 URL", placeholder="뉴스 기사 본문을 직접 입력하거나 URL을 입력하세요.", lines=10),
-            gr.Slider(0, 40, value=20, step=1, label="제목 생성 - 최대 길이"),
-            gr.Slider(1, 10, value=10, step=1, label="제목 생성 - 탐색 빔 수"),
-            gr.Checkbox(value=True, label="제목 생성 - 조기 종료"),
-            gr.Slider(0, 256, value=128, step=1, label="요약 생성 - 최대 길이"),
-            gr.Slider(1, 10, value=10, step=1, label="요약 생성 - 탐색 빔 수"),
-            gr.Checkbox(value=True, label="요약 생성 - 조기 종료")
-        ],
-        outputs=[
-            gr.Textbox(label="추출되거나 입력된 기사 본문"),
-            gr.Textbox(label="생성된 제목"),
-            gr.Textbox(label="생성된 요약문")
-        ],
-        title="네이버 뉴스 제목 및 요약 생성기",
-        description="뉴스 기사 본문 또는 URL을 입력하면, 본문을 자동 추출한 후 제목과 요약문을 생성합니다."
-    )
-    interface.launch(share=True)
-
-if __name__ == '__main__':
-    launch_gradio_interface()
+import torch
+import gradio as gr
+import requests
+from bs4 import BeautifulSoup
+from transformers import T5Tokenizer, T5ForConditionalGeneration, T5TokenizerFast, AutoTokenizer, AutoModelForSeq2SeqLM
+import nltk
+
+# nltk의 'punkt' 데이터 다운로드 (문장 토큰화를 위해)
+nltk.download('punkt')
+
+# 뉴스 기사 본문 또는 URL을 입력받아,
+# 만약 입력값이 "http"로 시작하면 네이버 뉴스 URL로 간주하여 본문을 추출하고,
+# 그렇지 않으면 그대로 텍스트를 반환하는 함수
+def extract_naver_news_article(text_or_url):
+    try:
+        text_or_url = text_or_url.strip()
+        if text_or_url.startswith("http"):
+            headers = {
+                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+            }
+            response = requests.get(text_or_url, headers=headers)
+            if response.status_code != 200:
+                return f"페이지 요청 실패 (HTTP {response.status_code})"
+            soup = BeautifulSoup(response.text, "html.parser")
+            article_content = soup.select_one("#dic_area")
+            if not article_content:
+                article_content = soup.select_one(".go_trans._article_content")
+            if not article_content:
+                return "기사 본문을 찾을 수 없습니다."
+            paragraphs = article_content.find_all("p")
+            article_text = "\n".join([p.get_text(strip=True) for p in paragraphs]) if paragraphs else article_content.get_text(strip=True)
+            return article_text
+        else:
+            return text_or_url
+    except Exception as e:
+        return f"오류 발생: {e}"
+
+# 제목 생성 모델 로드 함수 (파인튜닝된 T5 모델)
+def load_title_model(checkpoint_path):
+    try:
+        title_tokenizer = T5TokenizerFast.from_pretrained(checkpoint_path)
+        title_model = T5ForConditionalGeneration.from_pretrained(checkpoint_path)
+        return title_model, title_tokenizer
+    except Exception as e:
+        print(f"Error loading title model: {e}")
+        return None, None
+
+# 뉴스 기사 요약 모델 로드 함수 
+def load_summarization_model(model_dir):
+    try:
+        summarizer_tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+        return summarizer_model, summarizer_tokenizer
+    except Exception as e:
+        print(f"Error loading summarization model: {e}")
+        return None, None
+
+# 모델 경로 설정 (사용자 환경에 맞게 변경)
+checkpoint_path = 'onebeans/keT5-news-summarizer'  # 파인튜닝된 T5 제목 생성 모델 경로
+model_dir = 'onebeans/keT5-news-title-gen'                  # 뉴스 기사 요약 모델 경로
+
+title_model, title_tokenizer = load_title_model(checkpoint_path)
+summarizer_model, summarizer_tokenizer = load_summarization_model(model_dir)
+
+# 뉴스 기사 제목 생성 함수
+def generate_title(article_text, max_length=20, num_beams=10, early_stopping=True):
+    try:
+        input_ids = title_tokenizer.encode(f"summarize: {article_text}", return_tensors="pt", truncation=True)
+        outputs = title_model.generate(
+            input_ids,
+            max_length=max_length,
+            num_beams=num_beams,
+            early_stopping=early_stopping
+        )
+        title = title_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return title
+    except Exception as e:
+        return f"Error generating title: {e}"
+
+# 뉴스 기사 요약 생성 함수
+def generate_summary(article_text, max_length=40, num_beams=10, early_stopping=True):
+    try:
+        input_ids = summarizer_tokenizer.encode(article_text, return_tensors="pt", truncation=True)
+        outputs = summarizer_model.generate(
+            input_ids,
+            max_length=max_length,
+            num_beams=num_beams,
+            early_stopping=early_stopping
+        )
+        summary = summarizer_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return summary
+    except Exception as e:
+        return f"Error generating summary: {e}"
+
+# 뉴스 기사 본문 또는 URL을 입력받아 본문 추출 후 제목과 요약 생성하는 함수
+def process_news(input_text, title_max_length, title_num_beams, title_early_stopping,
+                 summary_max_length, summary_num_beams, summary_early_stopping):
+    article_text = extract_naver_news_article(input_text)
+    title = generate_title(article_text, max_length=title_max_length, num_beams=title_num_beams, early_stopping=title_early_stopping)
+    summary = generate_summary(article_text, max_length=summary_max_length, num_beams=summary_num_beams, early_stopping=summary_early_stopping)
+    return article_text, title, summary
+
+# Gradio 인터페이스 실행 함수
+def launch_gradio_interface():
+    interface = gr.Interface(
+        fn=process_news,
+        inputs=[
+            gr.Textbox(label="뉴스 기사 본문 또는 URL", placeholder="뉴스 기사 본문을 직접 입력하거나 URL을 입력하세요.", lines=10),
+            gr.Slider(0, 40, value=20, step=1, label="제목 생성 - 최대 길이"),
+            gr.Slider(1, 10, value=10, step=1, label="제목 생성 - 탐색 빔 수"),
+            gr.Checkbox(value=True, label="제목 생성 - 조기 종료"),
+            gr.Slider(0, 256, value=128, step=1, label="요약 생성 - 최대 길이"),
+            gr.Slider(1, 10, value=10, step=1, label="요약 생성 - 탐색 빔 수"),
+            gr.Checkbox(value=True, label="요약 생성 - 조기 종료")
+        ],
+        outputs=[
+            gr.Textbox(label="추출되거나 입력된 기사 본문"),
+            gr.Textbox(label="생성된 제목"),
+            gr.Textbox(label="생성된 요약문")
+        ],
+        title="네이버 뉴스 제목 및 요약 생성기",
+        description="뉴스 기사 본문 또는 URL을 입력하면, 본문을 자동 추출한 후 제목과 요약문을 생성합니다."
+    )
+    interface.launch(share=True)
+
+if __name__ == '__main__':
+    launch_gradio_interface()