Spaces:
Sleeping
Sleeping
File size: 2,038 Bytes
20b1f3c |
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 |
import streamlit as st
import base64
import os
import tempfile
import time
import utilities as util
def main():
st.set_page_config(page_title="PDF Viewer", layout="wide", page_icon='./logo.png')
# st.image('./logo.png', width=60)
st.title(":blue[DocuBot]",anchor=False)
st.write("View and chat with your PDF")
if 'messages' not in st.session_state:
st.session_state.messages = [{'role': 'assistant', "content": "Hello! Upload a document and let's get started."}]
state = True
# with st.sidebar:
uploaded_file = st.sidebar.file_uploader("Upload your PDF File", type="pdf")
if uploaded_file:
state = False
with tempfile.TemporaryDirectory() as tmp_dir:
file_path = os.path.join(tmp_dir, uploaded_file.name)
with open(file_path, "wb") as f:
f.write(uploaded_file.getvalue()) # Write the PDF content
pdf_text = util.get_pdf_text(file_path)
pdf_frame = util.display_pdf(file_path)
st.sidebar.markdown(pdf_frame, unsafe_allow_html=True)
user_prompt = st.chat_input("What do you wanna know about the document?", disabled=state)
if st.sidebar.button(label="summarize"):
st.session_state.messages.append({'role': 'user', "content": "Summarize the document"})
with st.spinner("..."):
summary = util.summarize(pdf_text, max_length=200)
st.session_state.messages.append({'role': 'assistant', "content": "Summary of "+uploaded_file.name+": <br>"+summary})
if user_prompt:
st.session_state.messages.append({'role': 'user', "content": user_prompt})
response = "You asked: "+user_prompt
with st.spinner("..."):
time.sleep(2)
st.session_state.messages.append({'role': 'assistant', "content": response})
for message in st.session_state.messages:
with st.chat_message(message['role']):
st.markdown(message['content'], unsafe_allow_html=True)
if __name__ == "__main__":
main()
|