import streamlit as st
from transformers import BartForConditionalGeneration, PreTrainedTokenizerFast
import requests
from bs4 import BeautifulSoup

# KoBART 모델 로딩
model = BartForConditionalGeneration.from_pretrained("digit82/kobart-summarization")
tokenizer = PreTrainedTokenizerFast.from_pretrained("digit82/kobart-summarization")

# 뉴스 본문 추출 함수
def get_article_text(naver_url):
    headers = {"User-Agent": "Mozilla/5.0"}
    response = requests.get(naver_url, headers=headers)
    soup = BeautifulSoup(response.text, "html.parser")
    article_body = soup.select_one("article") or soup.select_one("#dic_area")
    return article_body.get_text(strip=True) if article_body else None

# 뉴스 요약 함수
def summarize_korean_news(text):
    input_ids = tokenizer.encode(text, return_tensors="pt", max_length=1024, truncation=True)
    summary_ids = model.generate(input_ids, max_length=128, min_length=30, num_beams=4, early_stopping=True)
    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)

# Streamlit 앱 UI
st.title("HL만도 뉴스 요약기")
st.write("HL만도 관련 최신 뉴스를 요약해 드립니다!")

# 사용자로부터 URL 입력 받기
url = st.text_input("뉴스 URL을 입력하세요:")

if url:
    article = get_article_text(url)
    if article:
        summary = summarize_korean_news(article)
        st.subheader("뉴스 요약")
        st.write(summary)
    else:
        st.write("기사 본문을 가져올 수 없습니다.")