ashrma commited on
Commit
0b1eb85
·
1 Parent(s): 57f27a1

Added main file

Browse files
Files changed (1) hide show
  1. 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())