생성형 AI 프로젝트 3단계 - 챗봇 UI 만들기 (Streamlit)

728x90

개요

 레딧 스크랩, AI 모델 학습 등 핵심적인 기능은 모두 구현했고, UI를 챗봇 형태로 만드는 일만 남았다. 이를 위해서 Streamlit을 사용했다.

 

완성된 코드

# 챗봇 형태로 변환

import streamlit as st
import torch
from transformers import BertTokenizer, BertForSequenceClassification
from wordcloud import WordCloud
import matplotlib.pyplot as plt
import pandas as pd
from datetime import datetime, timedelta
import pytz
import praw


##############################################
# 1. 레딧 코멘트 스크랩
##############################################

# Reddit API 인증
reddit = praw.Reddit(
    client_id='본인 클라이언트 ID',
    client_secret='본인 클라이언트 Secret',
    user_agent='본인 계정 닉네임 넝보'
)

# 특정 티커에 대해 Reddit 코멘트를 수집하고, 최근 n일 이내의 것만 반환
def scrape_reddit_comments(ticker, subreddit_name='wallstreetbets', days=3, limit=1000):
    utc = pytz.UTC
    end_date = datetime.now(utc)
    start_date = end_date - timedelta(days=days)

    comments_data = []

    # Subreddit에서 인기 글 검색
    subreddit = reddit.subreddit(subreddit_name)
    for submission in subreddit.search(ticker, sort='new', time_filter='week', limit=limit):
        submission.comments.replace_more(limit=0)  # 더 많은 댓글 로딩
        for comment in submission.comments.list():
            comment_time = datetime.fromtimestamp(comment.created_utc, tz=utc)
            if start_date <= comment_time <= end_date:
                comments_data.append({
                    'comment_id': comment.id, # 댓글 ID
                    'created_utc': comment_time, # 댓글이 작성된 시간
                    'body': comment.body, # 댓글 내용
                    'score': comment.score, # 코멘트 추천수
                    'submission_title': submission.title, # 댓글이 달린 게시글 제목
                    'permalink': f"https://reddit.com{comment.permalink}" # 댓글 링크
                })

    df = pd.DataFrame(comments_data) # 스크랩한 코멘트 담겨 있는 데이터프레임

    # date 컬럼 추가 : '년도-월-일' 형식
    df['date'] = pd.to_datetime(df['created_utc']) # created_utc를 datetime 객체로 변환
    df['date'] = df['date'].dt.tz_convert('Asia/Seoul') # UTC -> KST(Asia/Seoul) 서울로 시간대 변환
    df['date'] = df['date'].dt.date # '년도-월-일' 형식의 날짜 정보만 추출한 새 컬럼

    return df

##############################################
# 2. 긍정, 부정, 중립 분석 모델
##############################################

def load_model_and_tokenizer(model_path):
    model = BertForSequenceClassification.from_pretrained(model_path) # 저장된 모델 불러오기
    tokenizer = BertTokenizer.from_pretrained(model_path) # 저장된 토크나이저 불러오기
    return model, tokenizer

def predict_sentiments(model, tokenizer, df):
    model.eval() # 평가 모드로 전환
    sentiments = []
    with torch.no_grad():
        for text in df["body"].tolist():
            inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=100)
            outputs = model(**inputs)
            preds = torch.argmax(outputs.logits, dim=1).item()
            sentiments.append(preds)

    df["label"] = sentiments # 긍정, 부정, 중립 판별해서 라벨링하기
    label_map = {0: "부정", 1: "중립", 2: "긍정"}
    df["sentiment"] = df["label"].map(label_map) # 라벨링된 판별 해석
    return df

##############################################
# 3. 시각화
##############################################

def show_daily_sentiment_chart(df): # 설정한 기간 동안 매일에 대해 감성 분석
    daily = df.groupby(["date", "sentiment"]).size().unstack().fillna(0) # 각 감정 별로 집계하기 
    st.line_chart(daily)

def show_cumulative_chart(df): # 설정한 기간 도안 누적 감성 분석
    cumulative = df["sentiment"].value_counts() # 각 감정 별로 집계하기
    st.bar_chart(cumulative)

