Spaces:
Runtime error
Runtime error
import os | |
from dotenv import load_dotenv | |
import tempfile | |
import streamlit as st | |
from langchain_community.document_loaders import PyPDFLoader | |
from main import summarize_pdf_document | |
from main import retrieve_documents | |
from main import embed_document_data | |
from main import is_document_embedded | |
load_dotenv() | |
def main(): | |
st.sidebar.title("PDF Management") | |
uploaded_files = st.sidebar.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True) | |
files_info = [] # Initialize files_info to an empty list before checking for uploaded files | |
if uploaded_files: | |
files_info = save_uploaded_files(uploaded_files) | |
process_documents(files_info) | |
if st.button('Add Uploaded Documents in Q nd A'): | |
embed_documents(files_info) | |
# Call to display the Q&A section unconditionally | |
display_qna_section(files_info) | |
def save_uploaded_files(uploaded_files): | |
"""Save uploaded files to temporary directory and return their file paths along with original filenames.""" | |
files_info = [] | |
for uploaded_file in uploaded_files: | |
# Create a temporary file | |
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmpfile: | |
# Write contents of the uploaded file to the temporary file | |
tmpfile.write(uploaded_file.getvalue()) | |
# Append both the temp file path and the original name | |
files_info.append((tmpfile.name, uploaded_file.name)) | |
return files_info | |
def process_documents(files_info): | |
"""Handle document processing for summarization.""" | |
st.header("Document Summaries") | |
summarize_button = st.button('Summarize Uploaded Documents') | |
if summarize_button: | |
for temp_path, original_name in files_info: | |
with st.container(): # Using container to better organize the display | |
st.write(f"Summary for {original_name}:") # Display the original file name | |
try: | |
summary = summarize_pdf_document(temp_path, os.getenv('OPENAI_API_KEY')) | |
st.text_area("", value=summary, height=200, key=f"summary_{original_name}") | |
except Exception as e: | |
st.error(f"Failed to summarize {original_name}: {str(e)}") | |
def embed_documents(files_info): | |
"""Embed each document with correct metadata, replacing temp path with original filename.""" | |
for temp_path, original_name in files_info: | |
if not is_document_embedded(original_name): # Check if already embedded | |
try: | |
# Load the document | |
loader = PyPDFLoader(temp_path) | |
documents = loader.load() | |
# Update the metadata for each document | |
updated_documents = [] | |
for doc in documents: | |
doc.metadata['source'] = original_name # Use original filename | |
updated_documents.append(doc) | |
embed_document_data(updated_documents) # Embed the documents | |
st.success(f"Embedded {original_name}") | |
except Exception as e: | |
st.error(f"Failed to embed {original_name}: {str(e)}") | |
else: | |
st.info(f"{original_name} is already embedded.") | |
def display_qna_section(files_info): | |
"""Display Q&A section.""" | |
st.header("Question and Answer") | |
with st.form("qa_form"): | |
user_query = st.text_input("Enter your question here:") | |
submit_button = st.form_submit_button('Get Answer') | |
if submit_button and user_query: | |
answer = handle_query(user_query) | |
st.write(answer) | |
elif submit_button and not user_query: | |
st.error("Please enter a question to get an answer.") | |
def handle_query(query): | |
"""Retrieve answers based on the query.""" | |
answer = retrieve_documents(query) | |
return answer | |
if __name__ == "__main__": | |
main() | |