validmind 2.0.1__py3-none-any.whl → 2.0.7__py3-none-any.whl

validmind/vm_models/dataset.py

@@ -11,6 +11,7 @@ from dataclasses import dataclass, field
 
 import numpy as np
 import pandas as pd
+import polars as pl
 
 from validmind.logging import get_logger
 from validmind.vm_models.model import VMModel
@@ -35,9 +36,216 @@ class VMDataset(ABC):
         pass
 
     @abstractmethod
-    def serialize(self):
+    def assign_predictions(
+        self,
+        model,
+        prediction_values: list = None,
+        prediction_column=None,
+    ):
         """
-        Serializes the dataset to a dictionary.
+        Assigns predictions to the dataset for a given model or prediction values.
+        The dataset is updated with a new column containing the predictions.
+        """
+        pass
+
+    @abstractmethod
+    def get_extra_column(self, column_name):
+        """
+        Returns the values of the specified extra column.
+
+        Args:
+            column_name (str): The name of the extra column.
+
+        Returns:
+            np.ndarray: The values of the extra column.
+        """
+        pass
+
+    @abstractmethod
+    def add_extra_column(self, column_name, column_values=None):
+        """
+        Adds an extra column to the dataset without modifying the dataset `features` and `target` columns.
+
+        Args:
+            column_name (str): The name of the extra column.
+            column_values (np.ndarray, optional): The values of the extra column.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def input_id(self) -> str:
+        """
+        Returns input id of dataset.
+
+        Returns:
+            str: input_id.
+        """
+        return self.input_id
+
+    @property
+    @abstractmethod
+    def columns(self) -> list:
+        """
+        Returns the the list of columns in the dataset.
+
+        Returns:
+            List[str]: The columns list.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def target_column(self) -> str:
+        """
+        Returns the target column name of the dataset.
+
+        Returns:
+            str: The target column name.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def feature_columns(self) -> list:
+        """
+        Returns the feature columns of the dataset. If _feature_columns is None,
+        it returns all columns except the target column.
+
+        Returns:
+            list: The list of feature column names.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def text_column(self) -> str:
+        """
+        Returns the text column of the dataset.
+
+        Returns:
+            str: The text column name.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def x(self) -> np.ndarray:
+        """
+        Returns the input features (X) of the dataset.
+
+        Returns:
+            np.ndarray: The input features.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def y(self) -> np.ndarray:
+        """
+        Returns the target variables (y) of the dataset.
+
+        Returns:
+            np.ndarray: The target variables.
+        """
+        pass
+
+    @abstractmethod
+    def y_pred(self, model_id) -> np.ndarray:
+        """
+        Returns the prediction values (y_pred) of the dataset for a given model_id.
+
+        Returns:
+            np.ndarray: The prediction values.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def df(self):
+        """
+        Returns the dataset as a pandas DataFrame.
+
+        Returns:
+            pd.DataFrame: The dataset as a DataFrame.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def copy(self):
+        """
+        Returns a copy of the raw_dataset dataframe.
+        """
+        pass
+
+    @abstractmethod
+    def x_df(self):
+        """
+        Returns the non target and prediction columns.
+
+        Returns:
+            pd.DataFrame: The non target and prediction columns .
+        """
+        pass
+
+    @abstractmethod
+    def y_df(self):
+        """
+        Returns the target columns (y) of the dataset.
+
+        Returns:
+            pd.DataFrame: The target columns.
+        """
+        pass
+
+    @abstractmethod
+    def y_pred_df(self, model_id):
+        """
+        Returns the target columns (y) of the dataset.
+
+        Returns:
+            pd.DataFrame: The target columns.
+        """
+        pass
+
+    @abstractmethod
+    def prediction_column(self, model_id) -> str:
+        """
+        Returns the prediction column name of the dataset.
+
+        Returns:
+            str: The prediction column name.
+        """
+        pass
+
+    @abstractmethod
+    def get_features_columns(self):
+        """
+        Returns the column names of the feature variables.
+
+        Returns:
+            List[str]: The column names of the feature variables.
+        """
+        pass
+
+    @abstractmethod
+    def get_numeric_features_columns(self):
+        """
+        Returns the column names of the numeric feature variables.
+
+        Returns:
+            List[str]: The column names of the numeric feature variables.
+        """
+        pass
+
+    @abstractmethod
+    def get_categorical_features_columns(self):
+        """
+        Returns the column names of the categorical feature variables.
+
+        Returns:
+            List[str]: The column names of the categorical feature variables.
         """
         pass
 
