nv-ingest-api 2025.4.20.dev20250420__py3-none-any.whl → 2025.4.22.dev20250422__py3-none-any.whl

nv_ingest_api/internal/transform/split_text.py

@@ -0,0 +1,157 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import os
+import copy
+import logging
+import uuid
+from typing import Any, Optional, Dict
+from typing import List
+
+import pandas as pd
+from transformers import AutoTokenizer
+
+from nv_ingest_api.internal.enums.common import ContentTypeEnum
+from nv_ingest_api.internal.schemas.transform.transform_text_splitter_schema import TextSplitterSchema
+from nv_ingest_api.util.exception_handlers.decorators import unified_exception_handler
+
+logger = logging.getLogger(__name__)
+
+
+def _build_split_documents(row, chunks: List[str]) -> List[dict[str, Any]]:
+    """Build documents from text chunks"""
+    documents: List[dict] = []
+
+    for i, text in enumerate(chunks):
+        if text is None or not text.strip():
+            continue
+
+        metadata = row.metadata if hasattr(row, "metadata") and isinstance(row.metadata, dict) else {}
+        metadata = copy.deepcopy(metadata)
+
+        metadata["content"] = text
+
+        documents.append({"document_type": ContentTypeEnum.TEXT.value, "metadata": metadata, "uuid": str(uuid.uuid4())})
+
+    return documents
+
+
+def _split_into_chunks(text, tokenizer, chunk_size=1024, chunk_overlap=20):
+    # Tokenize the text into token IDs
+    encoding = tokenizer.encode_plus(text, add_special_tokens=False, return_offsets_mapping=True)
+
+    # Get the token IDs and offsets for splitting
+    offsets = encoding["offset_mapping"]
+
+    # Split the tokens into chunks of the desired size with the desired overlap
+    chunks = [offsets[i : i + chunk_size] for i in range(0, len(offsets), chunk_size - chunk_overlap)]
+
+    # Convert token chunks back to text while preserving original spacing and case
+    text_chunks = []
+    for chunk in chunks:
+        text_chunk = text[chunk[0][0] : chunk[-1][0]]
+        text_chunks.append(text_chunk)
+
+    return text_chunks
+
+
+@unified_exception_handler
+def transform_text_split_and_tokenize_internal(
+    df_transform_ledger: pd.DataFrame,
+    task_config: Dict[str, Any],
+    transform_config: TextSplitterSchema,
+    execution_trace_log: Optional[Dict[str, Any]],
+) -> pd.DataFrame:
+    """
+    Internal function to split and tokenize text in a ledger DataFrame.
+
+    This function extracts text from documents that match a filter criteria based on source types,
+    splits the text into chunks using the specified tokenizer, and rebuilds document records with the
+    split text. The resulting DataFrame contains both split and unsplit documents.
+
+    Parameters
+    ----------
+    df_transform_ledger : pd.DataFrame
+        DataFrame containing documents to be processed. Expected to have columns 'document_type' and
+        'metadata', where 'metadata' includes a 'content' field and nested source information.
+    task_config : dict
+        Dictionary with task-specific configuration. Expected keys include:
+            - "tokenizer": Tokenizer identifier or path.
+            - "chunk_size": Maximum number of tokens per chunk.
+            - "chunk_overlap": Number of tokens to overlap between chunks.
+            - "params": A sub-dictionary that may contain:
+                - "hf_access_token": Hugging Face access token.
+                - "split_source_types": List of source types to filter for splitting.
+    transform_config : TextSplitterSchema
+        Configuration object providing default values for text splitting parameters.
+    execution_trace_log : Optional[dict]
+        Optional dictionary for logging execution trace information; may be None.
+
+    Returns
+    -------
+    pd.DataFrame
+        DataFrame with processed documents. Documents with text matching the filter are split into chunks,
+        and then merged with those that do not match the filter.
+
+    Raises
+    ------
+    ValueError
+        If the text splitting or tokenization process fails.
+    """
+    _ = execution_trace_log  # Placeholder for potential execution trace logging.
+
+    # Override parameters using task_config, with fallback to transform_config.
+    tokenizer_identifier: Optional[str] = task_config.get("tokenizer", transform_config.tokenizer)
+    chunk_size: int = task_config.get("chunk_size", transform_config.chunk_size)
+    chunk_overlap: int = task_config.get("chunk_overlap", transform_config.chunk_overlap)
+    params: Dict[str, Any] = task_config.get("params", {})
+
+    hf_access_token: Optional[str] = params.get("hf_access_token", None)
+    split_source_types: List[str] = params.get("split_source_types", ["text"])
+
+    logger.debug(
+        f"Splitting text with tokenizer: {tokenizer_identifier}, "
+        f"chunk_size: {chunk_size} tokens, "
+        f"chunk_overlap: {chunk_overlap}"
+    )
+
+    # Filter to documents with text content.
+    bool_index = (df_transform_ledger["document_type"] == ContentTypeEnum.TEXT) & (
+        pd.json_normalize(df_transform_ledger["metadata"])["source_metadata.source_type"].isin(split_source_types)
+    )
+    df_filtered: pd.DataFrame = df_transform_ledger.loc[bool_index]
+
+    if df_filtered.empty:
+        return df_transform_ledger
+
+    model_predownload_path = os.environ.get("MODEL_PREDOWNLOAD_PATH")
+
+    if os.path.exists(os.path.join(model_predownload_path, "llama-3.2-1b/tokenizer/tokenizer.json")) and (
+        tokenizer_identifier is None or tokenizer_identifier == "meta-llama/Llama-3.2-1B"
+    ):
+        tokenizer_identifier = os.path.join(model_predownload_path, "llama-3.2-1b/tokenizer/")
+    elif os.path.exists(os.path.join(model_predownload_path, "e5-large-unsupervised/tokenizer/tokenizer.json")) and (
+        tokenizer_identifier is None or tokenizer_identifier == "intfloat/e5-large-unsupervised"
+    ):
+        tokenizer_identifier = os.path.join(model_predownload_path, "e5-large-unsupervised/tokenizer/")
+
+    tokenizer_model = AutoTokenizer.from_pretrained(tokenizer_identifier, token=hf_access_token)
+
+    split_docs: List[Dict[str, Any]] = []
+    for _, row in df_filtered.iterrows():
+        content: str = row["metadata"]["content"] if row["metadata"]["content"] is not None else ""
+        chunks: List[str] = _split_into_chunks(content, tokenizer_model, chunk_size, chunk_overlap)
+        split_docs.extend(_build_split_documents(row, chunks))
+
+    split_docs_df: pd.DataFrame = pd.DataFrame(split_docs)
+
+    # Merge split documents with unsplit documents.
+    merged_df: pd.DataFrame = pd.concat([split_docs_df, df_transform_ledger[~bool_index]], axis=0).reset_index(
+        drop=True
+    )
+
+    result, execution_trace_log = merged_df, {}
+
+    return result

