nv-ingest-api 26.1.0rc4__py3-none-any.whl

nv_ingest_api/internal/extract/audio/audio_extraction.py

@@ -0,0 +1,202 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import pandas as pd
+import functools
+import uuid
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Tuple
+import base64
+from pathlib import Path
+
+from nv_ingest_api.internal.enums.common import ContentTypeEnum
+from nv_ingest_api.internal.primitives.nim.model_interface.parakeet import create_audio_inference_client
+from nv_ingest_api.internal.schemas.extract.extract_audio_schema import AudioExtractorSchema
+from nv_ingest_api.internal.schemas.meta.metadata_schema import MetadataSchema, AudioMetadataSchema
+from nv_ingest_api.util.exception_handlers.decorators import unified_exception_handler
+from nv_ingest_api.util.schema.schema_validator import validate_schema
+from nv_ingest_api.interface.utility import read_file_as_base64
+
+logger = logging.getLogger(__name__)
+
+
+@unified_exception_handler
+def _extract_from_audio(row: pd.Series, audio_client: Any, trace_info: Dict, segment_audio: bool = False) -> Dict:
+    """
+    Modifies the metadata of a row if the conditions for table extraction are met.
+
+    Parameters
+    ----------
+    row : pd.Series
+        A row from the DataFrame containing metadata for the audio extraction.
+
+    audio_client : Any
+        The client used to call the audio inference model.
+
+    trace_info : Dict
+        Trace information used for logging or debugging.
+
+    Returns
+    -------
+    Dict
+        The modified metadata if conditions are met, otherwise the original metadata.
+
+    Raises
+    ------
+    ValueError
+        If critical information (such as metadata) is missing from the row.
+    """
+
+    metadata = row.get("metadata")
+
+    if metadata is None:
+        logger.error("Row does not contain 'metadata'.")
+        raise ValueError("Row does not contain 'metadata'.")
+
+    base64_audio = metadata.pop("content")
+    try:
+        base64_file_path = base64_audio
+        if not base64_file_path:
+            return [row.to_list()]
+        base64_file_path = base64.b64decode(base64_file_path).decode("utf-8")
+        if not base64_file_path:
+            return [row.to_list()]
+        if Path(base64_file_path).exists():
+            base64_audio = read_file_as_base64(base64_file_path)
+    except (UnicodeDecodeError, base64.binascii.Error):
+        pass
+    content_metadata = metadata.get("content_metadata", {})
+
+    # Only extract transcript if content type is audio
+    if (content_metadata.get("type") != ContentTypeEnum.AUDIO) or (base64_audio in (None, "")):
+        return [row.to_list()]
+
+    logger.debug(f"Removing file {base64_file_path}")
+    Path(base64_file_path).unlink(missing_ok=True)
+
+    # Get the result from the inference model
+    segments, transcript = audio_client.infer(
+        base64_audio,
+        model_name="parakeet",
+        trace_info=trace_info,  # traceable_func arg
+        stage_name="audio_extraction",
+    )
+
+    extracted_data = []
+    if segment_audio:
+        for segment in segments:
+            segment_metadata = metadata.copy()
+            audio_metadata = {"audio_transcript": segment["text"]}
+            segment_metadata["audio_metadata"] = validate_schema(audio_metadata, AudioMetadataSchema).model_dump()
+            segment_metadata["content_metadata"]["start_time"] = segment["start"]
+            segment_metadata["content_metadata"]["end_time"] = segment["end"]
+
+            extracted_data.append(
+                [
+                    ContentTypeEnum.AUDIO,
+                    validate_schema(segment_metadata, MetadataSchema).model_dump(),
+                    str(uuid.uuid4()),
+                ]
+            )
+    else:
+        audio_metadata = {"audio_transcript": transcript}
+        metadata["audio_metadata"] = validate_schema(audio_metadata, AudioMetadataSchema).model_dump()
+        extracted_data.append(
+            [ContentTypeEnum.AUDIO, validate_schema(metadata, MetadataSchema).model_dump(), str(uuid.uuid4())]
+        )
+
+    return extracted_data
+
+
+def extract_text_from_audio_internal(
+    df_extraction_ledger: pd.DataFrame,
+    task_config: Dict[str, Any],
+    extraction_config: AudioExtractorSchema,
+    execution_trace_log: Optional[Dict] = None,
+) -> Tuple[pd.DataFrame, Dict]:
+    """
+    Extracts audio data from a DataFrame.
+
+    Parameters
+    ----------
+    df_extraction_ledger : pd.DataFrame
+        DataFrame containing the content from which audio data is to be extracted.
+
+    task_config : Dict[str, Any]
+        Dictionary containing task properties and configurations.
+
+    extraction_config : Any
+        The validated configuration object for audio extraction.
+
+    execution_trace_log : Optional[Dict], optional
+        Optional trace information for debugging or logging. Defaults to None.
+
+    Returns
+    -------
+    Tuple[pd.DataFrame, Dict]
+        A tuple containing the updated DataFrame and the trace information.
+
+    Raises
+    ------
+    Exception
+        If any error occurs during the audio data extraction process.
+    """
+    logger.debug(f"Entering audio extraction stage with {len(df_extraction_ledger)} rows.")
+
+    extract_params = task_config.get("params", {}).get("extract_audio_params", {})
+    audio_extraction_config = extraction_config.audio_extraction_config
+
+    grpc_endpoint = extract_params.get("grpc_endpoint") or audio_extraction_config.audio_endpoints[0]
+    http_endpoint = extract_params.get("http_endpoint") or audio_extraction_config.audio_endpoints[1]
+    infer_protocol = extract_params.get("infer_protocol") or audio_extraction_config.audio_infer_protocol
+    auth_token = extract_params.get("auth_token") or audio_extraction_config.auth_token
+    function_id = extract_params.get("function_id") or audio_extraction_config.function_id
+    use_ssl = extract_params.get("use_ssl") or audio_extraction_config.use_ssl
+    ssl_cert = extract_params.get("ssl_cert") or audio_extraction_config.ssl_cert
+    segment_audio = extract_params.get("segment_audio") or audio_extraction_config.segment_audio
+
+    parakeet_client = create_audio_inference_client(
+        (grpc_endpoint, http_endpoint),
+        infer_protocol=infer_protocol,
+        auth_token=auth_token,
+        function_id=function_id,
+        use_ssl=use_ssl,
+        ssl_cert=ssl_cert,
+    )
+
+    if execution_trace_log is None:
+        execution_trace_log = {}
+        logger.debug("No trace_info provided. Initialized empty trace_info dictionary.")
+
+    try:
+        # Create a partial function to extract using the provided configurations.
+        _extract_from_audio_partial = functools.partial(
+            _extract_from_audio,
+            audio_client=parakeet_client,
+            trace_info=execution_trace_log,
+            segment_audio=segment_audio,
+        )
+
+        # Apply the _extract_from_audio_partial function to each row in the DataFrame
+        extraction_series = df_extraction_ledger.apply(_extract_from_audio_partial, axis=1)
+
+        # Explode the results if the extraction returns lists.
+        extraction_series = extraction_series.explode().dropna()
+
+        # Convert the extracted results into a DataFrame.
+        if not extraction_series.empty:
+            extracted_df = pd.DataFrame(extraction_series.to_list(), columns=["document_type", "metadata", "uuid"])
+        else:
+            extracted_df = pd.DataFrame({"document_type": [], "metadata": [], "uuid": []})
+
+        return extracted_df, execution_trace_log
+
+    except Exception as e:
+        logger.exception(f"Error occurred while extracting audio data: {e}", exc_info=True)
+
+        raise

