urbanManul/DagTask2
0
1import streamlit as st2from dotenv import load_dotenv3from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter4from langchain.vectorstores import FAISS5from langchain.embeddings import HuggingFaceEmbeddings # General embeddings from HuggingFace models.6from langchain.memory import ConversationBufferMemory7from langchain.chains import ConversationalRetrievalChain8from htmlTemplates import css, bot_template, user_template9from langchain.llms import LlamaCpp # For loading transformer models.10from langchain.document_loaders import PyPDFLoader, TextLoader, JSONLoader, CSVLoader11import tempfile # 임시 파일을 생성하기 위한 라이브러리입니다.12import os13from huggingface_hub import hf_hub_download # Hugging Face Hub에서 모델을 다운로드하기 위한 함수입니다.14 15# PDF 문서로부터 텍스트를 추출하는 함수입니다.16def get_pdf_text(pdf_docs):17 temp_dir = tempfile.TemporaryDirectory() # 임시 디렉토리를 생성합니다.18 temp_filepath = os.path.join(temp_dir.name, pdf_docs.name) # 임시 파일 경로를 생성합니다.19 20 with open(temp_filepath, "wb") as f: # 임시 파일을 바이너리 쓰기 모드로 엽니다.21 f.write(pdf_docs.getvalue()) # PDF 문서의 내용을 임시 파일에 씁니다.22 23 pdf_loader = PyPDFLoader(temp_filepath) # PyPDFLoader를 사용해 PDF를 로드합니다.24 pdf_doc = pdf_loader.load() # 텍스트를 추출합니다.25 26 return pdf_doc # 추출한 텍스트를 반환합니다.27 28# 과제29# 아래 텍스트 추출 함수를 작성30def get_text_file(docs):31 temp_dir = tempfile.TemporaryDirectory()32 temp_filepath = os.path.join(temp_dir.name, docs.name)33 34 with open(temp_filepath, 'wb') as f:35 f.write(docs.getvalue())36 37 txt_loader = TextLoader(temp_filepath)38 txt_doc = txt_loader.load()39 40 return txt_doc41 42def get_csv_file(docs):43 temp_dir = tempfile.TemporaryDirectory()44 temp_filepath = os.path.join(temp_dir.name, docs.name)45 46 with open(temp_filepath, 'wb') as f:47 f.write(docs.getvalue())48 49 csv_loader = CSVLoader(temp_filepath)50 csv_doc = csv_loader.load()51 52 return csv_doc53 54def get_json_file(docs):55 temp_dir = tempfile.TemporaryDirectory()56 temp_filepath = os.path.join(temp_dir.name, docs.name)57 58 with open(temp_filepath, 'wb') as f:59 f.write(docs.getvalue())60 61 json_loader = JSONLoader(temp_filepath, jq_schema='.clerk[]')62 json_doc = json_loader.load()63 64 return json_doc 65 66 67# 문서들을 처리하여 텍스트 청크로 나누는 함수입니다.68def get_text_chunks(documents):69 text_splitter = RecursiveCharacterTextSplitter(70 chunk_size=1000, # 청크의 크기를 지정합니다.71 chunk_overlap=200, # 청크 사이의 중복을 지정합니다.72 length_function=len # 텍스트의 길이를 측정하는 함수를 지정합니다.73 )74 75 documents = text_splitter.split_documents(documents) # 문서들을 청크로 나눕니다.76 return documents # 나눈 청크를 반환합니다.77 78 79# 텍스트 청크들로부터 벡터 스토어를 생성하는 함수입니다.80def get_vectorstore(text_chunks):81 # 원하는 임베딩 모델을 로드합니다.82 embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L12-v2',83 model_kwargs={'device': 'cpu'}) # 임베딩 모델을 설정합니다.84 vectorstore = FAISS.from_documents(text_chunks, embeddings) # FAISS 벡터 스토어를 생성합니다.85 return vectorstore # 생성된 벡터 스토어를 반환합니다.86 87 88def get_conversation_chain(vectorstore):89 model_name_or_path = 'TheBloke/Llama-2-7B-chat-GGUF'90 model_basename = 'llama-2-7b-chat.Q2_K.gguf'91 model_path = hf_hub_download(repo_id=model_name_or_path, filename=model_basename)92 93 llm = LlamaCpp(model_path=model_path,94 n_ctx=4086,95 input={"temperature": 0.75, "max_length": 2000, "top_p": 1},96 verbose=True, )97 # 대화 기록을 저장하기 위한 메모리를 생성합니다.98 memory = ConversationBufferMemory(99 memory_key='chat_history', return_messages=True)100 # 대화 검색 체인을 생성합니다.101 conversation_chain = ConversationalRetrievalChain.from_llm(102 llm=llm,103 retriever=vectorstore.as_retriever(),104 memory=memory105 )106 return conversation_chain # 생성된 대화 체인을 반환합니다.107 108# 사용자 입력을 처리하는 함수입니다.109def handle_userinput(user_question):110 print('user_question => ', user_question)111 # 대화 체인을 사용하여 사용자 질문에 대한 응답을 생성합니다.112 response = st.session_state.conversation({'question': user_question})113 # 대화 기록을 저장합니다.114 st.session_state.chat_history = response['chat_history']115 116 for i, message in enumerate(st.session_state.chat_history):117 if i % 2 == 0:118 st.write(user_template.replace(119 "{{MSG}}", message.content), unsafe_allow_html=True)120 else:121 st.write(bot_template.replace(122 "{{MSG}}", message.content), unsafe_allow_html=True)123 124 125def main():126 load_dotenv()127 st.set_page_config(page_title="Chat with multiple Files",128 page_icon=":books:")129 st.write(css, unsafe_allow_html=True)130 131 if "conversation" not in st.session_state:132 st.session_state.conversation = None133 if "chat_history" not in st.session_state:134 st.session_state.chat_history = None135 136 st.header("Chat with multiple Files:")137 user_question = st.text_input("Ask a question about your documents:")138 if user_question:139 handle_userinput(user_question)140 141 with st.sidebar:142 st.subheader("Your documents")143 docs = st.file_uploader(144 "Upload your PDFs here and click on 'Process'", accept_multiple_files=True)145 if st.button("Process"):146 with st.spinner("Processing"):147 # get pdf text148 doc_list = []149 150 for file in docs:151 print('file - type : ', file.type)152 if file.type == 'text/plain':153 # file is .txt154 doc_list.extend(get_text_file(file))155 elif file.type in ['application/octet-stream', 'application/pdf']:156 # file is .pdf157 doc_list.extend(get_pdf_text(file))158 elif file.type == 'text/csv':159 # file is .csv160 doc_list.extend(get_csv_file(file))161 elif file.type == 'application/json':162 # file is .json163 doc_list.extend(get_json_file(file))164 165 # get the text chunks166 text_chunks = get_text_chunks(doc_list)167 168 # create vector store169 vectorstore = get_vectorstore(text_chunks)170 171 # create conversation chain172 st.session_state.conversation = get_conversation_chain(173 vectorstore)174 175 176if __name__ == '__main__':177 main()