Chat-with-Docs / app.py
ashrma's picture
Updated divider code in streamlit
2cf8f22
from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader
from llama_index import download_loader
from pandasai.llm.openai import OpenAI
from matplotlib import pyplot as plt
import streamlit as st
import pandas as pd
import os
documents_folder = "./documents"
# Load PandasAI loader, Which is a wrapper over PandasAI library
PandasAIReader = download_loader("PandasAIReader")
st.title("Welcome to `ChatwithDocs`")
st.header("Interact with Documents such as `PDFs/CSV/Docs` using the power of LLMs\nPowered by `LlamaIndex🦙` \nCheckout the [GITHUB Repo Here](https://github.com/anoopshrma/Chat-with-Docs) and Leave a star⭐")
def get_csv_result(df, query):
reader = PandasAIReader(llm=csv_llm)
response = reader.run_pandas_ai(
df,
query,
is_conversational_answer=False
)
return response
def save_file(doc):
fn = os.path.basename(doc.name)
# open read and write the file into the server
open(documents_folder+'/'+fn, 'wb').write(doc.read())
# Check for the current filename, If new filename
# clear the previous cached vectors and update the filename
# with current name
if st.session_state.get('file_name'):
if st.session_state.file_name != fn:
st.cache_resource.clear()
st.session_state['file_name'] = fn
else:
st.session_state['file_name'] = fn
return fn
def remove_file(file_path):
# Remove the file from the Document folder once
# vectors are created
if os.path.isfile(documents_folder+'/'+file_path):
os.remove(documents_folder+'/'+file_path)
@st.cache_resource
def create_index():
# Create vectors for the file stored under Document folder.
# NOTE: You can create vectors for multiple files at once.
documents = SimpleDirectoryReader(documents_folder).load_data()
index = GPTVectorStoreIndex.from_documents(documents)
return index
def query_doc(vector_index, query):
# Applies Similarity Algo, Finds the nearest match and
# take the match and user query to OpenAI for rich response
query_engine = vector_index.as_query_engine()
response = query_engine.query(query)
return response
api_key = st.text_input("Enter your OpenAI API key here:", type="password")
if api_key:
os.environ['OPENAI_API_KEY'] = api_key
csv_llm = OpenAI(api_token=api_key)
tab1, tab2= st.tabs(["CSV", "PDFs/Docs"])
with tab1:
st.write("Chat with CSV files using PandasAI loader with LlamaIndex")
input_csv = st.file_uploader("Upload your CSV file", type=['csv'])
if input_csv is not None:
st.info("CSV Uploaded Successfully")
df = pd.read_csv(input_csv)
st.dataframe(df, use_container_width=True)
st.write("---")
input_text = st.text_area("Ask your query")
if input_text is not None:
if st.button("Send"):
st.info("Your query: "+ input_text)
with st.spinner('Processing your query...'):
response = get_csv_result(df, input_text)
if plt.get_fignums():
st.pyplot(plt.gcf())
else:
st.success(response)
with tab2:
st.write("Chat with PDFs/Docs")
input_doc = st.file_uploader("Upload your Docs")
if input_doc is not None:
st.info("Doc Uploaded Successfully")
file_name = save_file(input_doc)
index = create_index()
remove_file(file_name)
st.write("---")
input_text = st.text_area("Ask your question")
if input_text is not None:
if st.button("Ask"):
st.info("Your query: \n" +input_text)
with st.spinner("Processing your query.."):
response = query_doc(index, input_text)
print(response)
st.success(response)
st.write("---")
# Shows the source documents context which
# has been used to prepare the response
st.write("Source Documents")
st.write(response.get_formatted_sources())