nv_ingest_api/internal/extract/docx/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright (c) 2024, NVIDIA CORPORATION.

nv_ingest_api/internal/extract/docx/docx_extractor.py

@@ -0,0 +1,232 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import base64
+import functools
+import io
+import logging
+from typing import Optional, Dict, Any, Union, Tuple
+
+import pandas as pd
+from pydantic import BaseModel
+
+from nv_ingest_api.internal.extract.docx.engines.docxreader_helpers.docx_helper import python_docx
+from nv_ingest_api.internal.extract.pdf.engines.pdfium import pdfium_extractor
+from nv_ingest_api.internal.extract.pptx.engines.pptx_helper import convert_stream_with_libreoffice
+from nv_ingest_api.internal.schemas.extract.extract_docx_schema import DocxExtractorSchema
+from nv_ingest_api.util.exception_handlers.decorators import unified_exception_handler
+
+
+logger = logging.getLogger(__name__)
+
+
+def _prepare_task_props(
+    task_config: Union[Dict[str, Any], BaseModel], base64_row: pd.Series
+) -> (Dict[str, Any], Optional[str]):
+    """
+    Prepares the task properties by converting a Pydantic model to a dictionary (if needed)
+    and injecting row-specific data.
+
+    Parameters
+    ----------
+    task_config : Union[Dict[str, Any], BaseModel]
+        A dictionary or Pydantic model containing instructions and parameters for extraction.
+    base64_row : pd.Series
+        A Series representing a row from the DataFrame that contains at least the "content"
+        key and optionally "source_id".
+
+    Returns
+    -------
+    Tuple[Dict[str, Any], Optional[str]]
+        A tuple where the first element is the prepared task properties dictionary with the key
+        "row_data" added under its "params" key, and the second element is the source_id (if present),
+        otherwise None.
+    """
+
+    if isinstance(task_config, BaseModel):
+        task_config = task_config.model_dump()
+    else:
+        task_config = dict(task_config)
+
+    # Extract all row data except the "content" field.
+    row_data = base64_row.drop(labels=["content"], errors="ignore")
+    if "params" not in task_config:
+        task_config["params"] = {}
+
+    task_config["params"]["row_data"] = row_data
+
+    source_id = base64_row.get("source_id", None)
+
+    return task_config, source_id
+
+
+@unified_exception_handler
+def _decode_and_extract_from_docx(
+    base64_row: pd.Series,
+    task_config: Union[Dict[str, Any], BaseModel],
+    extraction_config: Any,
+    execution_trace_log: Optional[Dict[str, Any]] = None,
+) -> Any:
+    """
+    Decodes base64 content from a DataFrame row and extracts data using the specified extraction method.
+
+    The function decodes the base64-encoded content from the "content" key in the row, prepares
+    extraction parameters (including additional row data and configuration), and invokes the extraction
+    function from the docx module. If an error occurs, an exception tag is returned.
+
+    Parameters
+    ----------
+    base64_row : pd.Series
+        A Series containing the base64-encoded content under the key "content" and optionally a "source_id".
+    task_config : Union[Dict[str, Any], BaseModel]
+        A dictionary or Pydantic model containing extraction instructions and parameters.
+        Expected to have a "params" key for additional parameters and optionally a "method" key specifying
+        the extraction method.
+    extraction_config : Any
+        A configuration object that contains extraction-specific settings, such as `docx_extraction_config`.
+    execution_trace_log : Optional[Dict[str, Any]], default=None
+        A dictionary containing trace information for debugging or logging.
+    default : str, optional
+        The default extraction method to use if the specified method is not available (default is "python_docx").
+
+    Returns
+    -------
+    Any
+        The extracted data, or an exception tag if extraction fails.
+
+    Raises
+    ------
+    Exception
+        If an unhandled exception occurs during extraction, it is logged and a tagged error is returned.
+    """
+    # Prepare task properties and extract source_id
+    task_config, source_id = _prepare_task_props(task_config, base64_row)
+
+    # Retrieve base64 content and decode it into a byte stream.
+    base64_content: str = base64_row["content"]
+    doc_bytes: bytes = base64.b64decode(base64_content)
+    doc_stream: io.BytesIO = io.BytesIO(doc_bytes)
+
+    extract_method = task_config.get("method", "python_docx")
+    extract_params: Dict[str, Any] = task_config.get("params", {})
+
+    # Extract required boolean flags from params.
+    try:
+        extract_text = extract_params.pop("extract_text", False)
+        extract_images = extract_params.pop("extract_images", False)
+        extract_tables = extract_params.pop("extract_tables", False)
+        extract_charts = extract_params.pop("extract_charts", False)
+        extract_infographics = extract_params.pop("extract_infographics", False)
+    except KeyError as e:
+        raise ValueError(f"Missing required extraction flag: {e}")
+
+    # Inject configuration and trace info into extraction parameters.
+    if getattr(extraction_config, "docx_extraction_config", None) is not None:
+        extract_params["docx_extraction_config"] = extraction_config.docx_extraction_config
+
+    if execution_trace_log is not None:
+        extract_params["trace_info"] = execution_trace_log
+
+    if extract_method == "render_as_pdf":
+        pdf_stream = convert_stream_with_libreoffice(doc_stream, "docx", "pdf")
+
+        pdf_extract_method = extract_params.get("pdf_extract_method", "pdfium")
+        pdf_extractor_config = extract_params.copy()
+        pdf_extractor_config["extract_method"] = pdf_extract_method
+        if getattr(extraction_config, "pdfium_config", None) is not None:
+            pdf_extractor_config["pdfium_config"] = extraction_config.pdfium_config
+
+        extracted_data: Any = pdfium_extractor(
+            pdf_stream=pdf_stream,
+            extract_text=extract_text,
+            extract_images=extract_images,
+            extract_infographics=extract_infographics,
+            extract_tables=extract_tables,
+            extract_charts=extract_charts,
+            extract_page_as_image=False,
+            extractor_config=pdf_extractor_config,
+            execution_trace_log=None,
+        )
+
+    elif extract_method == "python_docx":
+        extracted_data: Any = python_docx(
+            docx_stream=doc_stream,
+            extract_text=extract_text,
+            extract_images=extract_images,
+            extract_infographics=extract_infographics,
+            extract_tables=extract_tables,
+            extract_charts=extract_charts,
+            extraction_config=extract_params,
+            execution_trace_log=None,
+        )
+    else:
+        raise ValueError(f"Unsupported DOCX extraction method: {extract_method}")
+
+    return extracted_data
+
+
+@unified_exception_handler
+def extract_primitives_from_docx_internal(
+    df_extraction_ledger: pd.DataFrame,
+    task_config: Union[Dict[str, Any], BaseModel],
+    extraction_config: DocxExtractorSchema,
+    execution_trace_log: Optional[Dict[str, Any]] = None,
+) -> Tuple[pd.DataFrame, Union[Dict, None]]:
+    """
+    Processes a pandas DataFrame containing DOCX files encoded in base64, extracting text from
+    each document and replacing the original content with the extracted text.
+
+    This function applies a decoding and extraction routine to each row of the input DataFrame.
+    The routine is provided via the `decode_and_extract` function, which is partially applied with
+    task configuration, extraction configuration, and optional trace information. The results are
+    exploded and any missing values are dropped, then compiled into a new DataFrame with columns
+    for document type, metadata, and a UUID identifier.
+
+    Parameters
+    ----------
+    df_extraction_ledger : pd.DataFrame
+        The input DataFrame containing DOCX files in base64 encoding. Expected columns include
+        'source_id' and 'content'.
+    task_config : Union[Dict[str, Any], BaseModel]
+        Configuration instructions for the document processing task. This can be provided as a
+        dictionary or a Pydantic model.
+    extraction_config : Any
+        A configuration object for document extraction that guides the extraction process.
+    execution_trace_log : Optional[Dict[str, Any]], default=None
+        An optional dictionary containing trace information for debugging or logging.
+
+    Returns
+    -------
+    pd.DataFrame
+        A DataFrame with the original DOCX content replaced by the extracted text. The resulting
+        DataFrame contains the columns "document_type", "metadata", and "uuid".
+
+    Raises
+    ------
+    Exception
+        If an error occurs during the document extraction process, the exception is logged and
+        re-raised.
+    """
+    # Create a partial function to decode and extract using the provided configurations.
+    _decode_and_extract = functools.partial(
+        _decode_and_extract_from_docx,
+        task_config=task_config,
+        extraction_config=extraction_config,
+        execution_trace_log=execution_trace_log,
+    )
+
+    # Apply the decode_and_extract function to each row in the DataFrame.
+    sr_extraction = df_extraction_ledger.apply(_decode_and_extract, axis=1)
+
+    # Explode any list results and drop missing values.
+    sr_extraction = sr_extraction.explode().dropna()
+
+    # Convert the extraction results to a DataFrame if available.
+    if not sr_extraction.empty:
+        extracted_df = pd.DataFrame(sr_extraction.to_list(), columns=["document_type", "metadata", "uuid"])
+    else:
+        extracted_df = pd.DataFrame({"document_type": [], "metadata": [], "uuid": []})
+
+    return extracted_df, {}

