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

arize/embeddings/auto_generator.py

@@ -1,4 +1,4 @@
-from typing import Any
+"""Automatic embedding generation factory for various ML use cases."""
 
 import pandas as pd
 
@@ -30,7 +30,14 @@ UseCaseLike = str | UseCases.NLP | UseCases.CV | UseCases.STRUCTURED
 
 
 class EmbeddingGenerator:
-    def __init__(self, **kwargs: str):
+    """Factory class for creating embedding generators based on use case."""
+
+    def __init__(self, **kwargs: str) -> None:
+        """Raise error directing users to use from_use_case factory method.
+
+        Raises:
+            OSError: Always raised to prevent direct instantiation.
+        """
         raise OSError(
             f"{self.__class__.__name__} is designed to be instantiated using the "
             f"`{self.__class__.__name__}.from_use_case(use_case, **kwargs)` method."
@@ -38,23 +45,24 @@ class EmbeddingGenerator:
 
     @staticmethod
     def from_use_case(
-        use_case: UseCaseLike, **kwargs: Any
+        use_case: UseCaseLike, **kwargs: object
     ) -> BaseEmbeddingGenerator:
+        """Create an embedding generator for the specified use case."""
         if use_case == UseCases.NLP.SEQUENCE_CLASSIFICATION:
             return EmbeddingGeneratorForNLPSequenceClassification(**kwargs)
-        elif use_case == UseCases.NLP.SUMMARIZATION:
+        if use_case == UseCases.NLP.SUMMARIZATION:
             return EmbeddingGeneratorForNLPSummarization(**kwargs)
-        elif use_case == UseCases.CV.IMAGE_CLASSIFICATION:
+        if use_case == UseCases.CV.IMAGE_CLASSIFICATION:
             return EmbeddingGeneratorForCVImageClassification(**kwargs)
-        elif use_case == UseCases.CV.OBJECT_DETECTION:
+        if use_case == UseCases.CV.OBJECT_DETECTION:
             return EmbeddingGeneratorForCVObjectDetection(**kwargs)
-        elif use_case == UseCases.STRUCTURED.TABULAR_EMBEDDINGS:
+        if use_case == UseCases.STRUCTURED.TABULAR_EMBEDDINGS:
             return EmbeddingGeneratorForTabularFeatures(**kwargs)
-        else:
-            raise ValueError(f"Invalid use case {use_case}")
+        raise ValueError(f"Invalid use case {use_case}")
 
     @classmethod
     def list_default_models(cls) -> pd.DataFrame:
+        """Return a DataFrame of default models for each use case."""
         df = pd.DataFrame(
             {
                 "Area": ["NLP", "NLP", "CV", "CV", "STRUCTURED"],
@@ -74,13 +82,12 @@ class EmbeddingGenerator:
                 ],
             }
         )
-        df.sort_values(
-            by=[col for col in df.columns], ascending=True, inplace=True
-        )
+        df.sort_values(by=list(df.columns), ascending=True, inplace=True)
         return df.reset_index(drop=True)
 
     @classmethod
     def list_pretrained_models(cls) -> pd.DataFrame:
+        """Return a DataFrame of all available pretrained models."""
         data = {
             "Task": ["NLP" for _ in NLP_PRETRAINED_MODELS]
             + ["CV" for _ in CV_PRETRAINED_MODELS],
@@ -91,18 +98,15 @@ class EmbeddingGenerator:
             "Model Name": NLP_PRETRAINED_MODELS + CV_PRETRAINED_MODELS,
         }
         df = pd.DataFrame(data)
-        df.sort_values(
-            by=[col for col in df.columns], ascending=True, inplace=True
-        )
+        df.sort_values(by=list(df.columns), ascending=True, inplace=True)
         return df.reset_index(drop=True)
 
     @staticmethod
     def __parse_model_arch(model_name: str) -> str:
         if constants.GPT.lower() in model_name.lower():
             return constants.GPT
-        elif constants.BERT.lower() in model_name.lower():
+        if constants.BERT.lower() in model_name.lower():
             return constants.BERT
-        elif constants.VIT.lower() in model_name.lower():
+        if constants.VIT.lower() in model_name.lower():
             return constants.VIT
-        else:
-            raise ValueError("Invalid model_name, unknown architecture.")
+        raise ValueError("Invalid model_name, unknown architecture.")

arize/embeddings/base_generators.py

@@ -1,8 +1,9 @@
+"""Base embedding generator classes for NLP, CV, and tabular data."""
+
 import os
 from abc import ABC, abstractmethod
 from enum import Enum
 from functools import partial