@@ -134,7 +342,7 @@ class NumpyDataset(VMDataset):
         # initialize target column
         self._target_column = target_column
         # initialize extra columns
-        self.__set_extra_columns(extra_columns, model)
+        self.__set_extra_columns(extra_columns)
         # initialize feature columns
         self.__set_feature_columns(feature_columns)
         # initialize text column, target class labels and options
@@ -144,7 +352,7 @@ class NumpyDataset(VMDataset):
         if model:
             self.assign_predictions(model)
 
-    def __set_extra_columns(self, extra_columns, model):
+    def __set_extra_columns(self, extra_columns):
         if extra_columns is None:
             extra_columns = {
                 "prediction_columns": {},
@@ -191,14 +399,29 @@ class NumpyDataset(VMDataset):
         return model.input_id in self._extra_columns.get("prediction_columns", {})
 
     def __assign_prediction_values(self, model, pred_column, prediction_values):
+        # Link the prediction column with the model
         self._extra_columns.setdefault("prediction_columns", {})[
             model.input_id
         ] = pred_column
-        self._raw_dataset = np.hstack(
-            (self._raw_dataset, np.array(prediction_values).reshape(-1, 1))
+
+        # Check if the predictions are multi-dimensional (e.g., embeddings)
+        is_multi_dimensional = (
+            isinstance(prediction_values, np.ndarray) and prediction_values.ndim > 1
         )
-        self._columns.append(pred_column)
-        self._df[pred_column] = prediction_values
+
+        if is_multi_dimensional:
+            # For multi-dimensional outputs, convert to a list of lists to store in DataFrame
+            self._df[pred_column] = list(map(list, prediction_values))
+        else:
+            # If not multi-dimensional or a standard numpy array, reshape for compatibility
+            self._raw_dataset = np.hstack(
+                (self._raw_dataset, np.array(prediction_values).reshape(-1, 1))
+            )
+            self._df[pred_column] = prediction_values
+
+        # Update the dataset columns list
+        if pred_column not in self._columns:
+            self._columns.append(pred_column)
 
     def assign_predictions(  # noqa: C901 - we need to simplify this method
         self,
@@ -262,6 +485,59 @@ class NumpyDataset(VMDataset):
                 f"Prediction column {self._extra_columns['prediction_columns'][model.input_id]} already linked to the {model.input_id}"
             )
 
+    def get_extra_column(self, column_name):
+        """
+        Returns the values of the specified extra column.
+
+        Args:
+            column_name (str): The name of the extra column.
+
+        Returns:
+            np.ndarray: The values of the extra column.
+        """
+        if column_name not in self.extra_columns:
+            raise ValueError(f"Column {column_name} is not an extra column")
+
+        return self._df[column_name]
+
+    def add_extra_column(self, column_name, column_values=None):
+        """
+        Adds an extra column to the dataset without modifying the dataset `features` and `target` columns.
+
+        Args:
+            column_name (str): The name of the extra column.
+            column_values (np.ndarray, optional): The values of the extra column.
+        """
+        if column_name in self.extra_columns:
+            logger.info(f"Column {column_name} already registered as an extra column")
+            return
+
+        # The column name already exists in the dataset so we just assign the extra column
+        if column_name in self.columns:
+            self._extra_columns[column_name] = column_name
+            logger.info(
+                f"Column {column_name} exists in the dataset, registering as an extra column"
+            )
+            return
+
+        if column_values is None:
+            raise ValueError(
+                "Column values must be provided when the column doesn't exist in the dataset"
+            )
+
+        if len(column_values) != self.df.shape[0]:
+            raise ValueError(
+                "Length of column values doesn't match number of rows of the dataset"
+            )
+
+        self._raw_dataset = np.hstack(
+            (self._raw_dataset, np.array(column_values).reshape(-1, 1))
+        )
+        self._columns.append(column_name)
+        self._df[column_name] = column_values
+        self._extra_columns[column_name] = column_name
+        logger.info(f"Column {column_name} added as an extra column")
+
     @property
     def raw_dataset(self) -> np.ndarray:
         """
@@ -399,19 +675,42 @@ class NumpyDataset(VMDataset):
 
     def y_pred(self, model_id) -> np.ndarray:
         """
-        Returns the prediction variable (y_pred) of the dataset.
+        Returns the prediction variables for a given model_id, accommodating
+        both scalar predictions and multi-dimensional outputs such as embeddings.
+
+        Args:
+            model_id (str): The ID of the model whose predictions are sought.
 
         Returns:
-            np.ndarray: The prediction variables.
-        """
-        return self.raw_dataset[
-            :,
-            [
-                self.columns.index(name)
-                for name in self.columns
-                if name == self.prediction_column(model_id=model_id)
-            ],
-        ].flatten()
+            np.ndarray: The prediction variables, either as a flattened array for
+            scalar predictions or as an array of arrays for multi-dimensional outputs.
+        """
+        pred_column = self.prediction_column(model_id)
+
+        # First, attempt to retrieve the prediction data from the DataFrame
+        if hasattr(self, "_df") and pred_column in self._df.columns:
+            predictions = self._df[pred_column].to_numpy()
+
+            # Check if the predictions are stored as objects (e.g., lists for embeddings)
+            if self._df[pred_column].dtype == object:
+                # Attempt to convert lists to a numpy array
+                try:
+                    predictions = np.stack(predictions)
+                except ValueError as e:
+                    # Handling cases where predictions cannot be directly stacked
+                    raise ValueError(f"Error stacking prediction arrays: {e}")
+        else:
+            # Fallback to using the raw numpy dataset if DataFrame is not available or suitable
+            try:
+                predictions = self.raw_dataset[
+                    :, self.columns.index(pred_column)
+                ].flatten()
+            except IndexError as e:
+                raise ValueError(
+                    f"Prediction column '{pred_column}' not found in raw dataset: {e}"
+                )
+
+        return predictions
 
     @property
     def type(self) -> str:
@@ -608,10 +907,15 @@ class DataFrameDataset(NumpyDataset):
 
         Args:
             raw_dataset (pd.DataFrame): The raw dataset as a pandas DataFrame.
+            input_id (str, optional): Identifier for the dataset. Defaults to None.
+            model (VMModel, optional): Model associated with the dataset. Defaults to None.
             target_column (str, optional): The target column of the dataset. Defaults to None.
+            extra_columns (dict, optional): Extra columns to include in the dataset. Defaults to None.
             feature_columns (list, optional): The feature columns of the dataset. Defaults to None.
-            text_column (str, optional): The text column name of the dataset for nlp tasks. Defaults to None.
-            target_class_labels (Dict, optional): The class labels for the target columns. Defaults to None.
+            text_column (str, optional): The text column name of the dataset for NLP tasks. Defaults to None.
+            target_class_labels (dict, optional): The class labels for the target columns. Defaults to None.
+            options (dict, optional): Additional options for the dataset. Defaults to None.
+            date_time_index (bool, optional): Whether to use date-time index. Defaults to False.
         """
         index = None
         if isinstance(raw_dataset.index, pd.Index):
@@ -634,6 +938,57 @@ class DataFrameDataset(NumpyDataset):
         )
 
 
+@dataclass
+class PolarsDataset(NumpyDataset):
+    """
+    VM dataset implementation for Polars DataFrame.
+    """
+
+    def __init__(
+        self,
+        raw_dataset: pl.DataFrame,
+        input_id: str = None,
+        model: VMModel = None,
+        target_column: str = None,
+        extra_columns: dict = None,
+        feature_columns: list = None,
+        text_column: str = None,
+        target_class_labels: dict = None,
+        options: dict = None,
+        date_time_index: bool = False,
+    ):
+        """
+        Initializes a PolarsDataset instance.
+
+        Args:
+            raw_dataset (pl.DataFrame): The raw dataset as a Polars DataFrame.
+            input_id (str, optional): Identifier for the dataset. Defaults to None.
+            model (VMModel, optional): Model associated with the dataset. Defaults to None.
+            target_column (str, optional): The target column of the dataset. Defaults to None.
+            extra_columns (dict, optional): Extra columns to include in the dataset. Defaults to None.
+            feature_columns (list, optional): The feature columns of the dataset. Defaults to None.
+            text_column (str, optional): The text column name of the dataset for NLP tasks. Defaults to None.
+            target_class_labels (dict, optional): The class labels for the target columns. Defaults to None.
+            options (dict, optional): Additional options for the dataset. Defaults to None.
+            date_time_index (bool, optional): Whether to use date-time index. Defaults to False.
+        """
+        super().__init__(
+            raw_dataset=raw_dataset.to_numpy(),
+            input_id=input_id,
+            model=model,
+            index_name=None,
+            index=None,
+            columns=raw_dataset.columns,
+            target_column=target_column,
+            extra_columns=extra_columns,
+            feature_columns=feature_columns,
+            text_column=text_column,
+            target_class_labels=target_class_labels,
+            options=options,
+            date_time_index=date_time_index,
+        )
+
+
 @dataclass
 class TorchDataset(NumpyDataset):
     """

validmind/vm_models/figure.py

@@ -21,6 +21,18 @@ from ..errors import InvalidFigureForObjectError, UnsupportedFigureError
 from ..utils import get_full_typename
 
 
+def is_matplotlib_figure(figure) -> bool:
+    return isinstance(figure, matplotlib.figure.Figure)
+
+
+def is_plotly_figure(figure) -> bool:
+    return isinstance(figure, (go.Figure, go.FigureWidget))
+
+
+def is_png_image(figure) -> bool:
+    return isinstance(figure, bytes)
+
+
 @dataclass
 class Figure:
     """
@@ -52,22 +64,10 @@ class Figure:
         if (
             not client_config.running_on_colab
             and self.figure
-            and self.is_plotly_figure()
+            and is_plotly_figure(self.figure)
         ):
             self.figure = go.FigureWidget(self.figure)
 
-    def is_matplotlib_figure(self) -> bool:
-        """
-        Returns True if the figure is a matplotlib figure
-        """
-        return isinstance(self.figure, matplotlib.figure.Figure)
-
-    def is_plotly_figure(self) -> bool:
-        """
-        Returns True if the figure is a plotly figure
-        """
-        return isinstance(self.figure, (go.Figure, go.FigureWidget))
-
     def _get_for_object_type(self):
         """
         Returns the type of the object this figure is for
@@ -91,7 +91,7 @@ class Figure:
         we would render images as-is, but Plotly FigureWidgets don't work well
         on Google Colab when they are combined with ipywidgets.
         """
-        if self.is_matplotlib_figure():
+        if is_matplotlib_figure(self.figure):
             tmpfile = BytesIO()
             self.figure.savefig(tmpfile, format="png")
             encoded = base64.b64encode(tmpfile.getvalue()).decode("utf-8")
@@ -101,7 +101,7 @@ class Figure:
                 """
             )
 
-        elif self.is_plotly_figure():
+        elif is_plotly_figure(self.figure):
             # FigureWidget can be displayed as-is but not on Google Colab. In this case
             # we just return the image representation of the figure.
             if client_config.running_on_colab:
@@ -114,6 +114,15 @@ class Figure:
                 )
             else:
                 return self.figure
