투자 어시스턴트 개발

Streamlit과 Langchain을 활용한 실시간 온프레미스 투자 어시스턴트 구현하기 [19]

HeyTeddy 2024. 12. 25. 12:00
반응형

“투자 어시스턴트” 웹 애플리케이션을 Streamlit과 Hugging Face Transformers, Langchain 등의 라이브러리를 사용해 구축하는 과정을 보여줍니다. 실시간 질의에 대해 온프레미스(로컬) 환경에서 답변을 생성할 수 있도록 합니다.

import streamlit as st
import transformers
import torch

# Langchain 및 커뮤니티 모듈(베타)에서 BM25, FAISS, SentenceTransformer를 불러옵니다.
from langchain_community.retrievers import BM25Retriever
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.retrievers import EnsembleRetriever

# 예시 데이터
data = [
    {
        "기업명": "삼성전자",
        "날짜": "2024-03-02",
        "문서 카테고리": "인수합병",
        "요약": "삼성전자가 HVAC(냉난방공조) 사업 인수를 타진 중이며, 이는 기존 가전 사업의 약점 보완을 목적으로 한다.",
        "주요 이벤트": ["기업 인수합병"]
    },
    {
        "기업명": "삼성전자",
        "날짜": "2024-03-24",
        "문서 카테고리": "인수합병",
        "요약": "테스트 하나 둘 셋",
        "주요 이벤트": ["신제품 출시"]
    },
    {
        "기업명": "현대차",
        "날짜": "2024-04-02",
        "문서 카테고리": "인수합병",
        "요약": "삼성전자가 HVAC(냉난방공조) 사업 인수를 타진 중이며, 이는 기존 가전 사업의 약점 보완을 목적으로 한다.",
        "주요 이벤트": ["기업 인수합병", "신제품 출시"]
    },
]

# 요약 부분만 추출해 리스트화(문서 텍스트로 사용)
doc_list = [item['요약'] for item in data]

# BM25 Retriever 설정
bm25_retriever = BM25Retriever.from_texts(
    doc_list,
    metadatas=[{"source": i} for i in range(len(data))]
)
# 검색 결과 개수 설정
bm25_retriever.k = 1

# SentenceTransformer Embeddings & FAISS Vector DB
embedding = SentenceTransformerEmbeddings(model_name="distiluse-base-multilingual-cased-v1")

faiss_vectorstore = FAISS.from_texts(
    doc_list,
    embedding,
    metadatas=[{"source": i} for i in range(len(data))]
)
# FAISS 기반 검색 설정
faiss_retriever = faiss_vectorstore.as_retriever(search_kwargs={"k": 1})

# 앙상블 리트리버(EnsembleRetriever) 구성
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, faiss_retriever],
    weights=[0.5, 0.5]  # 두 검색 방식을 동일 가중치로 조합
)

# 로컬 모델 로드(42dot_LLM-SFT-1.3B)
model_id = "42dot/42dot_LLM-SFT-1.3B"

pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.float16}
)

# 모델 eval 모드로 전환
pipeline.model.eval()

# 검색 함수 정의
def search(query):
    """사용자의 질의(query)를 앙상블 리트리버에 전달하여 관련 문서들 검색"""
    ensemble_docs = ensemble_retriever.invoke(query)
    return ensemble_docs

# 모델을 사용해 텍스트 생성
def sllm_generate(query):
    """모델에 Prompt(query)를 넣어 텍스트 생성"""
    answer = pipeline(
        query,
        max_new_tokens=50,
        do_sample=True,
        temperature=0.3,
        top_p=0.9,
        repetition_penalty=1.2,
    )
    # query 길이 이후의 모델 생성 결과만 반환
    return answer[0]['generated_text'][len(query):]

# 최종 답변 함수
def prompt_and_generate(query):
    """검색된 문서를 바탕으로 최종 Prompt 구성 후 모델에 텍스트 생성 요청"""
    docs = [doc for doc in search(query)]
    prompt = f\"\"\"아래 질문을 기반으로 검색된 뉴스를 참고하여 질문에 대한 답변을 생성하시오.

