File size: 1,847 Bytes
0b03ede
e239fba
b075822
 
e239fba
 
 
0b03ede
b075822
 
0b03ede
 
 
 
 
 
 
b075822
0b03ede
 
b075822
ee8270e
b075822
 
 
 
 
 
 
 
 
 
 
ee8270e
b075822
 
 
 
 
 
 
 
 
 
 
 
 
 
ee8270e
b075822
 
0b03ede
 
b075822
0b03ede
b075822
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
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():
    st.title("Upload File to chat with file")
    uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
    if uploaded_file is None:
        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.retriever_engine = process_pdf(uploaded_file)
                st.success("PDF processed successfully!")

def ask_question():
    if st.session_state.retriever_engine :
        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.retriever_engine)
                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:")
            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 "retriever_engine" not in st.session_state:
        st.session_state.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()