Update app.py

app.py

@@ -30,9 +30,6 @@ MONGODB_URI = os.getenv("MONGODB_URI")
 DATABASE_NAME = os.getenv("DATABASE_NAME")
 COLLECTION_NAME = os.getenv("COLLECTION_NAME", "invoice_collection")
 
-# use_gpu = False
-# output_dir = 'output'
-
 # Check if environment variables are set
 if not MONGODB_URI:
     raise ValueError("MONGODB_URL is not set. Please add it to Hugging Face secrets.")
@@ -80,48 +77,36 @@ def fetch_file_from_s3(file_key):
     except Exception as e:
         raise Exception(f"Failed to fetch file from S3: {str(e)}")
 
-# def extract_text_from_pdf(file_data):
-#     """Extracts text from a PDF file."""
-#     try:
-#         doc = fitz.open(stream=file_data, filetype="pdf")
-#         text = "\n".join([page.get_text("text") for page in doc])
-#         return text.strip()
-#     except Exception as e:
-#         logger.error(f"Failed to extract text from PDF: {e}")
-#         return ""
+def extract_text_from_pdf(file_data):
+    try:
+        doc = fitz.open(stream=file_data, filetype="pdf")
+        return "\n".join([page.get_text("text") for page in doc]).strip()
+    except Exception as e:
+        logger.error(f"Failed to extract text from PDF: {e}")
+        return ""
 
 # Function to summarize text using OpenAI GPT
 def extract_invoice_data(file_data, content_type):
-    system_prompt = "You are an expert in document data extraction."
-
-    # Convert file to Base64
+    if content_type == "application/pdf":
+        text = extract_text_from_pdf(file_data)
+        if len(text.split()) > 500:  # Large document handling
+            return {"extracted_text": text}
+    
     base64_encoded = base64.b64encode(file_data).decode('utf-8')
-
-    # Determine the correct MIME type for OpenAI
-    if content_type.startswith("image/"):
-        mime_type = content_type  # e.g., image/png, image/jpeg
-    elif content_type == "application/pdf":
-        mime_type = "application/pdf"
-        # text = extract_text_from_pdf(file_data)
-        # mime_type = [{"role": "user", "content": text}]
-    else:
-        raise ValueError(f"Unsupported content type: {content_type}")
-
+    mime_type = content_type if content_type.startswith("image/") else "application/pdf"
+    
     try:
         response = openai.ChatCompletion.create(
             model="gpt-4o-mini",
             messages=[
-                {"role": "system", "content": system_prompt},
+                {"role": "system", "content": "You are an expert in document data extraction."},
                 {
                     "role": "user",
                     "content": [
                         {
                             "type": "image_url",
                             "image_url": {
-                                "url": f"data:image/{mime_type};base64,{base64_encoded}"
-                            },
-                            "image_url": {
-                                "url": f"data:application/pdf;base64,{base64_encoded}"
+                                "url": f"data:{mime_type};base64,{base64_encoded}"
                             }
                         }
                     ]