Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import os | |
| from PyPDF2 import PdfReader | |
| import docx | |
| from langchain.chat_models import ChatOpenAI | |
| from langchain.llms import OpenAI | |
| from dotenv import load_dotenv | |
| from langchain.embeddings import HuggingFaceEmbeddings | |
| from langchain.text_splitter import CharacterTextSplitter | |
| from langchain.vectorstores import FAISS | |
| from langchain.chains import ConversationalRetrievalChain | |
| from langchain.memory import ConversationBufferMemory | |
| from streamlit_chat import message | |
| from langchain.callbacks import get_openai_callback | |
| # Load environment variables | |
| load_dotenv() | |
| openapi_key = os.getenv("OPENAI_API_KEY") | |
| def main(): | |
| st.set_page_config(page_title="Chat with your file") | |
| st.header("DocumentGPT") | |
| 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: | |
| uploaded_files = st.file_uploader("Upload your file", type=['pdf', 'docx'], accept_multiple_files=True) | |
| process = st.button("Process") | |
| if process: | |
| if not openapi_key: | |
| st.info("Please add your OpenAI API key to continue.") | |
| st.stop() | |
| files_text = get_files_text(uploaded_files) | |
| st.write("File loaded...") | |
| text_chunks = get_text_chunks(files_text) | |
| st.write("File chunks created...") | |
| vectorstore = get_vectorstore(text_chunks) | |
| st.write("Vector Store Created...") | |
| st.session_state.conversation = get_conversation_chain(vectorstore, openapi_key) | |
| st.session_state.processComplete = True | |
| if st.session_state.processComplete: | |
| user_question = st.chat_input("Ask a question about your files.") | |
| if user_question: | |
| handle_user_input(user_question) | |
| def get_files_text(uploaded_files): | |
| text = "" | |
| for uploaded_file in uploaded_files: | |
| file_extension = os.path.splitext(uploaded_file.name)[1] | |
| if file_extension == ".pdf": | |
| text += get_pdf_text(uploaded_file) | |
| elif file_extension == ".docx": | |
| text += get_docx_text(uploaded_file) | |
| return text | |
| def get_pdf_text(pdf): | |
| pdf_reader = PdfReader(pdf) | |
| text = "" | |
| for page in pdf_reader.pages: | |
| text += page.extract_text() | |
| return text | |
| def get_docx_text(file): | |
| doc = docx.Document(file) | |
| return ' '.join([para.text for para in doc.paragraphs]) | |
| def get_text_chunks(text): | |
| text_splitter = CharacterTextSplitter( | |
| separator="\n", | |
| chunk_size=900, | |
| chunk_overlap=100, | |
| length_function=len | |
| ) | |
| return text_splitter.split_text(text) | |
| def get_vectorstore(text_chunks): | |
| embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2") | |
| return FAISS.from_texts(text_chunks, embeddings) | |
| def get_conversation_chain(vectorstore, openapi_key): | |
| llm = ChatOpenAI(openai_api_key=openapi_key, model_name='gpt-3.5-turbo', temperature=0) | |
| memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True) | |
| return ConversationalRetrievalChain.from_llm( | |
| llm=llm, | |
| retriever=vectorstore.as_retriever(), | |
| memory=memory | |
| ) | |
| def handle_user_input(user_question): | |
| with get_openai_callback() as cb: | |
| response = st.session_state.conversation({'question': user_question}) | |
| st.session_state.chat_history = response['chat_history'] | |
| response_container = st.container() | |
| with response_container: | |
| for i, message in enumerate(st.session_state.chat_history): | |
| message(message.content, is_user=(i % 2 == 0), key=str(i)) | |
| if __name__ == '__main__': | |
| main() |