Spaces:
Sleeping
Sleeping
File size: 5,333 Bytes
3a7d601 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import os, tempfile
import pinecone
from pathlib import Path
from langchain.chains import RetrievalQA, ConversationalRetrievalChain
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain import OpenAI
from langchain.llms.openai import OpenAIChat
from langchain.document_loaders import DirectoryLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma, Pinecone
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.memory import ConversationBufferMemory
from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
import streamlit as st
TMP_DIR = Path(__file__).resolve().parent.joinpath('data', 'tmp')
LOCAL_VECTOR_STORE_DIR = Path(__file__).resolve().parent.joinpath('data', 'vector_store')
st.set_page_config(page_title="RAG")
st.title("Retrieval Augmented Generation Engine")
def load_documents():
loader = DirectoryLoader(TMP_DIR.as_posix(), glob='**/*.pdf')
documents = loader.load()
return documents
def split_documents(documents):
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)
return texts
def embeddings_on_local_vectordb(texts):
vectordb = Chroma.from_documents(texts, embedding=OpenAIEmbeddings(),
persist_directory=LOCAL_VECTOR_STORE_DIR.as_posix())
vectordb.persist()
retriever = vectordb.as_retriever(search_kwargs={'k': 7})
return retriever
def embeddings_on_pinecone(texts):
pinecone.init(api_key=st.session_state.pinecone_api_key, environment=st.session_state.pinecone_env)
embeddings = OpenAIEmbeddings(openai_api_key=st.session_state.openai_api_key)
vectordb = Pinecone.from_documents(texts, embeddings, index_name=st.session_state.pinecone_index)
retriever = vectordb.as_retriever()
return retriever
def query_llm(retriever, query):
qa_chain = ConversationalRetrievalChain.from_llm(
llm=OpenAIChat(openai_api_key=st.session_state.openai_api_key),
retriever=retriever,
return_source_documents=True,
)
result = qa_chain({'question': query, 'chat_history': st.session_state.messages})
result = result['answer']
st.session_state.messages.append((query, result))
return result
def input_fields():
#
with st.sidebar:
#
if "openai_api_key" in st.secrets:
st.session_state.openai_api_key = st.secrets.openai_api_key
else:
st.session_state.openai_api_key = st.text_input("OpenAI API key", type="password")
#
if "pinecone_api_key" in st.secrets:
st.session_state.pinecone_api_key = st.secrets.pinecone_api_key
else:
st.session_state.pinecone_api_key = st.text_input("Pinecone API key", type="password")
#
if "pinecone_env" in st.secrets:
st.session_state.pinecone_env = st.secrets.pinecone_env
else:
st.session_state.pinecone_env = st.text_input("Pinecone environment")
#
if "pinecone_index" in st.secrets:
st.session_state.pinecone_index = st.secrets.pinecone_index
else:
st.session_state.pinecone_index = st.text_input("Pinecone index name")
#
st.session_state.pinecone_db = st.toggle('Use Pinecone Vector DB')
#
st.session_state.source_docs = st.file_uploader(label="Upload Documents", type="pdf", accept_multiple_files=True)
#
def process_documents():
if not st.session_state.openai_api_key or not st.session_state.pinecone_api_key or not st.session_state.pinecone_env or not st.session_state.pinecone_index or not st.session_state.source_docs:
st.warning(f"Please upload the documents and provide the missing fields.")
else:
try:
for source_doc in st.session_state.source_docs:
#
with tempfile.NamedTemporaryFile(delete=False, dir=TMP_DIR.as_posix(), suffix='.pdf') as tmp_file:
tmp_file.write(source_doc.read())
#
documents = load_documents()
#
for _file in TMP_DIR.iterdir():
temp_file = TMP_DIR.joinpath(_file)
temp_file.unlink()
#
texts = split_documents(documents)
#
if not st.session_state.pinecone_db:
st.session_state.retriever = embeddings_on_local_vectordb(texts)
else:
st.session_state.retriever = embeddings_on_pinecone(texts)
except Exception as e:
st.error(f"An error occurred: {e}")
def boot():
#
input_fields()
#
st.button("Submit Documents", on_click=process_documents)
#
if "messages" not in st.session_state:
st.session_state.messages = []
#
for message in st.session_state.messages:
st.chat_message('human').write(message[0])
st.chat_message('ai').write(message[1])
#
if query := st.chat_input():
st.chat_message("human").write(query)
response = query_llm(st.session_state.retriever, query)
st.chat_message("ai").write(response)
if __name__ == '__main__':
#
boot()
|