+
+        elif is_png_image(self.figure):
+            encoded = base64.b64encode(self.figure).decode("utf-8")
+            return widgets.HTML(
+                value=f"""
+                <img style="width:100%; height: auto;" src="data:image/png;base64,{encoded}"/>
+                """
+            )
+
         else:
             raise UnsupportedFigureError(
                 f"Figure type {type(self.figure)} not supported for plotting"
@@ -129,15 +138,38 @@ class Figure:
             "metadata": json.dumps(self.metadata, allow_nan=False),
         }
 
+    def _get_b64_url(self):
+        """
+        Returns a base64 encoded URL for the figure
+        """
+        if is_matplotlib_figure(self.figure):
+            buffer = BytesIO()
+            self.figure.savefig(buffer, format="png")
+            buffer.seek(0)
+
+            b64_data = base64.b64encode(buffer.read()).decode("utf-8")
+
+            return f"data:image/png;base64,{b64_data}"
+
+        elif is_plotly_figure(self.figure):
+            bytes = self.figure.to_image(format="png")
+            b64_data = base64.b64encode(bytes).decode("utf-8")
+
+            return f"data:image/png;base64,{b64_data}"
+
+        raise UnsupportedFigureError(
+            f"Unrecognized figure type: {get_full_typename(self.figure)}"
+        )
+
     def serialize_files(self):
         """Creates a `requests`-compatible files object to be sent to the API"""
