app.py

import os
import sys
import gradio as gr
from transformers import BartForConditionalGeneration, BartTokenizer, MBartForConditionalGeneration, MBart50TokenizerFast
from newspaper import Article
import requests
from urllib.parse import urlparse

# 전역 변수
trans_model = None
trans_tokenizer = None
summ_model = None
summ_tokenizer = None

def load_models():
    global trans_model, trans_tokenizer, summ_model, summ_tokenizer
    
    if trans_model is None:
        try:
            print("MBart 다국어 번역 모델 로딩 시도 중...")
            # 영어 → 한국어 번역 모델 로드 (MBart 다국어 모델 사용)
            try:
                trans_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50")
                trans_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50")
                print("번역 모델 로드 성공")
            except Exception as e:
                print(f"번역 모델 로드 오류: {e}")
                return f"번역 모델 로드 오류: {str(e)}"
            
            print("요약 모델 로딩 시도 중...")
            # 한국어 요약 모델 로드
            try:
                summ_tokenizer = BartTokenizer.from_pretrained("digit82/bart-korean-summarization")
                summ_model = BartForConditionalGeneration.from_pretrained("digit82/bart-korean-summarization")
                print("요약 모델 로드 성공")
            except Exception as e:
                print(f"요약 모델 로드 오류: {e}")
                return f"요약 모델 로드 오류: {str(e)}"
            
            return "모델 로드 완료"
        except Exception as e:
            print(f"모델 로드 중 예상치 못한 오류: {e}")
            return f"모델 로드 중 오류 발생: {str(e)}"

def is_valid_url(url):
    """URL 유효성 검사"""
    try:
        result = urlparse(url)
        return all([result.scheme, result.netloc])
    except:
        return False

def validate_url_accessibility(url):
    """URL 접근 가능 여부 확인"""
    try:
        response = requests.head(url, timeout=5)
        return 200 <= response.status_code < 400
    except:
        return False

def translate_text(text):
    """MBart 모델을 사용한 영어→한국어 번역"""
    # 영어로 소스 언어 설정
    trans_tokenizer.src_lang = "en_XX"
    
    # 인코딩
    encoded = trans_tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
    
    # 번역 생성
    generated_tokens = trans_model.generate(
        **encoded,
        forced_bos_token_id=trans_tokenizer.lang_code_to_id["ko_KR"],
        max_length=1024,
        num_beams=5,
        early_stopping=True
    )
    
    # 디코딩
    translated_text = trans_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
    return translated_text

def summarize_news(url):
    """뉴스 기사 요약 함수"""
    # 입력 검증
    if not url or not isinstance(url, str):
        return "오류: 유효한 URL을 입력해주세요."
    
    if not is_valid_url(url):
        return "오류: 입력한 텍스트가 유효한 URL 형식이 아닙니다."
    
    if not validate_url_accessibility(url):
        return "오류: 입력한 URL에 접근할 수 없습니다. URL이 올바른지 확인해주세요."
    
    # 모델 로드 (아직 로드되지 않은 경우)
    if trans_model is None:
        load_result = load_models()
        if isinstance(load_result, str) and "오류" in load_result:
            return load_result
    
    try:
        # 뉴스 기사 크롤링
        article = Article(url)
        article.download()
        article.parse()
        text = article.text
        
        if not text:
            return "오류: 기사에서 텍스트를 추출할 수 없습니다."
        
        # 긴 텍스트 처리
        if len(text) > 5000:
            print(f"텍스트가 너무 깁니다. 원본 길이: {len(text)}자, 5000자로 제한합니다.")
            text = text[:5000]
        
        print("텍스트 번역 중...")
        # 영어 → 한국어 번역 (MBart 모델 사용)
        translated_text = translate_text(text)
        
        print("번역된 텍스트 요약 중...")
        # 한국어 텍스트 요약
        inputs = summ_tokenizer([translated_text], return_tensors="pt", max_length=1024, truncation=True)
        summary_ids = summ_model.generate(
            inputs["input_ids"],
            num_beams=4,
            max_length=200,
            min_length=50,
            early_stopping=True
        )
        summary = summ_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        
        # 결과 반환 (기사 제목 포함)
        result = f"제목: {article.title}\n\n요약: {summary}"
        return result
        
    except Exception as e:
        print(f"처리 중 오류 발생: {e}")
        return f"처리 중 오류가 발생했습니다: {str(e)}"

# Gradio UI
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("# 영어 뉴스 요약기 (EN → KO → 요약)")
    gr.Markdown("### 첫 실행 시 모델을 다운로드하므로 시간이 소요될 수 있습니다.")
    
    with gr.Row():
        with gr.Column():
            input_url = gr.Textbox(
                label="뉴스 기사 URL 입력", 
                placeholder="https://example.com/news-article"
            )
            submit_btn = gr.Button("요약하기", variant="primary")
            model_load_btn = gr.Button("모델 미리 로드하기")
        
    with gr.Row():
        output_text = gr.Textbox(
            label="요약 결과",
            lines=10
        )
    
    gr.Markdown("""
    ### 사용 방법
    1. '모델 미리 로드하기' 버튼을 클릭하여 모델을 미리 로드할 수 있습니다.
    2. 영어 뉴스 기사 URL을 입력합니다.
    3. '요약하기' 버튼을 클릭합니다.
    4. 한국어로 번역 및 요약된 결과가 표시됩니다.
    
    ### 참고 사항
    - 처리 시간은 기사 길이에 따라 달라질 수 있습니다.
    - 모델이 처음 로드될 때 다소 시간이 소요될 수 있습니다.
    - 5,000자가 넘는 긴 기사는 일부만 처리됩니다.
    """)
    
    submit_btn.click(
        fn=summarize_news,
        inputs=input_url,
        outputs=output_text
    )
    
    model_load_btn.click(
        fn=load_models,
        inputs=None,
        outputs=output_text
    )

if __name__ == "__main__":
    demo.launch()