app.py

import streamlit as st
import tiktoken
import re
from loguru import logger

from langchain.chains import ConversationalRetrievalChain
from langchain_community.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

from langchain.document_loaders import PyPDFLoader
from langchain.document_loaders import Docx2txtLoader

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings

from langchain.memory import ConversationBufferMemory
from langchain.vectorstores import FAISS
from langchain.memory import StreamlitChatMessageHistory

def preprocess_korean_text(text):
    """한국어 텍스트 전처리 함수"""
    # 불필요한 특수문자 제거 (한국어, 영어, 숫자, 공백만 유지)
    text = re.sub(r'[^가-힣a-zA-Z0-9\s.,!?]', ' ', text)
    # 연속된 공백을 하나로 통합
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def main():
    st.set_page_config(
        page_title="한국어 문서 QA 챗봇",
        page_icon="🇰🇷",
        layout="wide"
    )

    st.title("🇰🇷 _한국어 전용 문서 :red[QA 챗봇]_ 📚")
    st.markdown("**최고 성능의 한국어 AI 모델로 구동되는 문서 질의응답 시스템**")

    if "conversation" not in st.session_state:
        st.session_state.conversation = None

    if "chat_history" not in st.session_state:
        st.session_state.chat_history = None

    if "processComplete" not in st.session_state:
        st.session_state.processComplete = None

    with st.sidebar:
        st.header("⚙️ 설정")
        
        uploaded_files = st.file_uploader(
            "📁 한국어 문서 업로드", 
            type=['pdf','docx'], 
            accept_multiple_files=True,
            help="PDF, DOCX 형식의 한국어 문서를 업로드하세요."
        )
        
        st.subheader("🤖 AI 모델 선택")
        # 최고 성능 한국어 모델들로 교체
        model_options = {
            "🥇 EEVE-Korean-10.8B (최고 성능)": "yanolja/EEVE-Korean-Instruct-10.8B-v1.0",
            "🥈 Llama3-Korean-Bllossom-8B": "MLP-KTLim/llama-3-Korean-Bllossom-8B", 
            "🥉 KoAlpaca-Polyglot-12.8B": "beomi/KoAlpaca-Polyglot-12.8B",
            "⚡ Kullm-Polyglot-5.8B (빠름)": "nlpai-lab/kullm-polyglot-5.8b-v2",
            "💎 Korean-Vicuna-13B": "kfkas/Llama-2-ko-7b-Chat"
        }
        
        selected_model_name = st.selectbox(
            "모델 선택:", 
            list(model_options.keys()),
            help="EEVE 모델이 한국어 지시사항 이해에 가장 뛰어납니다."
        )
        selected_model = model_options[selected_model_name]
        
        st.subheader("📊 임베딩 모델")
        embedding_options = {
            "🇰🇷 KoSBERT (추천)": "jhgan/ko-sroberta-multitask",
            "🔥 KoSimCSE": "BM-K/KoSimCSE-roberta-multitask", 
            "⭐ KR-SBERT": "snunlp/KR-SBERT-V40K-klueNLI-augSTS"
        }
        
        selected_embedding_name = st.selectbox(
            "임베딩 모델:",
            list(embedding_options.keys())
        )
        selected_embedding = embedding_options[selected_embedding_name]
        
        st.subheader("⚙️ 고급 설정")
        chunk_size = st.slider("청크 크기", 200, 1000, 400, help="한국어는 400-600자가 최적입니다.")
        chunk_overlap = st.slider("청크 겹침", 20, 200, 40, help="겹침이 클수록 문맥 연결성이 향상됩니다.")
        temperature = st.slider("창의성 (Temperature)", 0.1, 1.0, 0.3, help="낮을수록 정확, 높을수록 창의적")
        
        process = st.button("🚀 문서 처리 시작", type="primary")
    
    if process:
        if uploaded_files:
            with st.spinner("🔥 최고 성능 한국어 AI로 문서를 분석 중입니다..."):
                try:
                    files_text = get_text(uploaded_files)
                    text_chunks = get_text_chunks(files_text, chunk_size, chunk_overlap)
                    vectorstore = get_vectorstore(text_chunks, selected_embedding)
                    st.session_state.conversation = get_conversation_chain(vectorstore, selected_model, temperature)
                    st.session_state.processComplete = True
                    st.success(f"✅ {len(files_text)}개 문서, {len(text_chunks)}개 청크로 처리 완료!")
                    st.balloons()
                except Exception as e:
                    st.error(f"❌ 문서 처리 중 오류 발생: {str(e)}")
        else:
            st.error("📁 파일을 먼저 업로드해주세요!")

    if 'messages' not in st.session_state:
        st.session_state['messages'] = [{
            "role": "assistant", 
            "content": "안녕하세요! 🇰🇷 **한국어 전용 고성능 AI 챗봇**입니다.\n\n📚 **특징:**\n- 최신 한국어 특화 AI 모델 사용\n- 복잡한 지시사항 완벽 이해\n- 정확하고 자연스러운 한국어 답변\n\n📁 문서를 업로드하고 '문서 처리 시작'을 눌러주세요!"
        }]

    # 채팅 인터페이스
    st.subheader("💬 대화")
    
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if query := st.chat_input("🤔 문서에 대해 무엇이든 물어보세요... (복잡한 질문도 환영!)"):
        if st.session_state.conversation is None:
            st.error("먼저 파일을 업로드하고 '문서 처리 시작' 버튼을 눌러주세요!")
            st.stop()
            
        st.session_state.messages.append({"role": "user", "content": query})

        with st.chat_message("user"):
            st.markdown(query)

        with st.chat_message("assistant"):
            with st.spinner("🧠 한국어 AI가 깊이 분석하고 있습니다..."):
                try:
                    # 한국어 프롬프트 최적화
                    enhanced_query = f"다음 질문에 대해 문서 내용을 바탕으로 정확하고 상세하게 한국어로 답변해주세요: {query}"
                    
                    result = st.session_state.conversation({"question": enhanced_query})
                    response = result['answer']
                    source_documents = result.get('source_documents', [])

                    # 답변 후처리
                    if response:
                        # 불필요한 영어 제거 및 한국어 답변 추출
                        response = clean_korean_response(response)
                        st.markdown(response)
                    else:
                        st.markdown("죄송합니다. 해당 질문에 대한 답변을 문서에서 찾을 수 없습니다.")
                    
                    if source_documents:
                        with st.expander("📖 참고 문서 및 근거"):
                            for i, doc in enumerate(source_documents[:3]):
                                st.markdown(f"**📄 문서 {i+1}:** {doc.metadata.get('source', 'Unknown')}")
                                with st.container():
                                    st.text_area(
                                        f"관련 내용 {i+1}",
                                        doc.page_content[:400] + "...",
                                        height=120,
                                        disabled=True
                                    )
                    
                    st.session_state.messages.append({"role": "assistant", "content": response})
                except Exception as e:
                    error_msg = f"❌ 답변 생성 중 오류가 발생했습니다: {str(e)}"
                    st.error(error_msg)
                    st.session_state.messages.append({"role": "assistant", "content": "죄송합니다. 일시적인 오류가 발생했습니다. 다시 시도해주세요."})

