DEVOCEAN

import streamlit as st
from pathlib import Path
from langchain_community.chat_models import ChatOllama
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma, FAISS
from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough
from langchain.retrievers import BM25Retriever, EnsembleRetriever

class RagEnsemble:
    def __init__(self):
        """GemmaRAG 클래스의 초기화 메서드로 필요한 객체들을 초기화합니다."""
        self.model = ChatOllama(model="gemma:7b-instruct", temperature=0)
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1024,
            chunk_overlap=100
        )
        self.retriever = None
        self.rag_chain = None

    def format_docs(self, docs):
        """문서 리스트를 하나의 문자열로 포맷팅합니다."""
        return '\n\n'.join([d.page_content for d in docs])

    def setup_pdf_retriever(self, pdf_file_path: str):
        """PDF 파일을 읽어들여 리트리버를 설정합니다."""
        # PDF 파일 로드
        pages = PyMuPDFLoader(file_path=pdf_file_path).load()
        # 문서 분할
        docs = self.text_splitter.split_documents(pages)

        # 문장을 임베딩으로 변환하고 벡터 저장소에 저장
        embeddings = HuggingFaceEmbeddings(
            model_name='jhgan/ko-sroberta-nli',
            model_kwargs={'device':'cpu'},
            encode_kwargs={'normalize_embeddings':True}
        )

        # FAISS 벡터 스토어 생성
        vectorstore_faiss = FAISS.from_documents(docs, embeddings)
        # BM25 리트리버 생성
        bm25_retriever = BM25Retriever.from_documents(docs)
        # FAISS 리트리버 생성
        faiss_retriever = vectorstore_faiss.as_retriever(search_kwargs={'k': 5})

        # 앙상블 리트리버 생성
        self.retriever = EnsembleRetriever(
            retrievers=[bm25_retriever, faiss_retriever],
            weights=[0.5, 0.5]  # 가중치 설정
        )

        # 답변 생성을 위한 프롬프트 템플릿 설정
        template = '''Answer the question based only on the following context:
        {context}

        Question: {question}
        '''
        self.prompt = ChatPromptTemplate.from_template(template)
        self.rag_chain = (
            {"context": self.retriever | self.format_docs, "question": RunnablePassthrough()}
            | self.prompt
            | self.model
            | StrOutputParser()
        )

    def ask(self, query: str):
        """질문을 받아 답변을 생성합니다."""
        if not self.rag_chain:
            return "PDF 파일을 먼저 설정해주세요."
        return self.rag_chain.invoke(query)

    def clear(self):
        """설정을 초기화합니다."""
        self.retriever = None
        self.rag_chain = None

# Streamlit 앱 구성
def main():
    """Streamlit 앱의 메인 함수로, PDF Q&A 챗봇을 구성합니다."""
    st.title("PDF Q&A 챗봇")
    st.write("PDF 파일에서 정보를 추출하여 질문에 답변하는 챗봇입니다.")

    # GemmaRAG 클래스 인스턴스 생성
    gemma_rag = RagEnsemble()
    # PDF 파일 경로 설정
    pdf_file_path = '경남교육청_수목관리_업무매뉴얼.pdf'
    # PDF 리트리버 설정
    gemma_rag.setup_pdf_retriever(pdf_file_path)

    # 사용자 입력 받기
    user_input = st.text_input("질문을 입력하세요:")

    if st.button("질문하기"):
        if user_input:
            # 질문에 대한 답변 생성
            answer = gemma_rag.ask(user_input)
            st.write("답변: ", answer)
        else:
            st.write("질문을 입력하세요.")

if __name__ == "__main__":
    main()

import streamlit as st

from pathlib import Path

from langchain_community.chat_models import ChatOllama

from langchain_community.document_loaders import PyMuPDFLoader

from langchain.text_splitter import RecursiveCharacterTextSplitter

from langchain_community.embeddings import HuggingFaceEmbeddings

from langchain_community.vectorstores import FAISS

from langchain.prompts import ChatPromptTemplate

from langchain.schema.output_parser import StrOutputParser

from langchain.schema.runnable import RunnablePassthrough

from langchain.retrievers import BM25Retriever, EnsembleRetriever

import torch

  

class RagEnsembleRerank:

def __init__(self):

"""초기화 메서드로 필요한 객체들을 초기화합니다."""

self.model = ChatOllama(model="gemma:7b-instruct", temperature=0)

self.text_splitter = RecursiveCharacterTextSplitter(

chunk_size=1024,

chunk_overlap=100

)

self.retriever = None

self.rag_chain = None

self.embeddings = HuggingFaceEmbeddings(

model_name='jhgan/ko-sroberta-nli',

model_kwargs={'device': 'cpu'},

encode_kwargs={'normalize_embeddings': True}

)

  

def format_docs(self, docs):

"""문서 리스트를 하나의 문자열로 포맷팅합니다."""

return '\n\n'.join([d.page_content for d in docs])

  

def rerank(self, query, docs):

"""질문과 문서 리스트를 받아 유사도 기반으로 문서를 재정렬합니다."""

query_embedding = self.embeddings.embed_query(query) # 질문 임베딩 생성

doc_texts = [doc.page_content for doc in docs] # 문서 텍스트 추출

doc_embeddings = [self.embeddings.embed_query(doc) for doc in doc_texts] # 문서 임베딩 생성

scores = [torch.cosine_similarity(torch.tensor(query_embedding), torch.tensor(doc_embedding), dim=0).item()

for doc_embedding in doc_embeddings] # 코사인 유사도 계산

