File size: 2,487 Bytes
0b03ede
e239fba
b075822
 
e239fba
 
 
0b03ede
b075822
5c3777d
 
 
 
 
 
0b03ede
 
5c3777d
 
0b03ede
 
 
 
5c3777d
0b03ede
5c3777d
5e32aa7
0b03ede
b075822
5c3777d
 
b075822
 
5c3777d
b075822
 
 
 
 
 
5e32aa7
 
 
ee8270e
b075822
 
 
 
 
 
 
 
5c3777d
 
b075822
 
 
 
ee8270e
b075822
 
0b03ede
 
b075822
0b03ede
5e32aa7
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
import streamlit as st

from pdf_processing import process_pdf
from retrieve_and_display import retrieve_and_query, plot_images

from dotenv import load_dotenv
load_dotenv()

def upload_file():
    if not st.session_state.filename_and_retriever_engine:
        st.title("Upload File to chat with file")
    else:
        st.title(f"File {st.session_state.filename_and_retriever_engine[0]} loaded.")
        st.info("Click on Chat in sidebar")
        st.info("Upload another file if you want to chat with a different pdf")
    uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
    if uploaded_file is None:
        if not st.session_state.filename_and_retriever_engine:
            st.info("Please upload a PDF file.")
    else:
        st.info(f"Uploaded PDF: {uploaded_file.name}")
        if st.button("Process PDF"):
            with st.spinner("Processing PDF..."):
                st.session_state.filename_and_retriever_engine = uploaded_file.name, process_pdf(uploaded_file)
                st.success("PDF processed successfully!")
                st.success("Click on Chat in sidebar")
    

def ask_question():
    if st.session_state.filename_and_retriever_engine :
        st.title(f"Chat with {st.session_state.filename_and_retriever_engine[0]}")
        if user_question := st.chat_input("Ask a question"):
            with st.spinner("Retrieving information..."):
                response, retrieved_image_path_list = retrieve_and_query(user_question, st.session_state.filename_and_retriever_engine[1])
                print(retrieved_image_path_list)
            st.write("Retrieved Context:")
            for node in response.source_nodes:
                st.code(node.node.get_text())
            
            st.write("\nRetrieved Images:")
            # if len(retrieved_image_path_list) > 0:
            plot_images(retrieved_image_path_list)
                # st.pyplot()

            st.write("\nFinal Answer:")
            st.code(response.response)
    else:
        st.title("Upload File to chat with file")
        
    
    
def main():
    if "filename_and_retriever_engine" not in st.session_state:
        st.session_state.filename_and_retriever_engine = None
    page_names_to_funcs = {
        "Upload File": upload_file,
        "Chat": ask_question
    }

    demo_name = st.sidebar.selectbox("PDF Query Tool", page_names_to_funcs.keys())
    page_names_to_funcs[demo_name]()

if __name__ == "__main__":
    # login_page()
    main()