Spaces:
Sleeping
Sleeping
Added main file
Browse files- chat_with_docs.py +131 -0
chat_with_docs.py
ADDED
@@ -0,0 +1,131 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader
|
2 |
+
from llama_index import download_loader
|
3 |
+
from pandasai.llm.openai import OpenAI
|
4 |
+
from matplotlib import pyplot as plt
|
5 |
+
import streamlit as st
|
6 |
+
import pandas as pd
|
7 |
+
import os
|
8 |
+
|
9 |
+
|
10 |
+
documents_folder = "./documents"
|
11 |
+
|
12 |
+
# Load PandasAI loader, Which is a wrapper over PandasAI library
|
13 |
+
PandasAIReader = download_loader("PandasAIReader")
|
14 |
+
|
15 |
+
st.title("Welcome to `ChatwithDocs`")
|
16 |
+
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⭐")
|
17 |
+
|
18 |
+
|
19 |
+
def get_csv_result(df, query):
|
20 |
+
reader = PandasAIReader(llm=csv_llm)
|
21 |
+
response = reader.run_pandas_ai(
|
22 |
+
df,
|
23 |
+
query,
|
24 |
+
is_conversational_answer=False
|
25 |
+
)
|
26 |
+
return response
|
27 |
+
|
28 |
+
def save_file(doc):
|
29 |
+
fn = os.path.basename(doc.name)
|
30 |
+
# open read and write the file into the server
|
31 |
+
open(documents_folder+'/'+fn, 'wb').write(doc.read())
|
32 |
+
# Check for the current filename, If new filename
|
33 |
+
# clear the previous cached vectors and update the filename
|
34 |
+
# with current name
|
35 |
+
if st.session_state.get('file_name'):
|
36 |
+
if st.session_state.file_name != fn:
|
37 |
+
st.cache_resource.clear()
|
38 |
+
st.session_state['file_name'] = fn
|
39 |
+
else:
|
40 |
+
st.session_state['file_name'] = fn
|
41 |
+
|
42 |
+
return fn
|
43 |
+
|
44 |
+
def remove_file(file_path):
|
45 |
+
# Remove the file from the Document folder once
|
46 |
+
# vectors are created
|
47 |
+
if os.path.isfile(documents_folder+'/'+file_path):
|
48 |
+
os.remove(documents_folder+'/'+file_path)
|
49 |
+
|
50 |
+
|
51 |
+
|
52 |
+
@st.cache_resource
|
53 |
+
def create_index():
|
54 |
+
# Create vectors for the file stored under Document folder.
|
55 |
+
# NOTE: You can create vectors for multiple files at once.
|
56 |
+
documents = SimpleDirectoryReader(documents_folder).load_data()
|
57 |
+
index = GPTVectorStoreIndex.from_documents(documents)
|
58 |
+
return index
|
59 |
+
|
60 |
+
|
61 |
+
|
62 |
+
def query_doc(vector_index, query):
|
63 |
+
# Applies Similarity Algo, Finds the nearest match and
|
64 |
+
# take the match and user query to OpenAI for rich response
|
65 |
+
query_engine = vector_index.as_query_engine()
|
66 |
+
response = query_engine.query(query)
|
67 |
+
return response
|
68 |
+
|
69 |
+
|
70 |
+
api_key = st.text_input("Enter your OpenAI API key here:", type="password")
|
71 |
+
if api_key:
|
72 |
+
os.environ['OPENAI_API_KEY'] = api_key
|
73 |
+
csv_llm = OpenAI(api_token=api_key)
|
74 |
+
|
75 |
+
|
76 |
+
tab1, tab2= st.tabs(["CSV", "PDFs/Docs"])
|
77 |
+
|
78 |
+
with tab1:
|
79 |
+
|
80 |
+
st.write("Chat with CSV files using PandasAI loader with LlamaIndex")
|
81 |
+
input_csv = st.file_uploader("Upload your CSV file", type=['csv'])
|
82 |
+
|
83 |
+
if input_csv is not None:
|
84 |
+
st.info("CSV Uploaded Successfully")
|
85 |
+
df = pd.read_csv(input_csv)
|
86 |
+
st.dataframe(df, use_container_width=True)
|
87 |
+
|
88 |
+
|
89 |
+
st.divider()
|
90 |
+
|
91 |
+
input_text = st.text_area("Ask your query")
|
92 |
+
|
93 |
+
if input_text is not None:
|
94 |
+
if st.button("Send"):
|
95 |
+
st.info("Your query: "+ input_text)
|
96 |
+
with st.spinner('Processing your query...'):
|
97 |
+
response = get_csv_result(df, input_text)
|
98 |
+
if plt.get_fignums():
|
99 |
+
st.pyplot(plt.gcf())
|
100 |
+
else:
|
101 |
+
st.success(response)
|
102 |
+
|
103 |
+
|
104 |
+
with tab2:
|
105 |
+
st.write("Chat with PDFs/Docs")
|
106 |
+
input_doc = st.file_uploader("Upload your Docs")
|
107 |
+
|
108 |
+
if input_doc is not None:
|
109 |
+
st.info("Doc Uploaded Successfully")
|
110 |
+
file_name = save_file(input_doc)
|
111 |
+
index = create_index()
|
112 |
+
remove_file(file_name)
|
113 |
+
|
114 |
+
|
115 |
+
st.divider()
|
116 |
+
input_text = st.text_area("Ask your question")
|
117 |
+
|
118 |
+
if input_text is not None:
|
119 |
+
if st.button("Ask"):
|
120 |
+
st.info("Your query: \n" +input_text)
|
121 |
+
with st.spinner("Processing your query.."):
|
122 |
+
response = query_doc(index, input_text)
|
123 |
+
print(response)
|
124 |
+
|
125 |
+
st.success(response)
|
126 |
+
|
127 |
+
st.divider()
|
128 |
+
# Shows the source documents context which
|
129 |
+
# has been used to prepare the response
|
130 |
+
st.write("Source Documents")
|
131 |
+
st.write(response.get_formatted_sources())
|