nv-ingest-api 2025.4.17.dev20250417__py3-none-any.whl → 2025.4.19.dev20250419__py3-none-any.whl

nv_ingest_api/internal/transform/caption_image.py

@@ -1,205 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
-# All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import logging
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import pandas as pd
-from pydantic import BaseModel
-
-from nv_ingest_api.internal.primitives.nim.model_interface.vlm import VLMModelInterface
-from nv_ingest_api.internal.enums.common import ContentTypeEnum
-from nv_ingest_api.util.exception_handlers.decorators import unified_exception_handler
-from nv_ingest_api.util.image_processing import scale_image_to_encoding_size
-from nv_ingest_api.util.nim import create_inference_client
-
-logger = logging.getLogger(__name__)
-
-
-def _prepare_dataframes_mod(df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame, pd.Series]:
-    """
-    Prepares and returns three DataFrame-related objects from the input DataFrame.
-
-    The function performs the following:
-      1. Checks if the DataFrame is empty or if the "document_type" column is missing.
-         In such a case, returns the original DataFrame, an empty DataFrame, and an empty boolean Series.
-      2. Otherwise, it creates a boolean Series identifying rows where "document_type" equals IMAGE.
-      3. Extracts a DataFrame containing only those rows.
-
-    Parameters
-    ----------
-    df : pd.DataFrame
-        The input DataFrame that should contain a "document_type" column.
-
-    Returns
-    -------
-    Tuple[pd.DataFrame, pd.DataFrame, pd.Series]
-        A tuple containing:
-          - The original DataFrame.
-          - A DataFrame filtered to rows where "document_type" is IMAGE.
-          - A boolean Series indicating which rows in the original DataFrame are IMAGE rows.
-    """
-    try:
-        if df.empty or "document_type" not in df.columns:
-            return df, pd.DataFrame(), pd.Series(dtype=bool)
-
-        bool_index: pd.Series = df["document_type"] == ContentTypeEnum.IMAGE
-        df_matched: pd.DataFrame = df.loc[bool_index]
-
-        return df, df_matched, bool_index
-
-    except Exception as e:
-        err_msg = f"_prepare_dataframes_mod: Error preparing dataframes. Original error: {e}"
-        logger.error(err_msg, exc_info=True)
-        raise type(e)(err_msg) from e
-
-
-def _generate_captions(
-    base64_images: List[str], prompt: str, api_key: str, endpoint_url: str, model_name: str
-) -> List[str]:
-    """
-    Generates captions for a list of base64-encoded PNG images using the VLM model API.
-
-    This function performs the following steps:
-      1. Scales each image to meet encoding size requirements using `scale_image_to_encoding_size`.
-      2. Constructs the input payload containing the scaled images and the provided prompt.
-      3. Creates an inference client using the VLMModelInterface.
-      4. Calls the client's infer method to obtain a list of captions corresponding to the images.
-
-    Parameters
-    ----------
-    base64_images : List[str]
-        List of base64-encoded PNG image strings.
-    prompt : str
-        Text prompt to guide caption generation.
-    api_key : str
-        API key for authenticating with the VLM endpoint.
-    endpoint_url : str
-        URL of the VLM model HTTP endpoint.
-    model_name : str
-        The name of the model to use for inference.
-
-    Returns
-    -------
-    List[str]
-        A list of generated captions, each corresponding to an input image.
-
-    Raises
-    ------
-    Exception
-        Propagates any exception encountered during caption generation, with added context.
-    """
-    try:
-        # Scale each image to ensure it meets encoding size requirements.
-        scaled_images: List[str] = [scale_image_to_encoding_size(b64)[0] for b64 in base64_images]
-
-        # Build the input payload for the VLM model.
-        data: Dict[str, Any] = {
-            "base64_images": scaled_images,
-            "prompt": prompt,
-        }
-
-        # Create the inference client using the VLMModelInterface.
-        nim_client = create_inference_client(
-            model_interface=VLMModelInterface(),
-            endpoints=(None, endpoint_url),
-            auth_token=api_key,
-            infer_protocol="http",
-        )
-
-        logger.debug(f"Calling VLM endpoint: {endpoint_url} with model: {model_name}")
-        # Perform inference to generate captions.
-        captions: List[str] = nim_client.infer(data, model_name=model_name)
-        return captions
-
-    except Exception as e:
-        err_msg = f"_generate_captions: Error generating captions: {e}"
-        logger.error(err_msg, exc_info=True)
-        raise type(e)(err_msg) from e
-
-
-@unified_exception_handler
-def transform_image_create_vlm_caption_internal(
-    df_transform_ledger: pd.DataFrame,
-    task_config: Union[BaseModel, Dict[str, Any]],
-    transform_config: Any,
-    execution_trace_log: Optional[Dict[str, Any]] = None,
-) -> pd.DataFrame:
-    """
-    Extracts and adds captions for image content in a DataFrame using the VLM model API.
-
-    This function updates the 'metadata' column for rows where the content type is "image".
-    It uses configuration values from task_config (or falls back to transform_config defaults)
-    to determine the API key, prompt, endpoint URL, and model name for caption generation.
-    The generated captions are added under the 'image_metadata.caption' key in the metadata.
-
-    Parameters
-    ----------
-    df_transform_ledger : pd.DataFrame
-        The input DataFrame containing image data. Each row must have a 'metadata' column
-        with at least the 'content' and 'content_metadata' keys.
-    task_config : Union[BaseModel, Dict[str, Any]]
-        Configuration parameters for caption extraction. If provided as a Pydantic model,
-        it will be converted to a dictionary. Expected keys include "api_key", "prompt",
-        "endpoint_url", and "model_name".
-    transform_config : Any
-        A configuration object providing default values for caption extraction. It should have
-        attributes: api_key, prompt, endpoint_url, and model_name.
-    execution_trace_log : Optional[Dict[str, Any]], default=None
-        Optional trace information for debugging or logging purposes.
-
-    Returns
-    -------
-    pd.DataFrame
-        The updated DataFrame with generated captions added to the 'image_metadata.caption' field
-        within the 'metadata' column for each image row.
-
-    Raises
-    ------
-    Exception
-        Propagates any exception encountered during the caption extraction process, with added context.
-    """
-
-    _ = execution_trace_log  # Unused variable; placeholder to prevent linter warnings.
-
-    logger.debug("Attempting to caption image content")
-
-    # Convert task_config to dictionary if it is a Pydantic model.
-    if isinstance(task_config, BaseModel):
-        task_config = task_config.model_dump()
-
-    # Retrieve configuration values with fallback to transform_config defaults.
-    api_key: str = task_config.get("api_key") or transform_config.api_key
-    prompt: str = task_config.get("prompt") or transform_config.prompt
-    endpoint_url: str = task_config.get("endpoint_url") or transform_config.endpoint_url
-    model_name: str = task_config.get("model_name") or transform_config.model_name
-
-    # Create a mask for rows where the content type is "image".
-    df_mask: pd.Series = df_transform_ledger["metadata"].apply(
-        lambda meta: meta.get("content_metadata", {}).get("type") == "image"
-    )
-
-    # If no image rows exist, return the original DataFrame.
-    if not df_mask.any():
-        return df_transform_ledger
-
-    # Collect base64-encoded images from the rows where the content type is "image".
-    base64_images: List[str] = df_transform_ledger.loc[df_mask, "metadata"].apply(lambda meta: meta["content"]).tolist()
-
-    # Generate captions for the collected images.
-    captions: List[str] = _generate_captions(base64_images, prompt, api_key, endpoint_url, model_name)
-
-    # Update the DataFrame: assign each generated caption to the corresponding row.
-    for idx, caption in zip(df_transform_ledger.loc[df_mask].index, captions):
-        meta: Dict[str, Any] = df_transform_ledger.at[idx, "metadata"]
-        image_meta: Dict[str, Any] = meta.get("image_metadata", {})
-        image_meta["caption"] = caption
-        meta["image_metadata"] = image_meta
-        df_transform_ledger.at[idx, "metadata"] = meta
-
-    logger.debug("Image content captioning complete")
-    result, execution_trace_log = df_transform_ledger, {}
-    _ = execution_trace_log  # Unused variable; placeholder to prevent linter warnings.
-
-    return result