-        if self.is_matplotlib_figure():
+        if is_matplotlib_figure(self.figure):
             buffer = BytesIO()
             self.figure.savefig(buffer, bbox_inches="tight")
             buffer.seek(0)
             return {"image": (f"{self.key}.png", buffer, "image/png")}
 
-        elif self.is_plotly_figure():
+        elif is_plotly_figure(self.figure):
             # When using plotly, we need to use we will produce two files:
             # - a JSON file that will be used to display the figure in the UI
             # - a PNG file that will be used to display the figure in documents
@@ -154,6 +186,9 @@ class Figure:
                 ),
             }
 
+        elif is_png_image(self.figure):
+            return {"image": (f"{self.key}.png", self.figure, "image/png")}
+
         raise UnsupportedFigureError(
             f"Unrecognized figure type: {get_full_typename(self.figure)}"
         )

validmind/vm_models/test/metric.py

@@ -6,12 +6,14 @@
 Class for storing ValidMind metric objects and associated
 data for display and reporting purposes
 """
+import os
 from abc import abstractmethod
 from dataclasses import dataclass
 from typing import ClassVar, List, Optional, Union
 
 import pandas as pd
 
+from ...ai import generate_description
 from ...errors import MissingCacheResultsArgumentsError
 from ...utils import clean_docstring
 from ..figure import Figure
@@ -74,41 +76,42 @@ class Metric(Test):
                 "Metric must provide a metric value or figures to cache_results"
             )
 
-        # At a minimum, send the metric description
-        result_metadata = [
-            {
-                "content_id": f"metric_description:{self.test_id}",
-                "text": clean_docstring(self.description()),
-            }
-        ]
-
-        result_summary = self.summary(metric_value)
-
-        result_wrapper = MetricResultWrapper(
-            result_id=self.test_id,
-            result_metadata=result_metadata,
-            inputs=self.get_accessed_inputs(),
-            output_template=self.output_template,
-        )
-
-        # We can send an empty result to push an empty metric with a summary and plots
-        metric_result_value = metric_value if metric_value is not None else {}
-
-        result_wrapper.metric = MetricResult(
-            # key=self.key,
-            # Now using the fully qualified test ID as `key`.
-            # Ideally the backend is updated to use `test_id` instead of `key`.
+        metric = MetricResult(
             key=self.test_id,
             ref_id=self._ref_id,
-            value=metric_result_value,
+            value=metric_value if metric_value is not None else {},
             value_formatter=self.value_formatter,
-            summary=result_summary,
+            summary=self.summary(metric_value),
         )
 
-        # Allow metrics to attach figures to the test suite result
-        if figures:
-            result_wrapper.figures = figures
+        if (
+            os.environ.get("VALIDMIND_LLM_DESCRIPTIONS_ENABLED", "false").lower()
+            == "true"
+        ):
+            revision_name = "Generated by ValidMind AI"
+            description = generate_description(
+                test_name=self.test_id,
+                test_description=self.description().splitlines()[0],
+                test_results=metric.serialize()["value"],
+                test_summary=metric.serialize()["summary"],
+                figures=figures,
+            )
+        else:
+            revision_name = "Default Description"
+            description = clean_docstring(self.description())
+
+        description_metadata = {
+            "content_id": f"metric_description:{self.test_id}::{revision_name}",
+            "text": description,
+        }
 
-        self.result = result_wrapper
+        self.result = MetricResultWrapper(
+            result_id=self.test_id,
+            result_metadata=[description_metadata],
+            metric=metric,
+            figures=figures,
+            inputs=self.get_accessed_inputs(),
+            output_template=self.output_template,
+        )
 
         return self.result