arize 8.0.0a22__py3-none-any.whl → 8.0.0b0__py3-none-any.whl

arize/embeddings/tabular_generators.py

@@ -1,6 +1,7 @@
+"""Tabular data embedding generators for structured feature embeddings."""
+
 import logging
 from functools import partial
-from typing import Dict, List, Optional, Tuple, Union, cast
 
 import pandas as pd
 
@@ -10,7 +11,7 @@ from arize.embeddings.constants import (
     IMPORT_ERROR_MESSAGE,
 )
 from arize.embeddings.usecases import UseCases
-from arize.types import is_list_of
+from arize.utils.types import is_list_of
 
 try:
     from datasets import Dataset
@@ -27,7 +28,10 @@ TABULAR_PRETRAINED_MODELS = [
 
 
 class EmbeddingGeneratorForTabularFeatures(NLPEmbeddingGenerator):
+    """Embedding generator for tabular feature data using prompt-based LLM encoding."""
+
     def __repr__(self) -> str:
+        """Return a string representation of the tabular embedding generator."""
         return (
             f"{self.__class__.__name__}(\n"
             f"  use_case={self.use_case},\n"
@@ -41,8 +45,17 @@ class EmbeddingGeneratorForTabularFeatures(NLPEmbeddingGenerator):
     def __init__(
         self,
         model_name: str = DEFAULT_TABULAR_MODEL,
-        **kwargs,
-    ):
+        **kwargs: object,
+    ) -> None:
+        """Initialize the tabular features embedding generator.
+
+        Args:
+            model_name: Name of the pre-trained NLP model for tabular data.
+            **kwargs: Additional arguments for model initialization.
+
+        Raises:
+            ValueError: If model_name is not in supported models list.
+        """
         if model_name not in TABULAR_PRETRAINED_MODELS:
             raise ValueError(
                 "model_name not supported. Check supported models with "
@@ -57,25 +70,28 @@ class EmbeddingGeneratorForTabularFeatures(NLPEmbeddingGenerator):
     def generate_embeddings(
         self,
         df: pd.DataFrame,
-        selected_columns: List[str],
-        col_name_map: Optional[Dict[str, str]] = None,
+        selected_columns: list[str],
+        col_name_map: dict[str, str] | None = None,
         return_prompt_col: bool = False,
-    ) -> Union[pd.Series, Tuple[pd.Series, pd.Series]]:
-        """
-        Obtain embedding vectors from your tabular data. Prompts are generated from your
-        `selected_columns` and passed to a pre-trained large language model for embedding vector
-        computation.
-
-        :param df: pandas DataFrame containing the tabular data, not all columns will be
-        considered, see `selected_columns`.
-        :param selected_columns: columns to be considered to construct the prompt to be passed to
-        the LLM.
-        :param col_name_map: mapping between selected column names and a more verbose description of
-        the name. This helps the LLM understand the features better.
-        :param return_prompt_col: if set to True, an extra pandas Series will be returned
-        containing the constructed prompts. Defaults to False.
-        :return: a pandas Series containing the embedding vectors and, if `return_prompt_col` is
-        set to True, a pandas Seres containing the prompts created from tabular features.
+    ) -> pd.Series | tuple[pd.Series, pd.Series]:
+        """Obtain embedding vectors from your tabular data.
+
+        Prompts are generated from your `selected_columns` and passed to a pre-trained
+        large language model for embedding vector computation.
+
+        Args:
+            df: Pandas DataFrame containing the tabular data. Not all columns will be
+                considered, see `selected_columns`.
+            selected_columns: Columns to be considered to construct the prompt to be passed to
+                the LLM.
+            col_name_map: Mapping between selected column names and a more verbose description of
+                the name. This helps the LLM understand the features better.
+            return_prompt_col: If set to True, an extra pandas Series will be returned
+                containing the constructed prompts. Defaults to False.
+
+        Returns:
+            A pandas Series containing the embedding vectors and, if `return_prompt_col` is
+            set to True, a pandas Series containing the prompts created from tabular features.
         """
         if col_name_map is None:
             col_name_map = {}
@@ -99,7 +115,7 @@ class EmbeddingGeneratorForTabularFeatures(NLPEmbeddingGenerator):
             )
         for k, v in col_name_map.items():
             if not isinstance(k, str) or not isinstance(v, str):
-                raise ValueError(
+                raise TypeError(
                     "col_name_map dictionary keys and values should be strings"
                 )
         missing_cols = set(col_name_map.keys()).difference(df.columns)
@@ -109,7 +125,7 @@ class EmbeddingGeneratorForTabularFeatures(NLPEmbeddingGenerator):
                 f"The following columns are not found {missing_cols}"
             )
 
-        prompts = df.rename(columns=col_name_map).apply(
+        prompts: pd.Series = df.rename(columns=col_name_map).apply(
             partial(
                 self.__prompt_fn,
                 columns=[
@@ -129,22 +145,20 @@ class EmbeddingGeneratorForTabularFeatures(NLPEmbeddingGenerator):
             batch_size=self.batch_size,
         )
 
+        df: pd.DataFrame = ds.to_pandas()
         if return_prompt_col:
-            return (
-                cast(pd.DataFrame, ds.to_pandas())["embedding_vector"],
-                cast(pd.Series, prompts),
-            )
+            return df["embedding_vector"], prompts
 
-        return cast(pd.DataFrame, ds.to_pandas())["embedding_vector"]
+        return df["embedding_vector"]
 
     @staticmethod
-    def __prompt_fn(row: pd.DataFrame, columns: List[str]) -> str:
+    def __prompt_fn(row: pd.DataFrame, columns: list[str]) -> str:
         return " ".join(
             f"The {col.replace('_', ' ')} is {str(row[col]).strip()}."
             for col in columns
         )
 
-    def __get_method_for_embedding_calculation(self):
+    def __get_method_for_embedding_calculation(self) -> str:
         try:
             return {
                 "bert-base-uncased": "avg_token",
@@ -158,4 +172,5 @@ class EmbeddingGeneratorForTabularFeatures(NLPEmbeddingGenerator):
 
     @staticmethod
     def list_pretrained_models() -> pd.DataFrame:
+        """Return a DataFrame of available pretrained tabular models."""
         return pd.DataFrame({"Model Name": sorted(TABULAR_PRETRAINED_MODELS)})

arize/embeddings/usecases.py

@@ -1,26 +1,36 @@
+"""Use case definitions and enums for embedding generation."""
+
 from dataclasses import dataclass
 from enum import Enum, auto, unique
 
 
 @unique
 class NLPUseCases(Enum):
+    """Enum representing supported NLP use cases for embedding generation."""
+
     SEQUENCE_CLASSIFICATION = auto()
     SUMMARIZATION = auto()
 
 
 @unique
 class CVUseCases(Enum):
+    """Enum representing supported computer vision use cases for embedding generation."""
+
     IMAGE_CLASSIFICATION = auto()
     OBJECT_DETECTION = auto()
 
 
 @unique
-class TabularUsecases(Enum):
+class TabularUseCases(Enum):
+    """Enum representing supported tabular/structured data use cases for embedding generation."""
+
     TABULAR_EMBEDDINGS = auto()
 
 
 @dataclass
 class UseCases:
+    """Container grouping all use case enums for embedding generators."""
+
     NLP = NLPUseCases
     CV = CVUseCases
-    STRUCTURED = TabularUsecases
+    STRUCTURED = TabularUseCases

arize/exceptions/__init__.py

@@ -0,0 +1 @@
+"""Custom exceptions and error types for the Arize SDK."""

arize/exceptions/auth.py

@@ -1,11 +1,21 @@
+"""Authentication-related exception classes."""
+
 from arize.constants.config import ENV_API_KEY
 
 
 class MissingAPIKeyError(Exception):
-    def __init__(self, message: str = ""):
+    """Raised when API key is not provided via environment or configuration."""
+
+    def __init__(self, message: str = "") -> None:
+        """Initialize the exception with an optional custom message.
+
+        Args:
+            message: Custom error message, or empty string for default.
+        """
         self.message = message
 
     def __str__(self) -> str:
+        """Return the error message."""
         return self.message or self._default_message()
 
     @staticmethod

arize/exceptions/base.py

@@ -1,6 +1,7 @@
+"""Base exception classes and common error messages."""
+
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from typing import List
 
 INVALID_ARROW_CONVERSION_MSG = (
     "The dataframe needs to convert to pyarrow but has failed to do so. "
@@ -11,20 +12,30 @@ INVALID_ARROW_CONVERSION_MSG = (
 
 
 class ValidationError(Exception, ABC):
+    """Base exception for validation errors in data and schema validation."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     @abstractmethod
     def __repr__(self) -> str:
-        pass
+        """Return a string representation for debugging and logging."""
 
     @abstractmethod
     def error_message(self) -> str:
-        pass
+        """Return the error message for this exception."""
 
 
 class ValidationFailure(Exception):
-    def __init__(self, errors: List[ValidationError]) -> None:
+    """Raised when one or more validation errors occur during validation."""
+
+    def __init__(self, errors: list[ValidationError]) -> None:
+        """Initialize the exception with a list of validation errors.
+
+        Args:
+            errors: List of ValidationError instances that occurred.
+        """
         self.errors = errors
 
 
@@ -44,14 +55,24 @@ class ValidationFailure(Exception):
 
 
 class InvalidFieldTypeConversion(ValidationError):
+    """Raised when fields cannot be converted to required type."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Input_Type_Conversion"
 
     def __init__(self, fields: Iterable, type: str) -> None:
+        """Initialize the exception with type conversion context.
+
+        Args:
+            fields: Fields that failed type conversion.
+            type: Target type that fields should be convertible to.
+        """
         self.fields = fields
         self.type = type
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"The following fields must be convertible to {self.type}: "
             f"{', '.join(map(str, self.fields))}."
@@ -84,10 +105,14 @@ class InvalidFieldTypeConversion(ValidationError):
 
 
 class InvalidDataFrameIndex(ValidationError):
+    """Raised when DataFrame has an invalid index that needs to be reset."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Index"
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             "The index of the dataframe is invalid; "
             "reset the index by using df.reset_index(drop=True, inplace=True)"

arize/exceptions/models.py

@@ -1,8 +1,19 @@
+"""Model and project-related exception classes."""
+
+
 class MissingProjectNameError(Exception):
-    def __init__(self, message: str = ""):
+    """Raised when project name is required but not provided."""
+
+    def __init__(self, message: str = "") -> None:
+        """Initialize the exception with an optional custom message.
+
+        Args:
+            message: Custom error message, or empty string for default.
+        """
         self.message = message
 
     def __str__(self) -> str:
+        """Return the error message."""
         return self.message or self._default_message()
 
     @staticmethod
@@ -11,10 +22,18 @@ class MissingProjectNameError(Exception):
 
 
 class MissingModelNameError(Exception):
-    def __init__(self, message: str = ""):
+    """Raised when model name is required but not provided."""
+
+    def __init__(self, message: str = "") -> None:
+        """Initialize the exception with an optional custom message.
+
+        Args:
+            message: Custom error message, or empty string for default.
+        """
         self.message = message
 
     def __str__(self) -> str:
+        """Return the error message."""
         return self.message or self._default_message()
 
     @staticmethod

arize/exceptions/parameters.py

@@ -1,3 +1,5 @@
+"""Parameter validation exception classes."""
+
 from arize.constants.ml import MAX_NUMBER_OF_EMBEDDINGS
 from arize.exceptions.base import ValidationError
 
@@ -156,10 +158,14 @@ from arize.exceptions.base import ValidationError
 
 
 class InvalidModelVersion(ValidationError):
+    """Raised when model version is empty or invalid."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Model_Version"
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return "Model version must be a nonempty string."
 
 
@@ -172,10 +178,14 @@ class InvalidModelVersion(ValidationError):
 
 
 class InvalidProjectName(ValidationError):
+    """Raised when project name is empty or invalid."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Project_Name"
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             "Project Name must be a nonempty string. "
             "If Model ID was used instead of Project Name, "
@@ -352,13 +362,22 @@ class InvalidProjectName(ValidationError):
 
 
 class InvalidNumberOfEmbeddings(ValidationError):
+    """Raised when number of embedding features exceeds the maximum allowed."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Number_Of_Embeddings"
 
     def __init__(self, number_of_embeddings: int) -> None:
+        """Initialize the exception with embedding count context.
+
+        Args:
+            number_of_embeddings: The number of embeddings found in the schema.
+        """
         self.number_of_embeddings = number_of_embeddings
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"The schema contains {self.number_of_embeddings} different embeddings when a maximum of "
             f"{MAX_NUMBER_OF_EMBEDDINGS} is allowed."
@@ -366,23 +385,35 @@ class InvalidNumberOfEmbeddings(ValidationError):
 
 
 class InvalidValueType(Exception):
+    """Raised when a value has an invalid or unexpected type."""
+
     def __init__(
         self,
         value_name: str,
         value: bool | int | float | str,
         correct_type: str,
     ) -> None:
+        """Initialize the exception with value type validation context.
+
+        Args:
+            value_name: Name of the value with invalid type.
+            value: The actual value that has the wrong type.
+            correct_type: Description of the expected type.
+        """
         self.value_name = value_name
         self.value = value
         self.correct_type = correct_type
 
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Value_Type"
 
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"{self.value_name} with value {self.value} is of type {type(self.value).__name__}, "
             f"but expected {self.correct_type}"

arize/exceptions/spaces.py

@@ -1,8 +1,19 @@
+"""Space-related exception classes."""
+
+
 class MissingSpaceIDError(Exception):
-    def __init__(self, message: str = ""):
+    """Raised when space ID is required but not provided."""
+
+    def __init__(self, message: str = "") -> None:
+        """Initialize the exception with an optional custom message.
+
+        Args:
+            message: Custom error message, or empty string for default.
+        """
         self.message = message
 
     def __str__(self) -> str:
+        """Return the error message."""
         return self.message or self._default_message()
 
     @staticmethod

arize/exceptions/types.py

@@ -1,5 +1,6 @@
+"""Type validation exception classes."""
+
 from collections.abc import Iterable
-from typing import List
 
 from arize.constants.ml import (
     MAX_EMBEDDING_DIMENSIONALITY,
@@ -9,17 +10,28 @@ from arize.exceptions.base import ValidationError
 
 
 class InvalidType(ValidationError):
+    """Raised when a field has an invalid type compared to expected types."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Type"
 
     def __init__(
-        self, name: str, expected_types: List[str], found_data_type: str
+        self, name: str, expected_types: list[str], found_data_type: str
     ) -> None:
+        """Initialize the exception with type validation context.
+
+        Args:
+            name: Name of the field with invalid type.
+            expected_types: List of expected data types.
+            found_data_type: Actual data type found.
+        """
         self.name = name
         self.expected_types = expected_types
         self.found_data_type = found_data_type
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         type_list = (
             self.expected_types[0]
             if len(self.expected_types) == 1
@@ -33,16 +45,26 @@ class InvalidType(ValidationError):
 
 
 class InvalidTypeColumns(ValidationError):
+    """Raised when columns have invalid types compared to expected types."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Type_Columns"
 
     def __init__(
-        self, wrong_type_columns: List[str], expected_types: List[str]
+        self, wrong_type_columns: list[str], expected_types: list[str]
     ) -> None:
+        """Initialize the exception with column type validation context.
+
+        Args:
+            wrong_type_columns: Columns with incorrect data types.
+            expected_types: List of expected data types for the columns.
+        """
         self.wrong_type_columns = wrong_type_columns
         self.expected_types = expected_types
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         col_list = (
             self.wrong_type_columns[0]
             if len(self.wrong_type_columns) == 1
@@ -57,14 +79,24 @@ class InvalidTypeColumns(ValidationError):
 
 
 class InvalidTypeFeatures(ValidationError):
+    """Raised when feature columns have invalid types."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Type_Features"
 
-    def __init__(self, cols: Iterable, expected_types: List[str]) -> None:
+    def __init__(self, cols: Iterable, expected_types: list[str]) -> None:
+        """Initialize the exception with feature type validation context.
+
+        Args:
+            cols: Feature columns with unrecognized data types.
+            expected_types: List of expected data types for features.
+        """
         self.wrong_type_columns = cols
         self.expected_types = expected_types
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         type_list = (
             self.expected_types[0]
             if len(self.expected_types) == 1
@@ -78,30 +110,48 @@ class InvalidTypeFeatures(ValidationError):
 
 
 class InvalidFieldTypePromptTemplates(ValidationError):
+    """Raised when prompt template field has invalid type."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Input_Type_Prompt_Templates"
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return "prompt_template_column_names must be of type PromptTemplateColumnNames"
 
 
 class InvalidFieldTypeLlmConfig(ValidationError):
+    """Raised when LLM config field has invalid type."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Input_Type_LLM_Config"
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return "llm_config_column_names must be of type LLMConfigColumnNames"
 
 
 class InvalidTypeTags(ValidationError):
+    """Raised when tag columns have invalid types."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Type_Tags"
 
-    def __init__(self, cols: Iterable, expected_types: List[str]) -> None:
+    def __init__(self, cols: Iterable, expected_types: list[str]) -> None:
+        """Initialize the exception with tag type validation context.
+
+        Args:
+            cols: Tag columns with unrecognized data types.
+            expected_types: List of expected data types for tags.
+        """
         self.wrong_type_columns = cols
         self.expected_types = expected_types
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         type_list = (
             self.expected_types[0]
             if len(self.expected_types) == 1
@@ -115,14 +165,24 @@ class InvalidTypeTags(ValidationError):
 
 
 class InvalidValueEmbeddingVectorDimensionality(ValidationError):
+    """Raised when embedding vectors have invalid dimensionality."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Value_Embedding_Vector_Dimensionality"
 
-    def __init__(self, dim_1_cols: List[str], high_dim_cols: List[str]) -> None:
+    def __init__(self, dim_1_cols: list[str], high_dim_cols: list[str]) -> None:
+        """Initialize the exception with embedding dimensionality context.
+
+        Args:
+            dim_1_cols: Columns with dimensionality of 1.
+            high_dim_cols: Columns with dimensionality exceeding the maximum.
+        """
         self.dim_1_cols = dim_1_cols
         self.high_dim_cols = high_dim_cols
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         msg = (
             "Embedding vectors cannot have length (dimensionality) of 1 or higher "
             f"than {MAX_EMBEDDING_DIMENSIONALITY}. "
@@ -139,13 +199,22 @@ class InvalidValueEmbeddingVectorDimensionality(ValidationError):
 
 
 class InvalidValueEmbeddingRawDataTooLong(ValidationError):
+    """Raised when embedding raw data exceeds maximum character limit."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Value_Embedding_Raw_Data_Too_Long"
 
     def __init__(self, cols: Iterable) -> None:
+        """Initialize the exception with raw data length validation context.
+
+        Args:
+            cols: Columns with embedding raw data exceeding maximum characters.
+        """
         self.invalid_cols = cols
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"Embedding raw data cannot have more than {MAX_RAW_DATA_CHARACTERS} characters. "
             "The following columns do not satisfy this condition: "
@@ -154,14 +223,24 @@ class InvalidValueEmbeddingRawDataTooLong(ValidationError):
 
 
 class InvalidTypeShapValues(ValidationError):
+    """Raised when SHAP value columns have invalid types."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Type_SHAP_Values"
 
-    def __init__(self, cols: Iterable, expected_types: List[str]) -> None:
+    def __init__(self, cols: Iterable, expected_types: list[str]) -> None:
+        """Initialize the exception with SHAP value type validation context.
+
+        Args:
+            cols: SHAP value columns with unrecognized data types.
+            expected_types: List of expected data types for SHAP values.
+        """
         self.wrong_type_columns = cols
         self.expected_types = expected_types
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         type_list = (
             self.expected_types[0]
             if len(self.expected_types) == 1