data-science-document-ai 1.42.5__py3-none-any.whl → 1.43.1__py3-none-any.whl

{data_science_document_ai-1.42.5.dist-info → data_science_document_ai-1.43.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-science-document-ai
-Version: 1.42.5
+Version: 1.43.1
 Summary: "Document AI repo for data science"
 Author: Naomi Nguyen
 Author-email: naomi.nguyen@forto.com
@@ -38,7 +38,7 @@ Requires-Dist: pdf2image (>=1.17.0,<2.0.0)
 Requires-Dist: pgzip (>=0.3.5,<0.4.0)
 Requires-Dist: pyarrow (==16.1.0)
 Requires-Dist: pymupdf (>=1.23.26,<2.0.0)
-Requires-Dist: pypdf2 (>=3.0.1,<4.0.0)
+Requires-Dist: pypdf (>=6.1.2,<7.0.0)
 Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
 Requires-Dist: python-multipart (>=0.0.7,<0.0.8)
 Requires-Dist: rapidfuzz (>=3.12.2,<4.0.0)

{data_science_document_ai-1.42.5.dist-info → data_science_document_ai-1.43.1.dist-info}/RECORD

@@ -2,12 +2,12 @@ src/constants.py,sha256=TF_UblovdXZnKIb1lnyJwUqQncJCbzBVihoelI6foSU,3579
 src/constants_sandbox.py,sha256=Iu6HdjCoNSmOX0AwoL9qUQkhq_ZnIN5U9e-Q2UfNuGc,547
 src/docai.py,sha256=dHuR0ehVjUi1CnoNvdp_yxJtpU_HFXqAZ61ywdz7BEo,5655
 src/docai_processor_config.yaml,sha256=qOMmCIORpLQ_D-ytvejXxFvER0e0uGYuzPVdZBGv4Pc,2105
-src/excel_processing.py,sha256=8toKsafUvwE5QN3TOQO3zfLo0Wv2sGxZHKPsL7n5LkA,2771
+src/excel_processing.py,sha256=gzP7QFCp4-n0FTevhWmXm-2UoDF0w0y5v39gsby0IV8,3135
 src/io.py,sha256=tOJpMyI-mP1AaXKG4UFudH47MHWzjWBgVahFJUcjGfs,4749
 src/llm.py,sha256=OE4IEIqcM-hYK9U7e0x1rAfcqdpeo4iXPHBp64L5Qz0,8199
 src/log_setup.py,sha256=RhHnpXqcl-ii4EJzRt47CF2R-Q3YPF68tepg_Kg7tkw,2895
-src/pdf_processing.py,sha256=dxsYvNnONAjzS-T7K5aSo89rz7QcdW3ZDfeuFyeCeII,16294
-src/postprocessing/common.py,sha256=5W-u3lKbnPQRKT4h5EfegegMjSXOKik73X7kUx9ik0Y,21888
+src/pdf_processing.py,sha256=M0RUi20j481FcvMJ3xifavMNsPHqMFpbRo9w8hMpmL8,16970
+src/postprocessing/common.py,sha256=wvlYI1S75r0q5xp9Yll89nOVWtwDd7hV4Sf0MIButA0,22150
 src/postprocessing/postprocess_booking_confirmation.py,sha256=nK32eDiBNbauyQz0oCa9eraysku8aqzrcoRFoWVumDU,4827
 src/postprocessing/postprocess_commercial_invoice.py,sha256=3I8ijluTZcOs_sMnFZxfkAPle0UFQ239EMuvZfDZVPg,1028
 src/postprocessing/postprocess_partner_invoice.py,sha256=koGR7dN37FqJcepdzkrzNBHuBBUuCp_3CrteScASqyE,10590
@@ -51,7 +51,7 @@ src/prompts/library/shippingInstruction/other/prompt.txt,sha256=dT2e-dPuvuz0rVYp
 src/prompts/prompt_library.py,sha256=jPxybNPPGH7mzonqtAOqmw5WcT-RtbGP0pvMqqP22hg,2760
 src/setup.py,sha256=M-p5c8M9ejKcSZ9N86VtmtPc4TYLxe1_4_dxf6jpfVc,7262
 src/tms.py,sha256=UXbIo1QE--hIX6NZi5Qyp2R_CP338syrY9pCTPrfgnE,1741
-src/utils.py,sha256=nU69zR3TB7IZmCc19DD8H27Riek8GJAldmhJjCSwNEE,16090
-data_science_document_ai-1.42.5.dist-info/METADATA,sha256=FauluZfyiueEsYJsiMdiXv7yko2N3Xp5UTe8K0U1Toc,2153
-data_science_document_ai-1.42.5.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-data_science_document_ai-1.42.5.dist-info/RECORD,,
+src/utils.py,sha256=cTF2A12jugKjXxGlNXEZQtfgcsIoaTtaU7zhVOOvXXA,16634
+data_science_document_ai-1.43.1.dist-info/METADATA,sha256=qHuRdJZSrKPDOQvynxf69ssK3jjM-pnBmGyheEPt0Xw,2152
+data_science_document_ai-1.43.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+data_science_document_ai-1.43.1.dist-info/RECORD,,

src/excel_processing.py

@@ -2,6 +2,8 @@
 # flake8: noqa: E402
 import logging
 
+from ddtrace import tracer
+
 from src.postprocessing.common import llm_prediction_to_tuples
 
 logger = logging.getLogger(__name__)
@@ -13,7 +15,7 @@ import numpy as np
 import pandas as pd
 
 from src.llm import prompt_excel_extraction
-from src.utils import generate_schema_structure, get_excel_sheets
+from src.utils import estimate_page_count, generate_schema_structure, get_excel_sheets
 
 
 async def extract_data_from_sheet(
@@ -70,6 +72,15 @@ async def extract_data_from_excel(
     # Load the Excel file and get ONLY the "visible" sheet names
     sheets, workbook = get_excel_sheets(file_content, mime_type)
 
+    # Track the number of sheets in dd-trace
+    span = tracer.current_span()
+    if span:
+        estimated_page_counts = [
+            estimate_page_count(workbook[sheet]) for sheet in sheets
+        ]
+        est_page_count = sum(estimated_page_counts)
+        span.set_metric("est_page_count", est_page_count)
+
     # Excel files may contain multiple sheets. Extract data from each sheet
     sheet_extract_tasks = [
         extract_data_from_sheet(

src/pdf_processing.py

@@ -9,6 +9,7 @@ logger = logging.getLogger(__name__)
 import asyncio
 from collections import defaultdict
 
+from ddtrace import tracer
 from fastapi import HTTPException
 from google.cloud.documentai_v1 import Document as docaiv1_document
 
@@ -32,6 +33,7 @@ from src.prompts.prompt_library import prompt_library
 from src.utils import (
     extract_top_pages,
     generate_schema_structure,
+    get_pdf_page_count,
     get_processor_name,
     run_background_tasks,
     transform_schema_strings,
@@ -411,6 +413,7 @@ async def data_extraction_manual_flow(
     """
     # Get the start time for processing
     start_time = asyncio.get_event_loop().time()
+    page_count = None
     # Validate the file type
     if mime_type == "application/pdf":
         # Enable Doc Ai only for certain document types.
@@ -434,6 +437,7 @@ async def data_extraction_manual_flow(
             if_use_llm=if_use_llm,
             isBetaTest=False,
         )
+        page_count = get_pdf_page_count(file_content)
 
     elif "excel" in mime_type or "spreadsheet" in mime_type:
         # Extract data from the Excel file
@@ -444,6 +448,16 @@ async def data_extraction_manual_flow(
             mime_type=mime_type,
         )
 
+        # Get sheet count from dd-trace span (set in extract_data_from_excel)
+        # Note: we use the span metric instead of len(extracted_data) because
+        # some sheets may fail extraction and not appear in extracted_data
+        span = tracer.current_span()
+        page_count = span.get_metric("est_page_count") if span else len(extracted_data)
+        if page_count > 100:
+            logger.warning(
+                f"Check logic. Count of sheets in excel file is weirdly large: {page_count}"
+            )
+
     else:
         raise HTTPException(
             status_code=400,
@@ -477,6 +491,7 @@ async def data_extraction_manual_flow(
                 processor_version,
                 mime_type,
                 elapsed_time,
+                page_count,
             )
         )
     return result

src/postprocessing/common.py

@@ -319,6 +319,11 @@ def remove_unwanted_patterns(lineitem: str):
     # Remove "HIGH CUBE"
     lineitem = lineitem.replace("HIGH CUBE", "")
 
+    # Remove container size e.g., 20FT, 40HC, etc.
+    lineitem = re.sub(
+        r"\b(20|22|40|45)(FT|HC|DC|HD|GP|OT|RF|FR|TK|DV)?\b", "", lineitem
+    ).strip()
+
     return lineitem
 
 
@@ -349,18 +354,21 @@ def clean_item_description(lineitem: str, remove_numbers: bool = True):
     # Remove the currency codes
     lineitem = re.sub(currency_codes_pattern, "", lineitem, flags=re.IGNORECASE)
 
+    # remove other patterns
+    lineitem = remove_unwanted_patterns(lineitem)
+
     # Remove numbers from the line item
     if (
         remove_numbers
     ):  # Do not remove numbers for the reverse charge sentence as it contains Article number
         lineitem = re.sub(r"\d+", "", lineitem)
 
-    # remove other patterns
-    lineitem = remove_unwanted_patterns(lineitem)
-
     # remove special chars
     lineitem = re.sub(r"[^A-Za-z0-9\s]", " ", lineitem).strip()
 
+    # Remove x from lineitem like 10 x
+    lineitem = re.sub(r"\b[xX]\b", " ", lineitem).strip()
+
     return re.sub(r"\s{2,}", " ", lineitem).strip()
 
 

src/utils.py

@@ -8,15 +8,29 @@ import pickle
 from datetime import datetime
 from typing import Literal
 
+import numpy as np
 import openpyxl
 import pandas as pd
 import requests
 from google.cloud import documentai_v1beta3 as docu_ai_beta
-from PyPDF2 import PdfReader, PdfWriter
+from pypdf import PdfReader, PdfWriter
 
 from src.io import get_storage_client, logger
 
 
+def get_pdf_page_count(pdf_bytes):
+    """Get the number of pages in a PDF document efficiently.
+
+    Args:
+        pdf_bytes (bytes): The PDF content as bytes.
+
+    Returns:
+        int: The number of pages in the PDF.
+    """
+    reader = PdfReader(io.BytesIO(pdf_bytes))
+    return len(reader.pages)
+
+
 def bq_logs(data_to_insert, params):
     """Insert logs into Google BigQuery.
 
@@ -164,6 +178,7 @@ async def run_background_tasks(
     processor_version,
     mime_type,
     elapsed_time=None,
+    page_count=None,
 ):
     """
     Run background tasks asynchronously.
@@ -177,6 +192,7 @@ async def run_background_tasks(
         processor_version: The processor version used to extract the data.
         mime_type: The MIME type of the document.
         elapsed_time: The time taken to process the document.
+        page_count (int, optional): The number of pages in the document.
 
     Returns:
         None
@@ -185,13 +201,8 @@ async def run_background_tasks(
 
     await loop.run_in_executor(None, store_json_in_gcs, params, doc_id, store_data)
 
-    # Keep the page count as 1 for Excel files.
-    page_count = 1
-    # calculate the number of pages processed for PDFs
-    try:
-        if mime_type == "application/pdf":
-            page_count = len(json.loads(store_data.encode("utf-8"))["pages"])
-    except AttributeError:
+    # Use the passed page_count or default to 0 if not provided
+    if page_count is None:
         page_count = 0
 
     # Log the request in BigQuery
@@ -472,3 +483,14 @@ def transform_schema_strings(schema):
     # Base case: for non-dict/list values (e.g., None, bool, str)
     else:
         return schema
+
+
+def estimate_page_count(sheet):
+    """Assuming a page is 10 columns x 50 rows."""
+    if hasattr(sheet, "shape"):
+        pg_cnt = sheet.shape[0] * sheet.shape[1]
+    elif hasattr(sheet, "max_row"):
+        pg_cnt = sheet.max_column * sheet.max_row
+    else:
+        return None
+    return np.ceil(pg_cnt / 500)