sorted_docs = [doc for _, doc in sorted(zip(scores, docs), key=lambda x: x[0], reverse=True)] # 유사도 기준 정렬

return sorted_docs

  

def setup_pdf_retriever(self, pdf_file_path: str, bm25_weight: float, faiss_weight: float):

"""PDF 파일을 읽어들여 리트리버를 설정합니다."""

pages = PyMuPDFLoader(file_path=pdf_file_path).load() # PDF 파일 로드

docs = self.text_splitter.split_documents(pages) # 문서 분할

  

vectorstore_faiss = FAISS.from_documents(docs, self.embeddings) # FAISS 벡터 스토어 생성

bm25_retriever = BM25Retriever.from_documents(docs) # BM25 리트리버 생성

faiss_retriever = vectorstore_faiss.as_retriever(search_kwargs={'k': 5}) # FAISS 리트리버 생성

  

self.retriever = EnsembleRetriever(

retrievers=[bm25_retriever, faiss_retriever],

weights=[bm25_weight, faiss_weight] # 가중치 설정

)

  

# 답변 생성을 위한 프롬프트 템플릿 설정

template = '''Answer the question based only on the following context:

{context}

  

Question: {question}

'''

self.prompt = ChatPromptTemplate.from_template(template)

self.rag_chain = (

self.prompt | self.model | StrOutputParser()

)

  

def ask(self, query: str):

"""질문을 받아 답변을 생성합니다."""

if not self.rag_chain:

return "PDF 파일을 먼저 설정해주세요."

result = self.retriever.get_relevant_documents(query) # 관련 문서 검색

reranked_docs = self.rerank(query, result) # 문서 재정렬

formatted_docs = self.format_docs(reranked_docs) # 문서 포맷팅

response = self.rag_chain.invoke({"context": formatted_docs, "question": query}) # 답변 생성

return response

  

def clear(self):

"""설정을 초기화합니다."""

self.retriever = None

self.rag_chain = None

  

# Streamlit 앱 구성

def main():

# 페이지 설정

st.set_page_config(page_title="PDF Q&A Chatbot", page_icon=":books:")

  

# CSS 스타일링

st.markdown(

"""

<style>

body {

background-color: #f0f2f6;

color: #333;

}

.chat-container {

display: flex;

flex-direction: column;

align-items: flex-start;

margin-bottom: 100px;

}

.user-message, .bot-message {

padding: 10px 15px;

border-radius: 10px;

margin: 5px 0;

max-width: 80%;

}

.user-message {

background-color: #DCF8C6;

align-self: flex-end;

}

.bot-message {

background-color: #E8E8E8;

}

.input-container {

position: fixed;

bottom: 0;

width: 100%;

background-color: white;

padding: 10px;

box-shadow: 0 -2px 5px rgba(0, 0, 0, 0.1);

}

.slider-container {

display: flex;

justify-content: space-between;

margin-bottom: 20px;

}

</style>

""", unsafe_allow_html=True

)

  

# 타이틀과 설명

st.title("PDF Q&A 챗봇 :books:")

st.write("PDF 파일에서 정보를 추출하여 질문에 답변하는 챗봇입니다.")

  

# 가중치 슬라이더 추가

st.markdown('<div class="slider-container">', unsafe_allow_html=True)

bm25_weight = st.slider("BM25 Weight", 0.0, 1.0, 0.5)

faiss_weight = st.slider("FAISS Weight", 0.0, 1.0, 0.5)

st.markdown('</div>', unsafe_allow_html=True)

  

if bm25_weight + faiss_weight != 1.0:

st.error("BM25 Weight와 FAISS Weight의 합은 1.0이어야 합니다.")

return

  

gemma_rag = RagEnsembleRerank()

pdf_file_path = '경남교육청_수목관리_업무매뉴얼.pdf'

gemma_rag.setup_pdf_retriever(pdf_file_path, bm25_weight, faiss_weight)

  

# 대화 기록 초기화

if 'chat_history' not in st.session_state:

st.session_state.chat_history = []

  

# 대화 상자

chat_placeholder = st.empty()

for message in st.session_state.chat_history:

if message['role'] == 'user':

chat_placeholder.markdown(

f"<div class='chat-container'><div class='user-message'>{message['content']}</div></div>",

unsafe_allow_html=True)

else:

chat_placeholder.markdown(

f"<div class='chat-container'><div class='bot-message'>{message['content']}</div></div>",

unsafe_allow_html=True)

  

# 사용자 입력

user_input = st.text_input("질문을 입력하세요:", key="input", label_visibility="hidden")

  

# 입력 컨테이너

st.markdown('<div class="input-container">', unsafe_allow_html=True)

if st.button("질문하기", key="button"):

if user_input:

# 사용자 질문 기록

st.session_state.chat_history.append({"role": "user", "content": user_input})

answer = gemma_rag.ask(user_input)

if isinstance(answer, str):

# 챗봇 응답 기록

st.session_state.chat_history.append({"role": "bot", "content": answer})

chat_placeholder.markdown(

f"<div class='chat-container'><div class='user-message'>{user_input}</div><div class='bot-message'>{answer}</div></div>",

unsafe_allow_html=True)

else:

st.write("오류가 발생했습니다. 다시 시도해주세요.")

else:

st.write("질문을 입력하세요.")

st.markdown('</div>', unsafe_allow_html=True)

  

if __name__ == "__main__":

main()