-from typing import Dict, List, Union, cast
 
 import pandas as pd
 
@@ -31,9 +32,26 @@ transformer_logging.enable_progress_bar()
 
 
 class BaseEmbeddingGenerator(ABC):
+    """Abstract base class for all embedding generators."""
+
     def __init__(
-        self, use_case: Enum, model_name: str, batch_size: int = 100, **kwargs
-    ):
+        self,
+        use_case: Enum,
+        model_name: str,
+        batch_size: int = 100,
+        **kwargs: object,
+    ) -> None:
+        """Initialize the embedding generator with model and configuration.
+
+        Args:
+            use_case: Enum specifying the use case for embedding generation.
+            model_name: Name of the pre-trained model to use.
+            batch_size: Number of samples to process per batch.
+            **kwargs: Additional arguments for model initialization.
+
+        Raises:
+            HuggingFaceRepositoryNotFound: If the model name is not found on HuggingFace.
+        """
         self.__use_case = self._parse_use_case(use_case=use_case)
         self.__model_name = model_name
         self.__device = self.select_device()
@@ -45,43 +63,50 @@ class BaseEmbeddingGenerator(ABC):
             ).to(self.device)
         except OSError as e:
             raise err.HuggingFaceRepositoryNotFound(model_name) from e
-        except Exception as e:
-            raise e
+        except Exception:
+            raise
 
     @abstractmethod
-    def generate_embeddings(self, **kwargs) -> pd.Series: ...
+    def generate_embeddings(self, **kwargs: object) -> pd.Series:
+        """Generate embeddings for the input data."""
+        ...
 
     def select_device(self) -> torch.device:
+        """Select the best available device (CUDA, MPS, or CPU) for model execution."""
         if torch.cuda.is_available():
             return torch.device("cuda")
-        elif torch.backends.mps.is_available():
+        if torch.backends.mps.is_available():
             return torch.device("mps")
-        else:
-            logger.warning(
-                "No available GPU has been detected. The use of GPU acceleration is "
-                "strongly recommended. You can check for GPU availability by running "
-                "`torch.cuda.is_available()` or `torch.backends.mps.is_available()`."
-            )
-            return torch.device("cpu")
+        logger.warning(
+            "No available GPU has been detected. The use of GPU acceleration is "
+            "strongly recommended. You can check for GPU availability by running "
+            "`torch.cuda.is_available()` or `torch.backends.mps.is_available()`."
+        )
+        return torch.device("cpu")
 
     @property
     def use_case(self) -> str:
+        """Return the use case for this embedding generator."""
         return self.__use_case
 
     @property
     def model_name(self) -> str:
+        """Return the name of the model being used."""
         return self.__model_name
 
     @property
-    def model(self):
+    def model(self) -> object:
+        """Return the underlying model instance."""
         return self.__model
 
     @property
     def device(self) -> torch.device:
+        """Return the device (CPU/GPU) being used for computation."""
         return self.__device
 
     @property
     def batch_size(self) -> int:
+        """Return the batch size for processing."""
         return self.__batch_size
 
     @batch_size.setter
@@ -89,11 +114,10 @@ class BaseEmbeddingGenerator(ABC):
         err_message = "New batch size should be an integer greater than 0."
         if not isinstance(new_batch_size, int):
             raise TypeError(err_message)
-        elif new_batch_size <= 0:
+        if new_batch_size <= 0:
             raise ValueError(err_message)
-        else:
-            self.__batch_size = new_batch_size
-            logger.info(f"Batch size has been set to {new_batch_size}.")
+        self.__batch_size = new_batch_size
+        logger.info(f"Batch size has been set to {new_batch_size}.")
 
     @staticmethod
     def _parse_use_case(use_case: Enum) -> str:
@@ -102,8 +126,8 @@ class BaseEmbeddingGenerator(ABC):
         return f"{uc_area}.{uc_task}"
 
     def _get_embedding_vector(
-        self, batch: Dict[str, torch.Tensor], method
-    ) -> Dict[str, torch.Tensor]:
+        self, batch: dict[str, torch.Tensor], method: str
+    ) -> dict[str, torch.Tensor]:
         with torch.no_grad():
             outputs = self.model(**batch)
         # (batch_size, seq_length/or/num_tokens, hidden_size)
@@ -116,20 +140,23 @@ class BaseEmbeddingGenerator(ABC):
         return {"embedding_vector": embeddings.cpu().numpy().astype(float)}
 
     @staticmethod