def show_wordcloud(df): # Wordcloud 만들기
    text = " ".join(df["body"]) # 모든 댓글들 하나의 문자열로 만들기
    wordcloud = WordCloud(width=800, height=400, background_color='white').generate(text)
    plt.figure(figsize=(10, 5))
    plt.imshow(wordcloud, interpolation="bilinear")
    plt.axis("off")
    st.pyplot(plt)

def get_top_comment(df): # 추천수가 가장 많은 댓글 찾기
    top = df.sort_values(by="score", ascending=False).iloc[0]
    return top["body"]


##############################################
# 4. streamlit
##############################################

# 페이지 설정
st.set_page_config(page_title='종목 민심 살피기', layout='wide')
st.markdown("💬 **원하는 종목 티커와 기간을 입력하면 Reddit 코멘트를 기반으로 민심 분석을 해드려요!**")

# 상태 저장을 위한 세션 상태 초기화
if "messages" not in st.session_state:
    st.session_state.messages = []
if "last_contexts" not in st.session_state:
    st.session_state.last_contexts = []
if 'ticker_history' not in st.session_state:
    st.session_state.ticker_history = []
if "latest_ticker" not in st.session_state:
    st.session_state.latest_ticker = ""
# 페이지가 새로 렌더링 될 때마다 이전에 조사된 티커 관련 내용들을 출력하기 위해

# 기존 대화 내용 출력
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# 이전 출력된 그래프 띄워놓기
if "last_comments_df" in st.session_state:
    with st.chat_message("assistant"):
        st.markdown(f"📅 {st.session_state.latest_ticker} 분석 결과 - 일자별 민심 비율")
        show_daily_sentiment_chart(st.session_state.last_comments_df)
        st.markdown(f"📊 {st.session_state.latest_ticker} 분석 결과 - 누적 민심 비율")
        show_cumulative_chart(st.session_state.last_comments_df)
        st.markdown(f"💡 {st.session_state.latest_ticker} WordCloud")
        show_wordcloud(st.session_state.last_comments_df)

# 채팅 입력창
user_input = st.chat_input("민심이 궁금한 티커를 검색해보세요. 검색 양식 : 티커, 기간 (NVDA, 3)")

if user_input:
    # 사용자 메시지 기록
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)
    
    try:
        ticker, days = str(user_input).strip().split(',') # 입력받은 값을 티커, 스크랩 기간으로 변경
    except:
        bot_response14 = "올바른 형태로 입력해주세요."
        with st.chat_message("assistant"):
            st.markdown(bot_response14)
        st.session_state.messages.append({"role": "assistant", "content": bot_response14})
    else:
        if ticker and days:
            st.session_state.latest_ticker = str(ticker).upper() # 페이지가 새로 렌더링 될 때마다 이전에 조사된 티커명 출력하기 위해
            st.session_state.ticker_history.append(ticker.upper()) # 입력된 티커 모두 대문자로 변환

            with st.spinner(f"{ticker.upper()}에 대한 코멘트 수집 중..."):
                comments_df = scrape_reddit_comments(ticker, days=int(days)) # 레딧 코멘트 스크랩
                comments_df.to_csv(f"data/{ticker.upper()}_comments.csv", index=False) # 스크랩된 코멘트들 저장

            bot_response1 = f"{ticker.upper()} 관련 총 {len(comments_df)}개의 코멘트를 수집했습니다." # 봇 답변1. 수집 완료 문구
            bot_response2 = f"{ticker.upper()} 의 민심을 분석 중입니다." # 봇 답변2. 감성 분석 중 문구 띄우기
            with st.chat_message("assistant"): # 수집 완료 문구, 감성 분석 중 문구 출력
                st.markdown(bot_response1)
                st.markdown(bot_response2)
            
            # 감성 분석
            DIR_MODEL = 'C:/Users/User/Desktop/code/project/saved_model/final_model' # 불러올 모델 경로
            model, tokenizer = load_model_and_tokenizer(DIR_MODEL) # 저장된 모델, 토크나이저 불러오기
            comments_df = predict_sentiments(model, tokenizer, comments_df) # 판별 결과 컬럼 추가하여 반환
            
            bot_response3 = f"{ticker.upper()} 의 민심이 분석 완료되었습니다. 잠시만 기다리십시오." # 봇 답변3. 감성 분석 완료 문구 띄우기

            with st.chat_message("assistant"): # 감성 분석 완료 문구 출력
                st.markdown(bot_response3)

            # 시각화 및 챗봇 답변 출력
            bot_response4 = f"📅 {ticker.upper()} 일자별 민심 비율" # 봇 답변4. 일자별 민심 비율 문구
            bot_response5 = f"📊 {ticker.upper()} 누적 민심 비율" # 봇 답변5. 누적 민심 비율 문구
            bot_response6 = "💡 단어 Wordcloud" # 봇 답변6. wordcloud 문구
            bot_response7 = "🏆 추천 수 가장 높은 코멘트" # 봇 답변7. 베스트 코멘트 문구
            bot_response8 = get_top_comment(comments_df) # 봇 답변8. 베스트 코멘트
            bot_response9 = "✅ 코멘트 자료 다운로드" # 봇 답변9. 코멘트 자료 다운로드 버튼
            bot_response10 = "🔍 분석이 끝났습니다! 더 궁금한 종목이 있으면 말씀해 주세요 😊" # 봇 답변10. 마무리 안내 메시지

            # 챗봇 대답
            with st.chat_message("assistant"): # 일자별, 누적 민심, 베스트 코멘트 출력
                st.markdown(bot_response4)
                bot_response11 = show_daily_sentiment_chart(comments_df) # 일자별 민석
                st.markdown(bot_response5)
                bot_response12 = show_cumulative_chart(comments_df) # 누적 민심 출력
                st.markdown(bot_response6)
                bot_response13 = show_wordcloud(comments_df) # wordcloud 출력
                st.markdown(bot_response7)
                st.markdown(bot_response8)
                st.markdown(bot_response9)

            st.session_state.last_comments_df = comments_df # 새로운 텍스트 입력 후에도 그래프 사라지지 않게 하기 위해

            csv = comments_df.to_csv(index=False).encode('utf-8-sig') # 다운로드 위해 csv로 변환
            st.download_button(
                label = "코멘트 파일 다운로드",
                data = csv,
                file_name = f"{ticker}_comments.csv",
                mime="text/csv"
            )