nv_ingest_api/util/control_message/validators.py

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage
+
+
+def cm_ensure_payload_not_null(control_message: IngestControlMessage):
+    """
+    Ensures that the payload of a IngestControlMessage is not None.
+
+    Parameters
+    ----------
+    control_message : IngestControlMessage
+        The IngestControlMessage to check.
+
+    Raises
+    ------
+    ValueError
+        If the payload is None.
+    """
+
+    if control_message.payload() is None:
+        raise ValueError("Payload cannot be None")
+
+
+def cm_set_failure(control_message: IngestControlMessage, reason: str) -> IngestControlMessage:
+    """
+    Sets the failure metadata on a IngestControlMessage.
+
+    Parameters
+    ----------
+    control_message : IngestControlMessage
+        The IngestControlMessage to set the failure metadata on.
+    reason : str
+        The reason for the failure.
+
+    Returns
+    -------
+    control_message : IngestControlMessage
+        The modified IngestControlMessage with the failure metadata set.
+    """
+
+    control_message.set_metadata("cm_failed", True)
+    control_message.set_metadata("cm_failed_reason", reason)
+
+    return control_message

nv_ingest_api/util/converters/bytetools.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import base64
+
+
+def bytesfromhex(hex_input):
+    """
+    Function to convert hex to bytes.
+
+    Parameters
+    ----------
+    hex_input : hex
+        Hex string to store bytes in cuDF.
+
+    Returns
+    -------
+    bytes
+        Hex encoded object converted to bytes.
+    """
+
+    return bytes.fromhex(hex_input)
+
+
+def hexfrombytes(bytes_input):
+    """
+    Function to bytes to hex string.
+
+    Parameters
+    ----------
+    bytes_input : bytes
+        Raw bytes of object.
+
+    Returns
+    -------
+    hex
+        Hex string to store bytes in cuDF.
+    """
+
+    return bytes_input.hex()
+
+
+def bytesfrombase64(base64_input):
+    """
+    Function to convert base64 encoded string to bytes.
+
+    Parameters
+    ----------
+    base64_input : hex
+        Base64 encoded string to store bytes in cuDF.
+
+    Returns
+    -------
+    bytes
+        Base64 encoded string converted to bytes.
+    """
+
+    return base64.b64decode(base64_input)
+
+
+def base64frombytes(bytes_input, encoding="utf-8"):
+    """
+    Function to bytes to base64 string.
+
+    Parameters
+    ----------
+    bytes_input : bytes
+        Raw bytes of object.
+
+    Returns
+    -------
+    base64
+        base64 encoded string to store bytes in cuDF.
+    """
+
+    return base64.b64encode(bytes_input).decode(encoding)