질문: {query}
\"\"\"
    for i, d in enumerate(docs):
        prompt += f"뉴스{i+1}\n"
        prompt += f"요약: {d.page_content}\n\n"

    prompt += "답변: "

    # 최종 응답 생성
    answer = sllm_generate(prompt)
    return answer

# Streamlit UI 구성
st.title("투자 어시스턴트")

# 대화 히스토리 저장(세션 스테이트 활용)
if "messages" not in st.session_state:
    st.session_state.messages = []

# 기존 메시지 출력
for message in st.session_state.messages:
    with st.chat_message(message['role']):
        st.markdown(message['content'])

# 사용자 입력창
if prompt := st.chat_input("궁금한 점을 물어보세요."):
    # 사용자 메시지를 대화 히스토리에 저장
    st.chat_message("user").markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    # 모델 응답 생성
    response_text = prompt_and_generate(prompt.strip())
    response = f"bot: {response_text}"

    # 채팅창에 답변 표시
    with st.chat_message("assistant"):
        st.markdown(response)
    st.session_state.messages.append({"role": "assistant", "content": response})

터미널에서 streamlit run 16_streamlit.py명령어로 실행하면 웹 인터페이스가 열립니다.

사용하는 라이브러리, 기술, 스택

  1. Python
  2. Streamlit: 빠르게 웹 앱을 만들기 위한 라이브러리
  3. Hugging Face Transformers: 사전 학습된 모델(LLM) 로드를 위한 라이브러리
  4. Torch (PyTorch): 모델 추론을 위한 프레임워크
  5. LangchainLangchain Community Modules:
    • BM25Retriever: 전통적 정보 검색 알고리즘(BM25)을 활용한 텍스트 검색
    • SentenceTransformerEmbeddings: 임베딩 모델로 텍스트 벡터화
    • FAISS: Facebook AI Research에서 개발한 벡터 검색 라이브러리
    • EnsembleRetriever: 여러 검색 방법을 조합해서 결과를 반환

간단한 프로젝트 설명

이 프로젝트는 투자 정보를 요약한 문서들을 기반으로 사용자의 질의에 맞춰 가장 관련성이 높은 문서를 검색하고, 해당 결과를 바탕으로 LLM(Local Large Language Model)실시간 응답을 만들어내는 애플리케이션입니다.

  • 데이터를 BM25와 FAISS를 이용해 검색한 뒤, 두 결과물을 앙상블하여 최적의 검색 결과를 도출합니다.
  • 검색된 요약 문서를 이용해 Prompt를 생성하고, 로컬에서 구동되는 LLM에 입력하여 텍스트 답변을 생성합니다.
  • Streamlit을 통해 간단한 UI/UX를 구성하여 채팅 형태의 인터페이스로 질의응답이 가능합니다.

요약

  • BM25 + FAISS: 전통적 정보 검색 알고리즘과 임베딩 벡터 기반 검색을 함께 활용하여 문서 검색 정확도를 높임
  • 앙상블 리트리버: 서로 다른 검색 방법을 결합해 더 높은 검색 성능을 기대할 수 있음
  • 로컬 LLM 추론: Hugging Face 모델(42dot_LLM-SFT-1.3B)을 직접 로드하여 서버에 부담을 최소화하고 데이터 유출 없이 자체 모델 추론 가능
  • Streamlit UI: 누구나 쉽게 접근 가능한 웹 앱 형태로 구현

위 과정을 통해 투자 어시스턴트 웹 애플리케이션을 제작하고, 검색된 문서를 기반으로 한 맞춤형 응답을 제공할 수 있습니다. Streamlit으로 간단한 데모 페이지를 제작하여 온프레미스 환경에서 LLM을 가볍게 구동해볼 수 있습니다.

 

반응형