nv_ingest_api/internal/extract/docx/engines/docxreader_helpers/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

nv_ingest_api/internal/extract/docx/engines/docxreader_helpers/docx_helper.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2024, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=too-many-locals
+
+import logging
+from typing import IO, Optional, List
+
+from nv_ingest_api.internal.enums.common import AccessLevelEnum, DocumentTypeEnum
+from nv_ingest_api.internal.enums.common import TextTypeEnum
+from nv_ingest_api.internal.extract.docx.engines.docxreader_helpers.docxreader import DocxReader
+
+logger = logging.getLogger(__name__)
+
+
+def python_docx(
+    *,
+    docx_stream: IO,
+    extract_text: bool,
+    extract_images: bool,
+    extract_infographics: bool,
+    extract_tables: bool,
+    extract_charts: bool,
+    extraction_config: dict,
+    execution_trace_log: Optional[List] = None,
+):
+    """
+    Helper function that use python-docx to extract text from a bytestream document
+
+    A document has three levels - document, paragraphs and runs. To align with the
+    pdf extraction paragraphs are aliased as block. python-docx leaves the page number
+    and line number to the renderer so we assume that the entire document is a single
+    page.
+
+    Run level parsing has been skipped but can be added as needed.
+
+    Parameters
+    ----------
+    docx_stream:
+        Bytestream
+    extract_text : bool
+        Specifies whether to extract text.
+    extract_images : bool
+        Specifies whether to extract images.
+    extract_infographics : bool
+        Specifies whether to extract infographics.
+    extract_tables : bool
+        Specifies whether to extract tables.
+    extract_charts : bool
+        Specifies whether to extract charts.
+    extraction_config : dict
+        A dictionary of configuration parameters for the extraction process.
+    execution_trace_log : list, optional
+        A list for accumulating trace information during extraction. Defaults to None.
+
+    Returns
+    -------
+    str
+        A string of extracted text.
+    """
+
+    _ = execution_trace_log
+    _ = extract_infographics
+
+    row_data = extraction_config.get("row_data")
+    # get source_id
+    source_id = row_data["source_id"]
+    # get text_depth
+    text_depth = extraction_config.get("text_depth", "document")
+    text_depth = TextTypeEnum(text_depth)
+    # get base metadata
+    metadata_col = "metadata"
+
+    docx_extractor_config = extraction_config.get("docx_extraction_config", {})
+
+    base_unified_metadata = row_data[metadata_col] if metadata_col in row_data.index else {}
+
+    # get base source_metadata
+    base_source_metadata = base_unified_metadata.get("source_metadata", {})
+    # get source_location
+    source_location = base_source_metadata.get("source_location", "")
+    # get collection_id (assuming coming in from source_metadata...)
+    collection_id = base_source_metadata.get("collection_id", "")
+    # get partition_id (assuming coming in from source_metadata...)
+    partition_id = base_source_metadata.get("partition_id", -1)
+    # get access_level (assuming coming in from source_metadata...)
+    access_level = base_source_metadata.get("access_level", AccessLevelEnum.UNKNOWN)
+
+    # python-docx doesn't maintain filename; re-use source_id
+    source_metadata = {
+        "source_name": source_id,
+        "source_id": source_id,
+        "source_location": source_location,
+        "source_type": DocumentTypeEnum.DOCX,
+        "collection_id": collection_id,
+        "partition_id": partition_id,
+        "access_level": access_level,
+        "summary": "",
+    }
+
+    # Extract data from the document using python-docx
+    doc = DocxReader(docx_stream, source_metadata, extraction_config=docx_extractor_config)
+    extracted_data = doc.extract_data(
+        base_unified_metadata,
+        text_depth=text_depth,
+        extract_text=extract_text,
+        extract_tables=extract_tables,
+        extract_charts=extract_charts,
+        extract_infographics=extract_infographics,
+        extract_images=extract_images,
+    )
+
+    return extracted_data