-    def check_invalid_index(field: Union[pd.Series, pd.DataFrame]) -> None:
+    def check_invalid_index(field: pd.Series | pd.DataFrame) -> None:
+        """Check if the field has a valid index and raise error if invalid."""
         if (field.index != field.reset_index(drop=True).index).any():
             if isinstance(field, pd.DataFrame):
                 raise err.InvalidIndexError("DataFrame")
-            else:
-                raise err.InvalidIndexError(str(field.name))
+            raise err.InvalidIndexError(str(field.name))
 
     @abstractmethod
     def __repr__(self) -> str:
-        pass
+        """Return a string representation of the embedding generator."""
 
 
 class NLPEmbeddingGenerator(BaseEmbeddingGenerator):
+    """Base class for NLP embedding generators with text tokenization support."""
+
     def __repr__(self) -> str:
+        """Return a string representation of the NLP embedding generator."""
         return (
             f"{self.__class__.__name__}(\n"
             f"  use_case={self.use_case},\n"
@@ -146,8 +173,16 @@ class NLPEmbeddingGenerator(BaseEmbeddingGenerator):
         use_case: Enum,
         model_name: str,
         tokenizer_max_length: int = 512,
-        **kwargs,
-    ):
+        **kwargs: object,
+    ) -> None:
+        """Initialize the NLP embedding generator with tokenizer configuration.
+
+        Args:
+            use_case: Enum specifying the NLP use case.
+            model_name: Name of the pre-trained NLP model.
+            tokenizer_max_length: Maximum sequence length for the tokenizer.
+            **kwargs: Additional arguments for model initialization.
+        """
         super().__init__(use_case=use_case, model_name=model_name, **kwargs)
         self.__tokenizer_max_length = tokenizer_max_length
         # We don't check for the tokenizer's existence since it is coupled with the corresponding model
@@ -158,16 +193,19 @@ class NLPEmbeddingGenerator(BaseEmbeddingGenerator):
         )
 
     @property
-    def tokenizer(self):
+    def tokenizer(self) -> object:
+        """Return the tokenizer instance for text processing."""
         return self.__tokenizer
 
     @property
     def tokenizer_max_length(self) -> int:
+        """Return the maximum sequence length for the tokenizer."""
         return self.__tokenizer_max_length
 
     def tokenize(
-        self, batch: Dict[str, List[str]], text_feat_name: str
+        self, batch: dict[str, list[str]], text_feat_name: str
     ) -> BatchEncoding:
+        """Tokenize a batch of text inputs."""
         return self.tokenizer(
             batch[text_feat_name],
             padding=True,
@@ -178,7 +216,10 @@ class NLPEmbeddingGenerator(BaseEmbeddingGenerator):
 
 
 class CVEmbeddingGenerator(BaseEmbeddingGenerator):
+    """Base class for computer vision embedding generators with image preprocessing support."""
+
     def __repr__(self) -> str:
+        """Return a string representation of the computer vision embedding generator."""
         return (
             f"{self.__class__.__name__}(\n"
             f"  use_case={self.use_case},\n"
@@ -189,7 +230,16 @@ class CVEmbeddingGenerator(BaseEmbeddingGenerator):
             f")"
         )
 
-    def __init__(self, use_case: Enum, model_name: str, **kwargs):
+    def __init__(
+        self, use_case: Enum, model_name: str, **kwargs: object
+    ) -> None:
+        """Initialize the computer vision embedding generator with image processor.
+
+        Args:
+            use_case: Enum specifying the computer vision use case.
+            model_name: Name of the pre-trained vision model.
+            **kwargs: Additional arguments for model initialization.
+        """
         super().__init__(use_case=use_case, model_name=model_name, **kwargs)
         logger.info("Downloading image processor")
         # We don't check for the image processor's existence since it is coupled with the corresponding model
@@ -199,18 +249,21 @@ class CVEmbeddingGenerator(BaseEmbeddingGenerator):
         )
 
     @property
-    def image_processor(self):
+    def image_processor(self) -> object:
+        """Return the image processor instance for image preprocessing."""
         return self.__image_processor
 
     @staticmethod
     def open_image(image_path: str) -> Image.Image:
+        """Open and convert an image to RGB format."""
         if not os.path.exists(image_path):
             raise ValueError(f"Cannot find image {image_path}")
         return Image.open(image_path).convert("RGB")
 
     def preprocess_image(
-        self, batch: Dict[str, List[str]], local_image_feat_name: str
-    ):
+        self, batch: dict[str, list[str]], local_image_feat_name: str
+    ) -> object:
+        """Preprocess a batch of images for model input."""
         return self.image_processor(
             [
                 self.open_image(image_path)
@@ -220,8 +273,7 @@ class CVEmbeddingGenerator(BaseEmbeddingGenerator):
         ).to(self.device)
 
     def generate_embeddings(self, local_image_path_col: pd.Series) -> pd.Series:
-        """
-        Obtain embedding vectors from your image data using pre-trained image models.
+        """Obtain embedding vectors from your image data using pre-trained image models.
 
         :param local_image_path_col: a pandas Series containing the local path to the images to
         be used to generate the embedding vectors.
@@ -252,4 +304,5 @@ class CVEmbeddingGenerator(BaseEmbeddingGenerator):
             batched=True,
             batch_size=self.batch_size,
         )
