Merge branch 'main' of https://huggingface.co/spaces/louiecerv/openai_pdf_multimodal

README.md

@@ -11,3 +11,5 @@ short_description: Implement the Multimodal for PDFs
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# replace with a full README later

app.py

@@ -3,7 +3,7 @@ import base64
 import io
 from io import BytesIO
 import tempfile
-import shutil
+import time
 
 import streamlit as st
 from PIL import Image
@@ -12,22 +12,27 @@ import fitz  # PyMuPDF
 from openai import OpenAI
 
 # OpenAI API Key
-api_key = os.getenv("OPENAI_API_KEY")
-
-client = OpenAI(api_key=api_key)
+try:
+    api_key = os.getenv("OPENAI_API_KEY")
+    client = OpenAI(api_key=api_key)
+except Exception as e:
+    st.error(f"An error occurred during OpenAI client initialization: {e}")
+    st.stop()
 
 def extract_text_and_images_from_pdf(pdf_file_path):
     try:
         text_content = ""
         image_urls = []
 
+        pdf_stream = BytesIO(pdf_file.read())
+
         # Extract text using PdfReader
-        pdf_reader = PdfReader(pdf_file_path)
+        pdf_reader = PdfReader(pdf_stream)
         for page in pdf_reader.pages:
             text_content += page.extract_text()
 
         # Extract images using PyMuPDF
-        doc = fitz.open(pdf_file_path)
+        doc = fitz.open(stream=pdf_stream)
         for page_index in range(len(doc)):
             page = doc.load_page(page_index)
             image_list = page.get_images()
@@ -39,32 +44,33 @@ def extract_text_and_images_from_pdf(pdf_file_path):
                 # Resize image (optional)
                 image.thumbnail((512, 512))  # Adjust size as needed
 
-                # Encode the image as base64 and create a data URL
-                buffered = io.BytesIO()
-                image.save(buffered, format="JPEG")
-                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-                data_url = f"data:image/jpeg;base64,{img_str}"
-                image_urls.append(data_url)
+                    # Encode the image as base64 and create a data URL
+                    with open(image_path, "rb") as f:
+                        img_str = base64.b64encode(f.read()).decode("utf-8")
+                    data_url = f"data:image/jpeg;base64,{img_str}"
+                    image_urls.append(data_url)
+
+            # Close the fitz document
+            doc.close()
 
         return text_content, image_urls
     except Exception as e:
         st.error(f"An error occurred during PDF processing: {e}")
         return "", []
 
+
 def generate_ai_response(text_content, image_urls, text_prompt):
     try:
-
-        if len(image_urls) > 0:    
-            # Construct the messages list with the prompt and base64-encoded image URLs
-            messages = [
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "text", "text": f"Perform this task {text_prompt} on the images:"},
-                        *[{"type": "image_url", "image_url": {"url": url}} for url in image_urls]
-                    ]
-                }
-            ]
+        # Construct the messages list with the prompt and base64-encoded image URLs
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": text_prompt},
+                    *[{"type": "image_url", "image_url": {"url": url}} for url in image_urls]
+                ]
+            }
+        ]
 
         else:
             # Construct the prompt on the extracted text only
@@ -89,6 +95,7 @@ def generate_ai_response(text_content, image_urls, text_prompt):
         st.error(f"An error occurred during AI response generation: {e}")
         return ""
 
+
 def main():
     st.title("Multimodal PDF Processing using GPT-4 Turbo Model")
 
@@ -105,13 +112,7 @@ def main():
 
     uploaded_pdf = st.file_uploader("Upload a PDF", type=["pdf"])
     if uploaded_pdf is not None:
-        # Save the uploaded PDF to a temporary directory
-        temp_dir = tempfile.mkdtemp()
-        pdf_file_path = os.path.join(temp_dir, uploaded_pdf.name)
-        with open(pdf_file_path, "wb") as f:
-            f.write(uploaded_pdf.getvalue())
-
-        text_content, image_urls = extract_text_and_images_from_pdf(pdf_file_path)
+        text_content, image_urls = extract_text_and_images_from_pdf(uploaded_pdf)
 
         st.subheader("Extracted Text")
         st.text(text_content)
@@ -121,16 +122,17 @@ def main():
         if image_urls:
             st.subheader("Extracted Images")
             for img_url in image_urls:
-                st.image(img_url, caption="Extracted Image", use_container_width=True)
+                st.image(
+                    img_url, caption="Extracted Image", use_container_width=True
+                )
 
         if st.button("Generate Response"):
             with st.spinner("Processing..."):
-                ai_response = generate_ai_response(text_content, image_urls, text_prompt)
+                ai_response = generate_ai_response(
+                    text_content, image_urls, text_prompt
+                )
                 st.success("Response generated!")
                 st.markdown(f"AI Response: {ai_response}")
 
-        # Clean up the temporary directory
-        shutil.rmtree(temp_dir)
-
 if __name__ == "__main__":
     main()