모듈1. 레딧 코멘트 스크랩

 첫 번째 단계에서 작성한 코드이다. 레딧 댓글들을 데이터프레임에 저장하여 반환한다. 나중에 챗봇에서 다운로드할 수 있도록 하기 위해서다.

 

모듈2. 감성 분석 모델 불러오기

 두 번째 단계에서 학습시킨 모델을 불러오는 함수(load_model_and_tokenizer), 수집한 댓글들을 감성 분석하는 함수(predict_sentiments)를 작성했다.

 

모듈3. 시각화 함수들 작성

 나는 최종적으로 네 가지 결과물을 얻고 싶었고, 각각에 대해 아래와 같이 함수를 만들었다.

 

1. 설정한 기간 동안 특정 종목에 대해 매일의 감성 분석 -> show_daily_sentiment_chart 함수

2. 설정한 기간 동안 특정 종목의 누적 감성 분석 (모든 날들의 감성 분석 결과 합치기) -> show_cumulattive_chart

3. 댓글들 wordcloud 만들기 -> show_wordcloud

4. 가장 추천수를 많이 받은 댓글 -> get_top_comment

 

모듈4. Streamlit으로 챗봇 형태 구성하기

 이번 생성형 AI 과정에서 Streamlit으로 챗봇 형태의 UI를 제작하는 방법을 배웠고, 적용했다. 한 가지 아쉬운 점은, 내가 새로운 질문을 입력할 때마다 페이지가 새로 렌더링 된다는 것. 그렇기 때문에 매번 이전의 그래프들과 대화들을 띄워주는 코드를 추가해야 했다. (# 기존 대화 내용 출력, # 이전 출력된 그래프 띄워놓기)

 완성된 결과는 아래와 같다. (NVDA, 3) 이런 식으로 입력하면 3일 전부터 오늘까지 레딧 wallstreetbets에서 NVDA를 검색했을 때 나오는 게시물들의 댓글들을 스크랩한 뒤, 집계하여 시각화해 준다.

 

 

 

참고 자료

1. streamlit session state 설명글 : https://velog.io/@jomminii/streamlit-session-state

[Streamlit] 변수를 기억하고 싶다면 Session State를 사용하십시다.

 

2. streamlit button 설명글 :

https://velog.io/@wonjun12/Streamlit-%EB%B2%84%ED%8A%BC-%EA%B8%B0%EB%8A%A5

[Streamlit] 버튼 기능