lexsi-sdk 0.1.16__py3-none-any.whl

lexsi_sdk/core/synthetic.py

@@ -0,0 +1,387 @@
+from typing import Any, List, Optional
+from pydantic import BaseModel, ConfigDict
+
+import io
+import json
+import pandas as pd
+import plotly.graph_objects as go
+
+from lexsi_sdk.client.client import APIClient
+from lexsi_sdk.common.utils import poll_events, pretty_date
+from lexsi_sdk.common.validation import Validate
+from lexsi_sdk.common.xai_uris import (
+    AVAILABLE_SYNTHETIC_CUSTOM_SERVERS_URI,
+    GENERATE_ANONYMITY_SCORE_URI,
+    GENERATE_SYNTHETIC_DATA_URI,
+    GET_ANONYMITY_SCORE_URI,
+)
+
+
+class SyntheticDataTag(BaseModel):
+    """Representation of a generated synthetic data tag and its quality."""
+    api_client: APIClient
+    project_name: str
+    project: Any
+
+    model_name: str
+    tag: str
+    created_at: str
+
+    overall_quality_score: float
+    column_shapes: float
+    column_pair_trends: float
+
+    metadata: Optional[dict] = {}
+    plot_data: Optional[List[dict]] = []
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    def __init__(self, **kwargs):
+        """Bind API client reference for follow-up requests."""
+        super().__init__(**kwargs)
+        self.api_client = kwargs.get("api_client")
+
+    def get_model_name(self) -> str:
+        """get model type
+
+        :return: model type
+        """
+        return self.model_name
+
+    def view_metadata(self) -> dict:
+        """print metadata"""
+
+        print(json.dumps(self.metadata, indent=4))
+
+    def get_metadata(self) -> dict:
+        """get metadata"""
+
+        return self.metadata
+
+    def __print__(self) -> str:
+        """User-friendly string representation."""
+        created_at = pretty_date(self.created_at)
+        return f"SyntheticDataTag(model_name={self.model_name}, tag={self.tag}, created_at={created_at})"
+
+    def __str__(self) -> str:
+        """Return printable representation."""
+        return self.__print__()
+
+    def __repr__(self) -> str:
+        """Return developer-friendly representation."""
+        return self.__print__()
+
+
+class SyntheticModel(BaseModel):
+    """Synthetic Model Class
+
+    :param BaseModel: _description_
+    :return: _description_
+    """
+
+    api_client: APIClient
+    project_name: str
+    project: Any
+
+    model_name: str
+    status: str
+    created_at: str
+    created_by: str
+
+    overall_quality_score: Optional[float] = None
+    column_shapes: Optional[float] = None
+    column_pair_trends: Optional[float] = None
+
+    metadata: Optional[dict] = {}
+    plot_data: Optional[List[dict]] = []
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    def __init__(self, **kwargs):
+        """Bind API client reference for this synthetic model."""
+        super().__init__(**kwargs)
+        self.api_client = kwargs.get("api_client")
+
+    def get_model_type(self) -> str:
+        """get model type
+
+        :return: model type
+        """
+        return self.metadata["model_name"]
+
+    def get_data_quality(self) -> pd.DataFrame:
+        """get data quality
+
+        :return: data quality metrics
+        """
+        quality = {
+            "overall_quality_score": self.overall_quality_score,
+            "column_shapes": self.column_shapes,
+            "column_pair_trends": self.column_pair_trends,
+        }
+
+        df = pd.DataFrame(quality, index=[0])
+
+        return df
+
+    def quality_plot(self):
+        """plot psi chart"""
+        x_data = [item["Column"] for item in self.plot_data]
+        y_data = [item["Quality Score"] for item in self.plot_data]
+        metric_data = [item["Metric"] for item in self.plot_data]
+
+        traces = []
+        for metric in set(metric_data):
+            indices = [i for i, val in enumerate(metric_data) if val == metric]
+            traces.append(
+                go.Bar(
+                    x=[x_data[i] for i in indices],
+                    y=[y_data[i] for i in indices],
+                    name=metric,
+                )
+            )
+
+        fig = go.Figure(data=traces)
+
+        fig.update_layout(
+            barmode="relative",
+            # xaxis_title="Column Names",
+            # yaxis_title="Quality Score",
+            height=450,
+            bargap=0.01,
+            legend_orientation="h",
+            legend_x=0.1,
+            legend_y=1.1,
+        )
+
+        fig.show(config={"displaylogo": False})
+
+    '''
+    def get_training_logs(self) -> str:
+        """get model training logs
+
+        :return: logs of string type
+        """
+        url = f"{GET_SYNTHETIC_TRAINING_LOGS_URI}?project_name={self.project_name}&model_name={self.model_name}"
+
+        res = self.api_client.get(url)
+
+        if not res['success']:
+            raise Exception('Error while getting training logs.')
+
+        return res['details']
+    '''
+
+    def generate_synthetic_datapoints(
+        self, num_of_datapoints: int, instance_type: Optional[str] = "shared"
+    ):
+        """generate given number of synthetic datapoints
+
+        :param num_of_datapoints: total datapoints to generate
+        :param instance_type: type of instance to run training
+            for all available instances check xai.available_synthetic_custom_servers()
+            defaults to shared
+        :return: None
+        """
+        if instance_type != "shared":
+            available_servers = self.api_client.get(
+                AVAILABLE_SYNTHETIC_CUSTOM_SERVERS_URI
+            )["details"]
+            servers = list(
+                map(lambda instance: instance["instance_name"], available_servers)
+            )
+            Validate.value_against_list("instance_type", instance_type, servers)
+
+        payload = {
+            "project_name": self.project_name,
+            "model_name": self.model_name,
+            "instance_type": instance_type,
+            "num_of_datapoints": num_of_datapoints,
+        }
+
+        res = self.api_client.post(GENERATE_SYNTHETIC_DATA_URI, payload)
+
+        if not res["success"]:
+            raise Exception(res["details"])
+
+        print("Generating synthetic datapoints...")
+        poll_events(
+            self.api_client,
+            self.project_name,
+            res["event_id"],
+            progress_message="Synthetic Data generation progress",
+        )
+        print("Synthetic datapoints generated successfully.\n")
+
+    def generate_anonymity_score(
+        self,
+        aux_columns: List[str],
+        control_tag: str,
+        instance_type: Optional[str] = "shared",
+    ):
+        """generate anonymity score
+
+        :param aux_columns: list of features
+        :param control_tag: tag
+        :param instance_type: type of instance to run training
+            for all available instances check xai.available_synthetic_custom_servers()
+            defaults to shared
+
+        :return: None
+        """
+        if instance_type != "shared":
+            available_servers = self.api_client.get(
+                AVAILABLE_SYNTHETIC_CUSTOM_SERVERS_URI
+            )["details"]
+            servers = list(
+                map(lambda instance: instance["instance_name"], available_servers)
+            )
+            Validate.value_against_list("instance_type", instance_type, servers)
+
+        if len(aux_columns) < 2:
+            raise Exception("aux_columns requires minimum 2 columns.")
+
+        project_config = self.project.config()["metadata"]
+
+        Validate.value_against_list(
+            "feature", aux_columns, project_config["feature_include"]
+        )
+
+        all_tags = self.project.all_tags()
+
+        Validate.value_against_list("tag", [control_tag], all_tags)
+
+        payload = {
+            "aux_columns": aux_columns,
+            "control_tag": control_tag,
+            "model_name": self.model_name,
+            "project_name": self.project_name,
+            "instance_type": instance_type,
+        }
+
+        res = self.api_client.post(GENERATE_ANONYMITY_SCORE_URI, payload)
+
+        if not res["success"]:
+            raise Exception(res["details"])
+
+        print("Calculating anonymity score...")
+        poll_events(self.api_client, self.project_name, res["event_id"])
+        print("Anonymity score calculated successfully.\n")
+
+    def anonymity_score(self):
+        """get anonymity score
+
+        :raises Exception: _description_
+        :return: _description_
+        """
+        payload = {
+            "project_name": self.project_name,
+            "model_name": self.model_name,
+        }
+
+        res = self.api_client.post(GET_ANONYMITY_SCORE_URI, payload)
+
+        if not res["success"]:
+            print(res["details"])
+            raise Exception("Error while getting anonymity score.")
+
+        print("metadata:")
+        print(res["details"]["metadata"])
+        print("\n")
+
+        return pd.DataFrame(res["details"]["scores"], index=[0])
+
+    def __print__(self) -> str:
+        """User-friendly string representation."""
+        created_at = pretty_date(self.created_at)
+
+        return f"SyntheticModel(model_name={self.model_name}, status={self.status}, created_by={self.created_by}, created_at={created_at})"
+
+    def __str__(self) -> str:
+        """Return printable representation."""
+        return self.__print__()
+
+    def __repr__(self) -> str:
+        """Return developer-friendly representation."""
+        return self.__print__()
+
+
+class SyntheticPrompt(BaseModel):
+    """Metadata container for a reusable synthetic prompt definition."""
+    api_client: APIClient
+    project: Any
+
+    prompt_name: str
+    prompt_id: str
+    project_name: str
+    status: str
+    configuration: List[Any]
+    metadata: dict
+    created_by: str
+    updated_by: str
+    created_at: str
+    updated_at: str
+
+    def __init__(self, **kwargs):
+        """Bind API client reference for prompt actions."""
+        super().__init__(**kwargs)
+        self.api_client = kwargs.get("api_client")
+
+    def get_expression(self) -> str:
+        """construct prompt expression
+
+        :return: prompt expression
+        """
+        expression_list = []
+
+        if not self.metadata:
+            raise Exception("Expression not found.")
+
+        for item in self.metadata["expression"]:
+            if isinstance(item, dict):
+                expression_list.append(
+                    f"{item['column']} {item['expression']} {item['value']}"
+                )
+            else:
+                expression_list.append(item)
+
+        return " ".join(expression_list)
+
+    def get_config(self) -> List[dict]:
+        """get prompt configuration
+
+        :return: prompt configuration
+        """
+        return self.configuration
+
+    """
+    def delete(self) -> str:
+        payload = {
+            "delete": True,
+            "project_name": self.project_name,
+            "prompt_id": self.prompt_id,
+            "update_keys": {}
+        }
+
+        res = self.api_client.post(UPDATE_SYNTHETIC_PROMPT_URI, payload)
+
+        if not res['success']:
+            raise Exception(res['details'])
+        
+        return res['details']
+    """
+
+    def __print__(self) -> str:
+        """User-friendly string representation."""
+        created_at = pretty_date(self.created_at)
+        updated_at = pretty_date(self.updated_at)
+
+        return f"SyntheticPrompt(prompt_name={self.prompt_name}, prompt_id={self.prompt_id}, status={self.status}, created_by={self.created_by}, created_at={created_at}, updated_at={updated_at})"
+
+    def __str__(self) -> str:
+        """Return printable representation."""
+        return self.__print__()
+
+    def __repr__(self) -> str:
+        """Return developer-friendly representation."""
+        return self.__print__()