Spaces:
Sleeping
Sleeping
import streamlit as st | |
from dotenv import load_dotenv | |
from PyPDF2 import PdfReader | |
from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter | |
from langchain.embeddings import OpenAIEmbeddings, HuggingFaceInstructEmbeddings | |
from langchain.vectorstores import FAISS, Chroma | |
from langchain.embeddings import HuggingFaceEmbeddings # General embeddings from HuggingFace models. | |
from langchain.chat_models import ChatOpenAI | |
from langchain.memory import ConversationBufferMemory | |
from langchain.chains import ConversationalRetrievalChain | |
from htmlTemplates import css, bot_template, user_template | |
from langchain.llms import HuggingFaceHub, LlamaCpp, CTransformers # For loading transformer models. | |
from langchain.document_loaders import PyPDFLoader, TextLoader, JSONLoader, CSVLoader | |
import tempfile # μμ νμΌμ μμ±νκΈ° μν λΌμ΄λΈλ¬λ¦¬μ λλ€. | |
import os | |
# PDF λ¬Έμλ‘λΆν° ν μ€νΈλ₯Ό μΆμΆνλ ν¨μμ λλ€. | |
def get_pdf_text(pdf_docs): | |
temp_dir = tempfile.TemporaryDirectory() # μμ λλ ν 리λ₯Ό μμ±ν©λλ€. | |
temp_filepath = os.path.join(temp_dir.name, pdf_docs.name) # μμ νμΌ κ²½λ‘λ₯Ό μμ±ν©λλ€. | |
with open(temp_filepath, "wb") as f: # μμ νμΌμ λ°μ΄λ리 μ°κΈ° λͺ¨λλ‘ μ½λλ€. | |
f.write(pdf_docs.getvalue()) # PDF λ¬Έμμ λ΄μ©μ μμ νμΌμ μλλ€. | |
pdf_loader = PyPDFLoader(temp_filepath) # PyPDFLoaderλ₯Ό μ¬μ©ν΄ PDFλ₯Ό λ‘λν©λλ€. | |
pdf_doc = pdf_loader.load() # ν μ€νΈλ₯Ό μΆμΆν©λλ€. | |
return pdf_doc # μΆμΆν ν μ€νΈλ₯Ό λ°νν©λλ€. | |
# κ³Όμ | |
# μλ ν μ€νΈ μΆμΆ ν¨μλ₯Ό μμ± | |
def get_text_file(docs): | |
pass | |
def get_csv_file(docs): | |
pass | |
def get_json_file(docs): | |
pass | |
# λ¬Έμλ€μ μ²λ¦¬νμ¬ ν μ€νΈ μ²ν¬λ‘ λλλ ν¨μμ λλ€. | |
def get_text_chunks(documents): | |
text_splitter = RecursiveCharacterTextSplitter( | |
chunk_size=1000, # μ²ν¬μ ν¬κΈ°λ₯Ό μ§μ ν©λλ€. | |
chunk_overlap=200, # μ²ν¬ μ¬μ΄μ μ€λ³΅μ μ§μ ν©λλ€. | |
length_function=len # ν μ€νΈμ κΈΈμ΄λ₯Ό μΈ‘μ νλ ν¨μλ₯Ό μ§μ ν©λλ€. | |
) | |
documents = text_splitter.split_documents(documents) # λ¬Έμλ€μ μ²ν¬λ‘ λλλλ€ | |
return documents # λλ μ²ν¬λ₯Ό λ°νν©λλ€. | |
# ν μ€νΈ μ²ν¬λ€λ‘λΆν° λ²‘ν° μ€ν μ΄λ₯Ό μμ±νλ ν¨μμ λλ€. | |
def get_vectorstore(text_chunks): | |
# OpenAI μλ² λ© λͺ¨λΈμ λ‘λν©λλ€. (Embedding models - Ada v2) | |
embeddings = OpenAIEmbeddings() | |
vectorstore = FAISS.from_documents(text_chunks, embeddings) # FAISS λ²‘ν° μ€ν μ΄λ₯Ό μμ±ν©λλ€. | |
return vectorstore # μμ±λ λ²‘ν° μ€ν μ΄λ₯Ό λ°νν©λλ€. | |
def get_conversation_chain(vectorstore): | |
gpt_model_name = 'gpt-3.5-turbo' | |
llm = ChatOpenAI(model_name = gpt_model_name) #gpt-3.5 λͺ¨λΈ λ‘λ | |
# λν κΈ°λ‘μ μ μ₯νκΈ° μν λ©λͺ¨λ¦¬λ₯Ό μμ±ν©λλ€. | |
memory = ConversationBufferMemory( | |
memory_key='chat_history', return_messages=True) | |
# λν κ²μ 체μΈμ μμ±ν©λλ€. | |
conversation_chain = ConversationalRetrievalChain.from_llm( | |
llm=llm, | |
retriever=vectorstore.as_retriever(), | |
memory=memory | |
) | |
return conversation_chain | |
# μ¬μ©μ μ λ ₯μ μ²λ¦¬νλ ν¨μμ λλ€. | |
def handle_userinput(user_question): | |
# λν 체μΈμ μ¬μ©νμ¬ μ¬μ©μ μ§λ¬Έμ λν μλ΅μ μμ±ν©λλ€. | |
response = st.session_state.conversation({'question': user_question}) | |
# λν κΈ°λ‘μ μ μ₯ν©λλ€. | |
st.session_state.chat_history = response['chat_history'] | |
for i, message in enumerate(st.session_state.chat_history): | |
if i % 2 == 0: | |
st.write(user_template.replace( | |
"{{MSG}}", message.content), unsafe_allow_html=True) | |
else: | |
st.write(bot_template.replace( | |
"{{MSG}}", message.content), unsafe_allow_html=True) | |
def main(): | |
load_dotenv() | |
st.set_page_config(page_title="Chat with multiple Files", | |
page_icon=":books:") | |
st.write(css, unsafe_allow_html=True) | |
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 | |
st.header("Chat with multiple Files :") | |
user_question = st.text_input("Ask a question about your documents:") | |
if user_question: | |
handle_userinput(user_question) | |
with st.sidebar: | |
openai_key = st.text_input("Paste your OpenAI API key (sk-...)") | |
if openai_key: | |
os.environ["OPENAI_API_KEY"] = openai_key | |
st.subheader("Your documents") | |
docs = st.file_uploader( | |
"Upload your PDFs here and click on 'Process'", accept_multiple_files=True) | |
if st.button("Process"): | |
with st.spinner("Processing"): | |
# get pdf text | |
doc_list = [] | |
for file in docs: | |
print('file - type : ', file.type) | |
if file.type == 'text/plain': | |
# file is .txt | |
doc_list.extend(get_text_file(file)) | |
elif file.type in ['application/octet-stream', 'application/pdf']: | |
# file is .pdf | |
doc_list.extend(get_pdf_text(file)) | |
elif file.type == 'text/csv': | |
# file is .csv | |
doc_list.extend(get_csv_file(file)) | |
elif file.type == 'application/json': | |
# file is .json | |
doc_list.extend(get_json_file(file)) | |
# get the text chunks | |
text_chunks = get_text_chunks(doc_list) | |
# create vector store | |
vectorstore = get_vectorstore(text_chunks) | |
# create conversation chain | |
st.session_state.conversation = get_conversation_chain( | |
vectorstore) | |
if __name__ == '__main__': | |
main() | |