def clean_korean_response(response):
    """한국어 답변 정제"""
    # 영어 패턴 제거
    response = re.sub(r'\b[A-Za-z]+\b', '', response)
    # 불필요한 기호 정리
    response = re.sub(r'[\[\]\(\)\{\}]', '', response)
    # 연속 공백 정리
    response = re.sub(r'\s+', ' ', response).strip()
    return response

def get_text(docs):
    """문서에서 텍스트 추출 및 전처리"""
    doc_list = []
    
    for doc in docs:
        file_name = doc.name
        with open(file_name, "wb") as file:
            file.write(doc.getvalue())
            logger.info(f"Uploaded {file_name}")
            
        try:
            if '.pdf' in doc.name:
                loader = PyPDFLoader(file_name)
                documents = loader.load_and_split()
            elif '.docx' in doc.name:
                loader = Docx2txtLoader(file_name)
                documents = loader.load_and_split()
            
            # 각 문서의 텍스트 전처리
            for document in documents:
                document.page_content = preprocess_korean_text(document.page_content)
                # 너무 짧은 청크 제거
                if len(document.page_content.strip()) < 50:
                    continue
            
            doc_list.extend([doc for doc in documents if len(doc.page_content.strip()) >= 50])
        except Exception as e:
            st.error(f"파일 {file_name} 처리 중 오류: {str(e)}")
            
    return doc_list

def get_text_chunks(text, chunk_size=400, chunk_overlap=40):
    """한국어 최적화된 텍스트 청킹"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=["\n\n", "\n", ".", "!", "?", ";", ":", ",", " ", ""]  # 한국어 구분자 최적화
    )
    chunks = text_splitter.split_documents(text)
    return chunks

def get_vectorstore(text_chunks, embedding_model):
    """한국어 특화 임베딩 모델을 사용한 벡터 스토어 생성"""
    embeddings = HuggingFaceEmbeddings(
        model_name=embedding_model,
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True}
    )
    vectordb = FAISS.from_documents(text_chunks, embeddings)
    return vectordb

def get_conversation_chain(vectorstore, model_name, temperature):
    """한국어 특화 대화 체인 생성"""
    try:
        # 한국어 특화 토크나이저 및 모델 로딩
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        
        # 패딩 토큰 설정
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            
        model = AutoModelForCausalLM.from_pretrained(
            model_name, 
            trust_remote_code=True,
            torch_dtype="auto",
            device_map=None  # GPU 사용 설정 제거
        )
        
        # 한국어 최적화 파이프라인
        pipe = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            max_new_tokens=512,
            temperature=temperature,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.1,
            device=-1,  # CPU 사용
            pad_token_id=tokenizer.eos_token_id
        )
        
        llm = HuggingFacePipeline(pipeline=pipe)
        
        # 한국어 특화 검색 설정
        conversation_chain = ConversationalRetrievalChain.from_llm(
            llm=llm,
            chain_type="stuff",
            retriever=vectorstore.as_retriever(
                search_type='mmr',
                search_kwargs={
                    'k': 4,  # 더 많은 문서 검색
                    'fetch_k': 8,
                    'lambda_mult': 0.7  # 다양성과 관련성 균형
                }
            ),
            memory=ConversationBufferMemory(
                memory_key='chat_history', 
                return_messages=True, 
                output_key='answer'
            ),
            return_source_documents=True,
            verbose=True
        )
        
        return conversation_chain
    except Exception as e:
        st.error(f"모델 로딩 중 오류: {str(e)}")
        st.info("더 가벼운 모델을 선택하거나 메모리를 확인해주세요.")
        return None

if __name__ == '__main__':
    main()