|
import os |
|
import uuid |
|
import tempfile |
|
import streamlit as st |
|
import openai |
|
from langchain.retrievers.multi_vector import MultiVectorRetriever |
|
from langchain.vectorstores import Chroma |
|
from langchain.embeddings import OpenAIEmbeddings |
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
from langchain.storage import InMemoryStore |
|
from langchain.memory import ConversationBufferMemory |
|
from langchain.llms import OpenAI |
|
from langchain.chains import ConversationalRetrievalChain |
|
from langchain.chat_models import ChatOpenAI |
|
from langchain.prompts import ChatPromptTemplate |
|
from langchain.schema.output_parser import StrOutputParser |
|
import uuid |
|
from langchain.schema.document import Document |
|
from langchain.output_parsers.openai_functions import JsonKeyOutputFunctionsParser |
|
from langchain.document_loaders import PyPDFLoader |
|
|
|
|
|
OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"] |
|
if not OPENAI_API_KEY: |
|
st.error("OPENAI_API_KEY not set in environment variables!") |
|
raise SystemExit |
|
openai.api_key = OPENAI_API_KEY |
|
|
|
|
|
def process_pdf(uploaded_file): |
|
with st.spinner("Processing PDF..."): |
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: |
|
tmp.write(uploaded_file.getvalue()) |
|
tmp_path = tmp.name |
|
loaders = [PyPDFLoader(tmp_path)] |
|
docs = [] |
|
for l in loaders: |
|
docs.extend(l.load()) |
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000) |
|
docs = text_splitter.split_documents(docs) |
|
return docs |
|
|
|
|
|
def smaller_chunks_strategy(docs): |
|
prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="1") |
|
if prompt: |
|
with st.spinner('Processing with smaller_chunks_strategy'): |
|
vectorstore = Chroma( |
|
collection_name="full_documents", |
|
embedding_function=OpenAIEmbeddings() |
|
) |
|
store = InMemoryStore() |
|
id_key = "doc_id" |
|
retriever = MultiVectorRetriever( |
|
vectorstore=vectorstore, |
|
docstore=store, |
|
id_key=id_key, |
|
) |
|
doc_ids = [str(uuid.uuid4()) for _ in docs] |
|
child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=400) |
|
sub_docs = [] |
|
for i, doc in enumerate(docs): |
|
_id = doc_ids[i] |
|
_sub_docs = child_text_splitter.split_documents([doc]) |
|
for _doc in _sub_docs: |
|
_doc.metadata[id_key] = _id |
|
sub_docs.extend(_sub_docs) |
|
|
|
retriever.vectorstore.add_documents(sub_docs) |
|
retriever.docstore.mset(list(zip(doc_ids, docs))) |
|
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True) |
|
qa = ConversationalRetrievalChain.from_llm(OpenAI(temperature=0), retriever, memory=memory) |
|
st.info(prompt, icon="🧐") |
|
result = qa({"question": prompt}) |
|
st.success(result['answer'], icon="🤖") |
|
|
|
|
|
def summary_strategy(docs): |
|
prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="2") |
|
if prompt: |
|
with st.spinner('Processing with summary_strategy'): |
|
chain = ( |
|
{"doc": lambda x: x.page_content} |
|
| ChatPromptTemplate.from_template("Summarize the following document:\n\n{doc}") |
|
| ChatOpenAI(max_retries=0) |
|
| StrOutputParser() |
|
) |
|
summaries = chain.batch(docs, {"max_concurrency": 5}) |
|
vectorstore = Chroma( |
|
collection_name="summaries", |
|
embedding_function= OpenAIEmbeddings() |
|
) |
|
store = InMemoryStore() |
|
id_key = "doc_id" |
|
retriever = MultiVectorRetriever( |
|
vectorstore=vectorstore, |
|
docstore=store, |
|
id_key=id_key, |
|
) |
|
doc_ids = [str(uuid.uuid4()) for _ in docs] |
|
summary_docs = [Document(page_content=s, metadata={id_key: doc_ids[i]}) for i, s in enumerate(summaries)] |
|
retriever.vectorstore.add_documents(summary_docs) |
|
retriever.docstore.mset(list(zip(doc_ids, docs))) |
|
qa = ConversationalRetrievalChain.from_llm(OpenAI(temperature=0), retriever, memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True)) |
|
st.info(prompt, icon="🧐") |
|
result = qa({"question": prompt}) |
|
st.success(result['answer'], icon="🤖") |
|
|
|
|
|
def hypothetical_questions_strategy(docs): |
|
prompt = st.text_input("Enter Your Question:", placeholder="Ask something", key="3") |
|
if prompt: |
|
with st.spinner('Processing with hypothetical_questions_strategy'): |
|
functions = [ |
|
{ |
|
"name": "hypothetical_questions", |
|
"description": "Generate hypothetical questions", |
|
"parameters": { |
|
"type": "object", |
|
"properties": { |
|
"questions": { |
|
"type": "array", |
|
"items": { |
|
"type": "string" |
|
}, |
|
}, |
|
}, |
|
"required": ["questions"] |
|
} |
|
} |
|
] |
|
chain = ( |
|
{"doc": lambda x: x.page_content} |
|
| ChatPromptTemplate.from_template("Generate a list of 3 hypothetical questions that the below document could be used to answer:\n\n{doc}") |
|
| ChatOpenAI(max_retries=0, model="gpt-4").bind(functions=functions, function_call={"name": "hypothetical_questions"}) |
|
| JsonKeyOutputFunctionsParser(key_name="questions") |
|
) |
|
hypothetical_questions = chain.batch(docs, {"max_concurrency": 5}) |
|
vectorstore = Chroma( |
|
collection_name="hypo-questions", |
|
embedding_function=OpenAIEmbeddings() |
|
) |
|
store = InMemoryStore() |
|
id_key = "doc_id" |
|
retriever = MultiVectorRetriever( |
|
vectorstore=vectorstore, |
|
docstore=store, |
|
id_key=id_key, |
|
) |
|
doc_ids = [str(uuid.uuid4()) for _ in docs] |
|
question_docs = [] |
|
for i, question_list in enumerate(hypothetical_questions): |
|
question_docs.extend([Document(page_content=s, metadata={id_key: doc_ids[i]}) for s in question_list]) |
|
retriever.vectorstore.add_documents(question_docs) |
|
retriever.docstore.mset(list(zip(doc_ids, docs))) |
|
qa = ConversationalRetrievalChain.from_llm(OpenAI(temperature=0), retriever, memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True)) |
|
st.info(prompt, icon="🧐") |
|
result = qa({"question": prompt}) |
|
st.success(result['answer'], icon="🤖") |
|
|
|
|
|
|
|
def app(): |
|
image_path = "icon.png" |
|
st.sidebar.image(image_path, caption="icon", use_column_width=True) |
|
st.title("VecDBCompare 0.0.1") |
|
st.sidebar.markdown(""" |
|
# 🚀 **VecDBCompare: Your Vector DB Strategy Tester** |
|
## 📌 **What is it?** |
|
VecDBCompare lets you evaluate and compare three vector database retrieval strategies in a snap! |
|
## 📤 **How to Use?** |
|
1. **Upload a PDF** 📄 |
|
2. Get **Three QABots** 🤖🤖🤖, each with a different strategy. |
|
3. **Ask questions** ❓ and see how each bot responds differently. |
|
4. **Decide** ✅ which strategy works best for you! |
|
## 🌟 **Why VecDBCompare?** |
|
- **Simple & Fast** ⚡: Upload, ask, and compare! |
|
- **Real-time Comparison** 🔍: See strategies in action side-by-side. |
|
- **Empower Your Choice** 💡: Pick the best strategy for your needs. |
|
Dive in and discover with VecDBCompare! 🌐 |
|
""") |
|
uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"]) |
|
if uploaded_file: |
|
docs = process_pdf(uploaded_file) |
|
option = st.selectbox( |
|
"Which retrieval strategy would you like to use?", |
|
("Smaller Chunks", "Summary", "Hypothetical Questions") |
|
) |
|
if option == 'Smaller Chunks': |
|
smaller_chunks_strategy(docs) |
|
elif option == 'Summary': |
|
summary_strategy(docs) |
|
elif option == 'Hypothetical Questions': |
|
hypothetical_questions_strategy(docs) |
|
|
|
|
|
if __name__ == "__main__": |
|
st.set_page_config(layout="wide") |
|
app() |
|
|