-        return cast(pd.DataFrame, ds.to_pandas())["embedding_vector"]
+        df: pd.DataFrame = ds.to_pandas()
+        return df["embedding_vector"]

arize/embeddings/constants.py

@@ -1,3 +1,5 @@
+"""Embedding generation constants and pre-trained model definitions."""
+
 DEFAULT_NLP_SEQUENCE_CLASSIFICATION_MODEL = "distilbert-base-uncased"
 DEFAULT_NLP_SUMMARIZATION_MODEL = "distilbert-base-uncased"
 DEFAULT_TABULAR_MODEL = "distilbert-base-uncased"

arize/embeddings/cv_generators.py

@@ -1,3 +1,5 @@
+"""Computer vision embedding generators for image classification and object detection."""
+
 from arize.embeddings.base_generators import CVEmbeddingGenerator
 from arize.embeddings.constants import (
     DEFAULT_CV_IMAGE_CLASSIFICATION_MODEL,
@@ -7,9 +9,19 @@ from arize.embeddings.usecases import UseCases
 
 
 class EmbeddingGeneratorForCVImageClassification(CVEmbeddingGenerator):
+    """Embedding generator for computer vision image classification tasks."""
+
     def __init__(
-        self, model_name: str = DEFAULT_CV_IMAGE_CLASSIFICATION_MODEL, **kwargs
-    ):
+        self,
+        model_name: str = DEFAULT_CV_IMAGE_CLASSIFICATION_MODEL,
+        **kwargs: object,
+    ) -> None:
+        """Initialize the image classification embedding generator.
+
+        Args:
+            model_name: Name of the pre-trained vision model.
+            **kwargs: Additional arguments for model initialization.
+        """
         super().__init__(
             use_case=UseCases.CV.IMAGE_CLASSIFICATION,
             model_name=model_name,
@@ -18,9 +30,19 @@ class EmbeddingGeneratorForCVImageClassification(CVEmbeddingGenerator):
 
 
 class EmbeddingGeneratorForCVObjectDetection(CVEmbeddingGenerator):
+    """Embedding generator for computer vision object detection tasks."""
+
     def __init__(
-        self, model_name: str = DEFAULT_CV_OBJECT_DETECTION_MODEL, **kwargs
-    ):
+        self,
+        model_name: str = DEFAULT_CV_OBJECT_DETECTION_MODEL,
+        **kwargs: object,
+    ) -> None:
+        """Initialize the object detection embedding generator.
+
+        Args:
+            model_name: Name of the pre-trained vision model.
+            **kwargs: Additional arguments for model initialization.
+        """
         super().__init__(
             use_case=UseCases.CV.OBJECT_DETECTION,
             model_name=model_name,

arize/embeddings/errors.py

@@ -1,37 +1,59 @@
+"""Embedding generation exception classes."""
+
+
 class InvalidIndexError(Exception):
+    """Raised when DataFrame or Series has an invalid index."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "Invalid_Index_Error"
 
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(self, field_name: str) -> None:
+        """Initialize the exception with field name context.
+
+        Args:
+            field_name: Name of the DataFrame or Series field with invalid index.
+        """
         self.field_name = field_name
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         if self.field_name == "DataFrame":
             return (
                 f"The index of the {self.field_name} is invalid; "
                 f"reset the index by using df.reset_index(drop=True, inplace=True)"
             )
-        else:
-            return (
-                f"The index of the Series given by the column '{self.field_name}' is invalid; "
-                f"reset the index by using df.reset_index(drop=True, inplace=True)"
-            )
+        return (
+            f"The index of the Series given by the column '{self.field_name}' is invalid; "
+            f"reset the index by using df.reset_index(drop=True, inplace=True)"
+        )
 
 
 class HuggingFaceRepositoryNotFound(Exception):
+    """Raised when HuggingFace model repository is not found."""
+
     def __repr__(self) -> str:
+        """Return a string representation for debugging and logging."""
         return "HuggingFace_Repository_Not_Found_Error"
 
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(self, model_name: str) -> None:
+        """Initialize the exception with model name context.
+
+        Args:
+            model_name: Name of the HuggingFace model that was not found.
+        """
         self.model_name = model_name
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"The given model name '{self.model_name}' is not a valid model identifier listed on "
             "'https://huggingface.co/models'. "

arize/embeddings/nlp_generators.py

@@ -1,6 +1,7 @@
+"""NLP embedding generators for text classification and summarization tasks."""
+
 import logging
 from functools import partial
-from typing import Optional, cast
 
 import pandas as pd
 
@@ -22,11 +23,19 @@ logger = logging.getLogger(__name__)
 
 
 class EmbeddingGeneratorForNLPSequenceClassification(NLPEmbeddingGenerator):
+    """Embedding generator for NLP sequence classification tasks."""
+
     def __init__(
         self,
         model_name: str = DEFAULT_NLP_SEQUENCE_CLASSIFICATION_MODEL,
-        **kwargs,
-    ):
+        **kwargs: object,
+    ) -> None:
+        """Initialize the sequence classification embedding generator.
+
+        Args:
+            model_name: Name of the pre-trained NLP model.
+            **kwargs: Additional arguments for model initialization.
+        """
         super().__init__(
             use_case=UseCases.NLP.SEQUENCE_CLASSIFICATION,
             model_name=model_name,
@@ -36,15 +45,17 @@ class EmbeddingGeneratorForNLPSequenceClassification(NLPEmbeddingGenerator):
     def generate_embeddings(
         self,
         text_col: pd.Series,
-        class_label_col: Optional[pd.Series] = None,
+        class_label_col: pd.Series | None = None,
     ) -> pd.Series:
-        """
-        Obtain embedding vectors from your text data using pre-trained large language models.
+        """Obtain embedding vectors from your text data using pre-trained large language models.
+
+        Args:
+            text_col: A pandas Series containing the different pieces of text.
+            class_label_col: If this column is passed, the sentence "The classification label
+                is <class_label>" will be appended to the text in the `text_col`.
 
-        :param text_col: a pandas Series containing the different pieces of text.
-        :param class_label_col: if this column is passed, the sentence "The classification label
-        is <class_label>" will be appended to the text in the `text_col`.
-        :return: a pandas Series containing the embedding vectors.
+        Returns:
+            A pandas Series containing the embedding vectors.
         """
         if not isinstance(text_col, pd.Series):
             raise TypeError("text_col must be a pandas Series")
@@ -72,13 +83,24 @@ class EmbeddingGeneratorForNLPSequenceClassification(NLPEmbeddingGenerator):
             batched=True,
             batch_size=self.batch_size,
         )
