Update app.py

app.py

@@ -5,6 +5,7 @@ from PIL import Image
 from transformers import AutoTokenizer, AutoModel
 import torch
 from pdf2image import convert_from_path
+import io
 
 # CSS styles
 css = """
@@ -21,13 +22,11 @@ css = """
 
 # Define layout with custom styles
 layout = [
- gr.Row([gr.File(label="Upload PDF", type="filepath")]),
- gr.Row([gr.Button("Generate Insights")]), # Removed the type="submit"
+ gr.Row([gr.File(label="Upload PDF", type="file")]),
+ gr.Row([gr.Button("Generate Insights")]),
  gr.Row([gr.Textbox("Placeholder for PDF insights", label="Insights", type="text")])
 ]
 
-
-
 # Function to get image embeddings using ViT
 def get_image_embeddings(image_path, model_name='google/vit-base-patch16-224'):
  feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
@@ -64,35 +63,43 @@ def get_text_embeddings(text, model_name='bert-base-uncased'):
 
 # Function to process PDF and generate a response
 def process_pdf_and_generate_response(pdf_file):
- # Convert PDF to images
- img_dir = "pdf_images"
- pdf_to_images(pdf_file, img_dir)
-
- # Generate embeddings for each image
- image_embeddings = []
- for filename in os.listdir(img_dir):
- if filename.endswith(".png"):
- image_path = os.path.join(img_dir, filename)
- image_embeddings.append(get_image_embeddings(image_path))
-
- # Perform some text analysis on the PDF content (replace with your logic)
- pdf_text = "PDF content analysis placeholder"
- text_embeddings = get_text_embeddings(pdf_text)
-
- # Combine image and text embeddings and generate a response (replace with your logic)
- combined_embeddings = torch.cat([*image_embeddings, text_embeddings], dim=0)
- response = "Response based on the processed PDF"
+ try:
+ # Save the uploaded PDF to a temporary file
+ tmp_pdf_path = "/tmp/uploaded_file.pdf"
+ with open(tmp_pdf_path, 'wb') as tmp_pdf:
+ tmp_pdf.write(pdf_file.read())
+
+ # Convert PDF to images
+ img_dir = "pdf_images"
+ pdf_to_images(tmp_pdf_path, img_dir)
+
+ # Generate embeddings for each image
+ image_embeddings = []
+ for filename in os.listdir(img_dir):
+ if filename.endswith(".png"):
+ image_path = os.path.join(img_dir, filename)
+ image_embeddings.append(get_image_embeddings(image_path))
+
+ # Perform some text analysis on the PDF content (replace with your logic)
+ pdf_text = "PDF content analysis placeholder"
+ text_embeddings = get_text_embeddings(pdf_text)
+
+ # Combine image and text embeddings and generate a response (replace with your logic)
+ combined_embeddings = torch.cat([*image_embeddings, text_embeddings], dim=0)
+ response = "Response based on the processed PDF"
+ except Exception as e:
+ response = f"An error occurred: {str(e)}"
  return response
 
 iface = gr.Interface(
  fn=process_pdf_and_generate_response,
- inputs=gr.File(label="Upload PDF", type="filepath"), # Changed 'file' to 'filepath'
+ inputs=gr.File(label="Upload PDF", type="file"),
  outputs=gr.Textbox("Placeholder for PDF insights", label="Insights", type="text"),
  title="pdf-chatbot",
  description="Upload a PDF and receive insights based on its content.",
- css=css # Add the CSS styles here
+ css=css
 )
 
-
 if __name__ == "__main__":
  iface.launch()
+