nv_ingest_api/internal/transform/embed_text.py

@@ -1,496 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
-# All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import logging
-from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, Tuple, Optional, Iterable, List
-
-import pandas as pd
-from openai import OpenAI
-
-from nv_ingest_api.internal.enums.common import ContentTypeEnum, StatusEnum, TaskTypeEnum
-from nv_ingest_api.internal.schemas.meta.metadata_schema import (
-    InfoMessageMetadataSchema,
-)
-from nv_ingest_api.internal.schemas.transform.transform_text_embedding_schema import TextEmbeddingSchema
-from nv_ingest_api.util.schema.schema_validator import validate_schema
-
-logger = logging.getLogger(__name__)
-
-
-# ------------------------------------------------------------------------------
-# Asynchronous Embedding Requests
-# ------------------------------------------------------------------------------
-
-
-def _make_async_request(
-    prompts: List[str],
-    api_key: str,
-    embedding_nim_endpoint: str,
-    embedding_model: str,
-    encoding_format: str,
-    input_type: str,
-    truncate: str,
-    filter_errors: bool,
-) -> list:
-    """
-    Interacts directly with the NIM embedding service to calculate embeddings for a batch of prompts.
-
-    Parameters
-    ----------
-    prompts : List[str]
-        A list of prompt strings for which embeddings are to be calculated.
-    api_key : str
-        API key for authentication with the embedding service.
-    embedding_nim_endpoint : str
-        Base URL for the NIM embedding service.
-    embedding_model : str
-        The model to use for generating embeddings.
-    encoding_format : str
-        The desired encoding format.
-    input_type : str
-        The type of input data.
-    truncate : str
-        Truncation setting for the input data.
-    filter_errors : bool
-        Flag indicating whether to filter errors in the response.
-
-    Returns
-    -------
-    list
-        A dictionary with keys "embedding" (the embedding results) and "info_msg" (any error info).
-
-    Raises
-    ------
-    RuntimeError
-        If an error occurs during the embedding request, with an info message attached.
-    """
-    response = {}
-
-    try:
-        client = OpenAI(
-            api_key=api_key,
-            base_url=embedding_nim_endpoint,
-        )
-
-        resp = client.embeddings.create(
-            input=prompts,
-            model=embedding_model,
-            encoding_format=encoding_format,
-            extra_body={"input_type": input_type, "truncate": truncate},
-        )
-
-        response["embedding"] = resp.data
-        response["info_msg"] = None
-
-    except Exception as err:
-        info_msg = {
-            "task": TaskTypeEnum.EMBED.value,
-            "status": StatusEnum.ERROR.value,
-            "message": f"Embedding error: {err}",
-            "filter": filter_errors,
-        }
-        validated_info_msg = validate_schema(info_msg, InfoMessageMetadataSchema).model_dump()
-
-        response["embedding"] = [None] * len(prompts)
-        response["info_msg"] = validated_info_msg
-
-        raise RuntimeError(f"Embedding error occurred. Info message: {validated_info_msg}") from err
-
-    return response
-
-
-def _async_request_handler(
-    prompts: List[str],
-    api_key: str,
-    embedding_nim_endpoint: str,
-    embedding_model: str,
-    encoding_format: str,
-    input_type: str,
-    truncate: str,
-    filter_errors: bool,
-) -> List[dict]:
-    """
-    Gathers calculated embedding results from the NIM embedding service concurrently.
-
-    Parameters
-    ----------
-    prompts : List[str]
-        A list of prompt batches.
-    api_key : str
-        API key for authentication.
-    embedding_nim_endpoint : str
-        Base URL for the NIM embedding service.
-    embedding_model : str
-        The model to use for generating embeddings.
-    encoding_format : str
-        The desired encoding format.
-    input_type : str
-        The type of input data.
-    truncate : str
-        Truncation setting for the input data.
-    filter_errors : bool
-        Flag indicating whether to filter errors in the response.
-
-    Returns
-    -------
-    List[dict]
-        A list of response dictionaries from the embedding service.
-    """
-    with ThreadPoolExecutor() as executor:
-        futures = [
-            executor.submit(
-                _make_async_request,
-                prompts=prompt_batch,
-                api_key=api_key,
-                embedding_nim_endpoint=embedding_nim_endpoint,
-                embedding_model=embedding_model,
-                encoding_format=encoding_format,
-                input_type=input_type,
-                truncate=truncate,
-                filter_errors=filter_errors,
-            )
-            for prompt_batch in prompts
-        ]
-        results = [future.result() for future in futures]
-
-    return results
-
-
-def _async_runner(
-    prompts: List[str],
-    api_key: str,
-    embedding_nim_endpoint: str,
-    embedding_model: str,
-    encoding_format: str,
-    input_type: str,
-    truncate: str,
-    filter_errors: bool,
-) -> dict:
-    """
-    Concurrently launches all NIM embedding requests and flattens the results.
-
-    Parameters
-    ----------
-    prompts : List[str]
-        A list of prompt batches.
-    api_key : str
-        API key for authentication.
-    embedding_nim_endpoint : str
-        Base URL for the NIM embedding service.
-    embedding_model : str
-        The model to use for generating embeddings.
-    encoding_format : str
-        The desired encoding format.
-    input_type : str
-        The type of input data.
-    truncate : str
-        Truncation setting for the input data.
-    filter_errors : bool
-        Flag indicating whether to filter errors in the response.
-
-    Returns
-    -------
-    dict
-        A dictionary with keys "embeddings" (flattened embedding results) and "info_msgs" (error messages).
-    """
-    results = _async_request_handler(
-        prompts,
-        api_key,
-        embedding_nim_endpoint,
-        embedding_model,
-        encoding_format,
-        input_type,
-        truncate,
-        filter_errors,
-    )
-
-    flat_results = {"embeddings": [], "info_msgs": []}
-    for batch_dict in results:
-        info_msg = batch_dict["info_msg"]
-        for embedding in batch_dict["embedding"]:
-            if not isinstance(embedding, list):
-                if embedding is not None:
-                    flat_results["embeddings"].append(embedding.embedding)
-                else:
-                    flat_results["embeddings"].append(embedding)
-            else:
-                flat_results["embeddings"].append(embedding)
-            flat_results["info_msgs"].append(info_msg)
-
-    return flat_results
-
-
-# ------------------------------------------------------------------------------
-# Pandas UDFs for Content Extraction
-# ------------------------------------------------------------------------------
-
-
-def _add_embeddings(row, embeddings, info_msgs):
-    """
-    Updates a DataFrame row with embedding data and associated error info.
-
-    Parameters
-    ----------
-    row : pandas.Series
-        A row of the DataFrame.
-    embeddings : list
-        List of embeddings corresponding to DataFrame rows.
-    info_msgs : list
-        List of info message dictionaries corresponding to DataFrame rows.
-
-    Returns
-    -------
-    pandas.Series
-        The updated row with embedding and info message metadata added.
-    """
-    row["metadata"]["embedding"] = embeddings[row.name]
-    if info_msgs[row.name] is not None:
-        row["metadata"]["info_message_metadata"] = info_msgs[row.name]
-        row["document_type"] = ContentTypeEnum.INFO_MSG
-        row["_contains_embeddings"] = False
-    else:
-        row["_contains_embeddings"] = True
-
-    return row
-
-
-def _get_pandas_text_content(row):
-    """
-    Extracts text content from a DataFrame row.
-
-    Parameters
-    ----------
-    row : pandas.Series
-        A row containing the 'content' key.
-
-    Returns
-    -------
-    str
-        The text content from the row.
-    """
-    return row["content"]
-
-
-def _get_pandas_table_content(row):
-    """
-    Extracts table/chart content from a DataFrame row.
-
-    Parameters
-    ----------
-    row : pandas.Series
-        A row containing 'table_metadata' with 'table_content'.
-
-    Returns
-    -------
-    str
-        The table/chart content from the row.
-    """
-    return row["table_metadata"]["table_content"]
-
-
-def _get_pandas_image_content(row):
-    """
-    Extracts image caption content from a DataFrame row.
-
-    Parameters
-    ----------
-    row : pandas.Series
-        A row containing 'image_metadata' with 'caption'.
-
-    Returns
-    -------
-    str
-        The image caption from the row.
-    """
-    return row["image_metadata"]["caption"]
-
-
-# ------------------------------------------------------------------------------
-# Batch Processing Utilities
-# ------------------------------------------------------------------------------
-
-
-def _batch_generator(iterable: Iterable, batch_size: int = 10):
-    """
-    Yields batches of a specified size from an iterable.
-
-    Parameters
-    ----------
-    iterable : Iterable
-        The iterable to batch.
-    batch_size : int, optional
-        The size of each batch (default is 10).
-
-    Yields
-    ------
-    list
-        A batch of items from the iterable.
-    """
-    iter_len = len(iterable)
-    for idx in range(0, iter_len, batch_size):
-        yield iterable[idx : min(idx + batch_size, iter_len)]
-
-
-def _generate_batches(prompts: List[str], batch_size: int = 100) -> List[str]:
-    """
-    Splits a list of prompts into batches.
-
-    Parameters
-    ----------
-    prompts : List[str]
-        The list of prompt strings.
-    batch_size : int, optional
-        The desired batch size (default is 100).
-
-    Returns
-    -------
-    List[List[str]]
-        A list of batches, each containing a subset of the prompts.
-    """
-    return [batch for batch in _batch_generator(prompts, batch_size)]
-
-
-def _get_pandas_audio_content(row):
-    """
-    A pandas UDF used to select extracted audio transcription to be used to create embeddings.
-    """
-    return row["audio_metadata"]["audio_transcript"]
-
-
-# ------------------------------------------------------------------------------
-# DataFrame Concatenation Utility
-# ------------------------------------------------------------------------------
-
-
-def _concatenate_extractions_pandas(
-    base_df: pd.DataFrame, dataframes: List[pd.DataFrame], masks: List[pd.Series]
-) -> pd.DataFrame:
-    """
-    Concatenates processed DataFrame rows (with embeddings) with unprocessed rows from the base DataFrame.
-
-    Parameters
-    ----------
-    base_df : pd.DataFrame
-        The original DataFrame.
-    dataframes : List[pd.DataFrame]
-        List of DataFrames that have been enriched with embeddings.
-    masks : List[pd.Series]
-        List of boolean masks indicating the rows that were processed.
-
-    Returns
-    -------
-    pd.DataFrame
-        The concatenated DataFrame with embeddings applied where available.
-    """
-    unified_mask = pd.Series(False, index=base_df.index)
-    for mask in masks:
-        unified_mask = unified_mask | mask
-
-    df_no_text = base_df.loc[~unified_mask].copy()
-    df_no_text["_contains_embeddings"] = False
-
-    dataframes.append(df_no_text)
-    combined_df = pd.concat(dataframes, axis=0, ignore_index=True).reset_index(drop=True)
-    return combined_df
-
-
-# ------------------------------------------------------------------------------
-# Embedding Extraction Pipeline
-# ------------------------------------------------------------------------------
-
-
-def transform_create_text_embeddings_internal(
-    df_transform_ledger: pd.DataFrame,
-    task_config: Dict[str, Any],
-    transform_config: TextEmbeddingSchema = TextEmbeddingSchema(),
-    execution_trace_log: Optional[Dict] = None,
-) -> Tuple[pd.DataFrame, Dict]:
-    """
-    Generates text embeddings for supported content types (TEXT, STRUCTURED, IMAGE)
-    from a pandas DataFrame using asynchronous requests.
-
-    Parameters
-    ----------
-    df_transform_ledger : pd.DataFrame
-        The DataFrame containing content for embedding extraction.
-    task_config : Dict[str, Any]
-        Dictionary containing task properties (e.g., filter error flag).
-    transform_config : Any
-        Validated configuration for text embedding extraction (EmbedExtractionsSchema).
-    execution_trace_log : Optional[Dict], optional
-        Optional trace information for debugging or logging (default is None).
-
-    Returns
-    -------
-    Tuple[pd.DataFrame, Dict]
-        A tuple containing:
-            - The updated DataFrame with embeddings applied.
-            - A dictionary with trace information.
-    """
-    _ = task_config  # Currently unused.
-
-    if execution_trace_log is None:
-        execution_trace_log = {}
-        logger.debug("No trace_info provided. Initialized empty trace_info dictionary.")
-
-    # TODO(Devin)
-    if df_transform_ledger.empty:
-        return df_transform_ledger, {"trace_info": execution_trace_log}
-
-    embedding_dataframes = []
-    content_masks = []  # List of pandas boolean Series
-
-    # Define pandas content extractors for supported content types.
-    pandas_content_extractor = {
-        ContentTypeEnum.TEXT: _get_pandas_text_content,
-        ContentTypeEnum.STRUCTURED: _get_pandas_table_content,
-        ContentTypeEnum.IMAGE: _get_pandas_image_content,
-        ContentTypeEnum.AUDIO: _get_pandas_audio_content,
-        ContentTypeEnum.VIDEO: lambda x: None,  # Not supported yet.
-    }
-
-    logger.debug("Generating text embeddings for supported content types: TEXT, STRUCTURED, IMAGE.")
-
-    # Process each supported content type.
-    for content_type, content_getter in pandas_content_extractor.items():
-        if not content_getter:
-            logger.debug(f"Skipping unsupported content type: {content_type}")
-            continue
-
-        content_mask = df_transform_ledger["document_type"] == content_type.value
-        if not content_mask.any():
-            continue
-
-        # Extract content from metadata and filter out rows with empty content.
-        extracted_content = df_transform_ledger.loc[content_mask, "metadata"].apply(content_getter)
-        non_empty_mask = extracted_content.notna() & (extracted_content.str.strip() != "")
-        final_mask = content_mask & non_empty_mask
-        if not final_mask.any():
-            continue
-
-        df_content = df_transform_ledger.loc[final_mask].copy().reset_index(drop=True)
-        filtered_content = df_content["metadata"].apply(content_getter)
-        filtered_content_batches = _generate_batches(filtered_content.tolist(), batch_size=transform_config.batch_size)
-        content_embeddings = _async_runner(
-            filtered_content_batches,
-            transform_config.api_key,
-            transform_config.embedding_nim_endpoint,
-            transform_config.embedding_model,
-            transform_config.encoding_format,
-            transform_config.input_type,
-            transform_config.truncate,
-            False,
-        )
-        # Apply the embeddings (and any error info) to each row.
-        df_content[["metadata", "document_type", "_contains_embeddings"]] = df_content.apply(
-            _add_embeddings, **content_embeddings, axis=1
-        )[["metadata", "document_type", "_contains_embeddings"]]
-        df_content["_content"] = filtered_content
-
-        embedding_dataframes.append(df_content)
-        content_masks.append(final_mask)
-
-    combined_df = _concatenate_extractions_pandas(df_transform_ledger, embedding_dataframes, content_masks)
-    return combined_df, {"trace_info": execution_trace_log}