-        return cast(pd.DataFrame, ds.to_pandas())["embedding_vector"]
+        df: pd.DataFrame = ds.to_pandas()
+        return df["embedding_vector"]
 
 
 class EmbeddingGeneratorForNLPSummarization(NLPEmbeddingGenerator):
+    """Embedding generator for NLP text summarization tasks."""
+
     def __init__(
-        self, model_name: str = DEFAULT_NLP_SUMMARIZATION_MODEL, **kwargs
-    ):
+        self,
+        model_name: str = DEFAULT_NLP_SUMMARIZATION_MODEL,
+        **kwargs: object,
+    ) -> None:
+        """Initialize the text summarization embedding generator.
+
+        Args:
+            model_name: Name of the pre-trained NLP model.
+            **kwargs: Additional arguments for model initialization.
+        """
         super().__init__(
             use_case=UseCases.NLP.SUMMARIZATION,
             model_name=model_name,
@@ -89,11 +111,13 @@ class EmbeddingGeneratorForNLPSummarization(NLPEmbeddingGenerator):
         self,
         text_col: pd.Series,
     ) -> pd.Series:
-        """
-        Obtain embedding vectors from your text data using pre-trained large language models.
+        """Obtain embedding vectors from your text data using pre-trained large language models.
+
+        Args:
+            text_col: A pandas Series containing the different pieces of text.
 
-        :param text_col: a pandas Series containing the different pieces of text.
-        :return: a pandas Series containing the embedding vectors.
+        Returns:
+            A pandas Series containing the embedding vectors.
         """
         if not isinstance(text_col, pd.Series):
             raise TypeError("text_col must be a pandas Series")
@@ -108,4 +132,5 @@ class EmbeddingGeneratorForNLPSummarization(NLPEmbeddingGenerator):
             batched=True,
             batch_size=self.batch_size,
         )
-        return cast(pd.DataFrame, ds.to_pandas())["embedding_vector"]
+        df: pd.DataFrame = ds.to_pandas()
+        return df["embedding_vector"]