nv_ingest_api/util/converters/containers.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from typing import Any
+from typing import Dict
+
+logger = logging.getLogger(__name__)
+
+
+def merge_dict(defaults: Dict[str, Any], overrides: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Recursively merges two dictionaries, with values from the `overrides` dictionary taking precedence.
+
+    This function merges the `overrides` dictionary into the `defaults` dictionary. If a key in both dictionaries
+    has a dictionary as its value, the function will recursively merge those dictionaries. Otherwise, the value
+    from the `overrides` dictionary will overwrite the value in the `defaults` dictionary.
+
+    Parameters
+    ----------
+    defaults : dict of {str: Any}
+        The default dictionary that will be updated with values from the `overrides` dictionary.
+    overrides : dict of {str: Any}
+        The dictionary containing values that will override or extend those in the `defaults` dictionary.
+
+    Returns
+    -------
+    dict of {str: Any}
+        The merged dictionary, with values from the `overrides` dictionary taking precedence.
+
+    Examples
+    --------
+    >>> defaults = {
+    ...     "a": 1,
+    ...     "b": {
+    ...         "c": 3,
+    ...         "d": 4
+    ...     },
+    ...     "e": 5
+    ... }
+    >>> overrides = {
+    ...     "b": {
+    ...         "c": 30
+    ...     },
+    ...     "f": 6
+    ... }
+    >>> result = merge_dict(defaults, overrides)
+    >>> result
+    {'a': 1, 'b': {'c': 30, 'd': 4}, 'e': 5, 'f': 6}
+
+    Notes
+    -----
+    - The `merge_dict` function modifies the `defaults` dictionary in place. If you need to preserve the original
+      `defaults` dictionary, consider passing a copy instead.
+    - This function is particularly useful when combining configuration dictionaries where certain settings should
+      override defaults.
+    """
+    for key, value in overrides.items():
+        if isinstance(value, dict) and value:
+            defaults[key] = merge_dict(defaults.get(key, {}), value)
+        else:
+            defaults[key] = overrides[key]
+    return defaults

nv_ingest_api/util/converters/datetools.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+from datetime import datetime
+from datetime import timezone
+
+from dateutil.parser import parse
+
+from nv_ingest_api.util.exception_handlers.converters import datetools_exception_handler
+
+
+@datetools_exception_handler
+def datetimefrompdfmeta(pdf_formated_date: str, keep_tz: bool = False) -> str:
+    """
+    Convert PDF metadata formatted date string to a datetime object.
+
+    Parameters
+    ----------
+    pdf_formated_date : str
+        A date string in standard PDF metadata format.
+        Example: `str("D:20211222141131-07'00'")`
+    keep_tz : bool, optional
+        Keep or remove the timezone attribute of the parsed datetime object. If `False` (necessary for arrow format),
+        the timezone offset will be added to the datetime. Parsed datetimes will be in the same local time.
+
+    Returns
+    -------
+    str
+        A datetime object parsed from the input date string in ISO 8601 format.
+
+    """
+
+    try:
+        # standard pdf date format
+        pattern = "D:%Y%m%d%H%M%S%z"
+        # clean up date string
+        cleaned_date_string = pdf_formated_date[:-1].replace("'", ":")
+        parsed_dt_tz = datetime.strptime(cleaned_date_string, pattern)
+    except ValueError:
+        parsed_dt_tz = parse(pdf_formated_date, fuzzy=True)
+
+    if not keep_tz:
+        return remove_tz(parsed_dt_tz).isoformat()
+
+    return parsed_dt_tz.isoformat()
+
+
+def remove_tz(datetime_obj: datetime) -> datetime:
+    """
+    Remove timezone and add offset to a datetime object.
+
+    Parameters
+    ----------
+    datetime_obj : datetime.datetime
+        A datetime object with or without the timezone attribute set.
+
+    Returns
+    -------
+    datetime.datetime
+        A datetime object with the timezone offset added and the timezone attribute removed.
+
+    """
+
+    if datetime_obj.tzinfo is not None:  # If timezone info is present
+        # Convert to UTC
+        datetime_obj = datetime_obj.astimezone(timezone.utc)
+        # Remove timezone information
+        datetime_obj = datetime_obj.replace(tzinfo=None)
+
+    return datetime_obj
+
+
+def validate_iso8601(date_string: str) -> None:
+    """
+    Verify that the given date string is in ISO 8601 format.
+
+    Parameters
+    ----------
+    date_string : str
+        A date string in human-readable format, ideally ISO 8601.
+
+    Raises
+    ------
+    ValueError
+        If the date string is not in a valid ISO 8601 format.
+    """
+
+    assert datetime.fromisoformat(date_string)

nv_ingest_api/util/converters/dftools.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import io
+import json
+
+import fastparquet
+import pandas as pd
+
+import cudf
+
+
+class MemoryFiles:
+    def __init__(self):
+        self.output = {}
+
+    def open(self, fn, mode="rb"):
+        if mode != "wb":
+            try:
+                self.output[fn].seek(0)
+            except KeyError:
+                raise FileNotFoundError
+            return self.output[fn]
+
+        i = io.BytesIO()
+        self.output[fn] = i
+        self.output[fn].close = lambda: None
+
+        return i
+
+
+def pandas_to_cudf(
+    df: pd.DataFrame,
+    deserialize_cols: list = [],
+    default_cols: dict = {"document_type": str, "metadata": str},
+    default_type: type = str,
+) -> cudf.DataFrame:
+    """
+    Helper function to convert from pandas to cudf until https://github.com/apache/arrow/pull/40412 is resolved.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        A pandas dataframe.
+    Returns
+    -------
+    cudf.DataFrame
+        A cuDF dataframe.
+    """
+
+    if not df.empty:
+        files = MemoryFiles()
+        for col in deserialize_cols:
+            df[col] = df[col].apply(lambda x: json.loads(x))
+        df = pd.concat([df, df.iloc[0:1]], axis=0)
+
+        fastparquet.write("_", df, open_with=files.open, compression="UNCOMPRESSED", object_encoding="json")
+
+        with files.output["_"] as bytes_buf:
+            gdf = cudf.read_parquet(bytes_buf).iloc[:-1]
+            gdf.index.name = None
+
+            return gdf
+    else:
+        gdf = cudf.DataFrame({col: [] for col in default_cols})
+        for col in df.columns:
+            field_type = default_cols.get(col, default_type)
+            gdf[col] = gdf[col].astype(field_type)
+
+        return gdf
+
+
+def cudf_to_pandas(gdf: cudf.DataFrame, deserialize_cols: list = []) -> pd.DataFrame:
+    """
+    Helper function to convert from cudf to pandas until https://github.com/apache/arrow/pull/40412 is resolved.
+
+    Parameters
+    ----------
+    gdf : cudf.DataFrame
+        A cuDF dataframe.
+    nested_cols : list
+        A list of columns containing nested data.
+    Returns
+    -------
+    pd.DataFrame
+        A pandas dataframe.
+    """
+
+    with io.BytesIO() as bytes_buf:
+        gdf.to_parquet(bytes_buf)
+        df = pd.read_parquet(bytes_buf, engine="fastparquet", index=None)
+
+        for col in deserialize_cols:
+            if col in df.columns:
+                df[col] = df[col].apply(lambda x: json.loads(x))
+
+        return df
+
+
+def cudf_to_json(gdf: cudf.DataFrame, deserialize_cols: list = []) -> str:
+    """
+    Helper function to convert from cudf to json until https://github.com/apache/arrow/pull/40412 is resolved.
+
+    Parameters
+    ----------
+    gdf : cudf.DataFrame
+        A cuDF dataframe.
+    nested_cols : list
+        A list of columns containing nested data.
+    Returns
+    -------
+    str
+        A JSON formated string.
+    """
+
+    records = []
+    dict_vals = cudf_to_pandas(gdf).to_dict(orient="records")
+    for d in dict_vals:
+        temp = {}
+        for key, val in d.items():
+            if key in deserialize_cols:
+                val = json.loads(val)
+            temp[key] = val
+        records.append(temp)
+
+    return records

nv_ingest_api/util/converters/formats.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# pylint: skip-file
+
+import json
+
+
+def ingest_json_results_to_blob(result_content):
+    """
+    Parse a JSON string or BytesIO object, combine and sort entries, and create a blob string.
+
+    Returns:
+        str: The generated blob string.
+    """
+    try:
+        # Load the JSON data
+        data = json.loads(result_content) if isinstance(result_content, str) else json.loads(result_content)
+        data = data["data"]
+
+        # Smarter sorting: by page, then structured objects by x0, y0
+        def sorting_key(entry):
+            page = entry["metadata"]["content_metadata"]["page_number"]
+            if entry["document_type"] == "structured":
+                # Use table location's x0 and y0 as secondary keys
+                x0 = entry["metadata"]["table_metadata"]["table_location"][0]
+                y0 = entry["metadata"]["table_metadata"]["table_location"][1]
+            else:
+                # Non-structured objects are sorted after structured ones
+                x0 = float("inf")
+                y0 = float("inf")
+            return page, x0, y0
+
+        data.sort(key=sorting_key)
+
+        # Initialize the blob string
+        blob = []
+
+        for entry in data:
+            document_type = entry.get("document_type", "")
+
+            if document_type == "structured":
+                # Add table content to the blob
+                blob.append(entry["metadata"]["table_metadata"]["table_content"])
+                blob.append("\n")
+
+            elif document_type == "text":
+                # Add content to the blob
+                blob.append(entry["metadata"]["content"])
+                blob.append("\n")
+
+            elif document_type == "image":
+                # Add image caption to the blob
+                caption = entry["metadata"]["image_metadata"].get("caption", "")
+                blob.append(f"image_caption:[{caption}]")
+                blob.append("\n")
+
+        # Join all parts of the blob into a single string
+        return "".join(blob)
+
+    except Exception as e:
+        print(f"[ERROR] An error occurred while processing JSON content: {e}")
+        return ""

nv_ingest_api/util/converters/type_mappings.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import DocumentTypeEnum
+from nv_ingest_api.internal.enums.common import ContentTypeEnum
+
+DOC_TO_CONTENT_MAP = {
+    DocumentTypeEnum.BMP: ContentTypeEnum.IMAGE,
+    DocumentTypeEnum.DOCX: ContentTypeEnum.STRUCTURED,
+    DocumentTypeEnum.HTML: ContentTypeEnum.STRUCTURED,
+    DocumentTypeEnum.JPEG: ContentTypeEnum.IMAGE,
+    DocumentTypeEnum.MP3: ContentTypeEnum.AUDIO,
+    DocumentTypeEnum.PDF: ContentTypeEnum.STRUCTURED,
+    DocumentTypeEnum.PNG: ContentTypeEnum.IMAGE,
+    DocumentTypeEnum.PPTX: ContentTypeEnum.STRUCTURED,
+    DocumentTypeEnum.SVG: ContentTypeEnum.IMAGE,
+    DocumentTypeEnum.TIFF: ContentTypeEnum.IMAGE,
+    DocumentTypeEnum.TXT: ContentTypeEnum.TEXT,
+    DocumentTypeEnum.WAV: ContentTypeEnum.AUDIO,
+}
+
+
+def doc_type_to_content_type(doc_type: DocumentTypeEnum) -> ContentTypeEnum:
+    """
+    Convert DocumentTypeEnum to ContentTypeEnum
+    """
+    return DOC_TO_CONTENT_MAP[doc_type]

nv_ingest_api/util/detectors/__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.