Skip to content

ModelQuery

ModelQuery

A class which allows for querying a model serving endpoint on databricks.

This class is used to query a model serving endpoint on databricks with a dataset.

Example:

url = "https://example.com/model_endpoint"
token = "your_auth_token"

# Create an instance of ModelQuery
model_query = ModelQuery(url, token)

# Example dataset
dataset = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [4, 5, 6]})

try:
    # Score the model using the dataset
    response = model_query.score_model(dataset)
    print(response)
except requests.exceptions.HTTPError as e:
    print(f"Error: {str(e)}")

Parameters:

Name Type Description Default
url str

The URL of the model serving endpoint.

 required
token str

The authorization token for the model serving endpoint.

 required
Source code in src/fleming/discovery/model_query.py 
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
class ModelQuery:
    """
    A class which allows for querying a model serving endpoint on databricks.

    This class is used to query a model serving endpoint on databricks with a dataset.

    Example:
    --------
    ```python

    url = "https://example.com/model_endpoint"
    token = "your_auth_token"

    # Create an instance of ModelQuery
    model_query = ModelQuery(url, token)

    # Example dataset
    dataset = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [4, 5, 6]})

    try:
        # Score the model using the dataset
        response = model_query.score_model(dataset)
        print(response)
    except requests.exceptions.HTTPError as e:
        print(f"Error: {str(e)}")

    ```

    Parameters:
        url (str): The URL of the model serving endpoint.
        token (str): The authorization token for the model serving endpoint.
    """

    url: str
    token: str

    def __init__(self, url, token):
        self.url = url
        self.token = token

    def create_tf_serving_json(self, data):
        """
        Creates a JSON object for TensorFlow serving.

        Parameters:
            data (Union[dict, pd.DataFrame, np.ndarray]): The input data.

        Returns:
            dict: The JSON object for TensorFlow serving.
        """
        return {
            "inputs": (
                {name: data[name].tolist() for name in data.keys()}
                if isinstance(data, dict)
                else data.tolist()
            )
        }

    def score_model(self, dataset):
        """
        Scores the model using the provided dataset.

        Parameters:
            dataset (Union[pd.DataFrame, np.ndarray]): The dataset to be scored.

        Returns:
            dict: The response JSON from the model serving endpoint.

        Raises:
            requests.exceptions.HTTPError: If the request to the model serving endpoint fails.
        """
        headers = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
        }
        ds_dict = (
            {"dataframe_split": dataset.to_dict(orient="split")}
            if isinstance(dataset, pd.DataFrame)
            else self.create_tf_serving_json(dataset)
        )
        data_json = json.dumps(ds_dict, allow_nan=True)
        response = requests.request(
            method="POST", headers=headers, url=self.url, data=data_json
        )
        if response.status_code != 200:
            raise requests.exceptions.HTTPError(
                f"Request failed with status {response.status_code}, {response.text}"
            )
        return response.json()

create_tf_serving_json(data)

Creates a JSON object for TensorFlow serving.

Parameters:

Name Type Description Default
data Union[dict, DataFrame, ndarray]

The input data.

 required

Returns:

Name Type Description
dict

The JSON object for TensorFlow serving.
Source code in src/fleming/discovery/model_query.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def create_tf_serving_json(self, data):
    """
    Creates a JSON object for TensorFlow serving.

    Parameters:
        data (Union[dict, pd.DataFrame, np.ndarray]): The input data.

    Returns:
        dict: The JSON object for TensorFlow serving.
    """
    return {
        "inputs": (
            {name: data[name].tolist() for name in data.keys()}
            if isinstance(data, dict)
            else data.tolist()
        )
    }

score_model(dataset)

Scores the model using the provided dataset.

Parameters:

Name Type Description Default
dataset Union[DataFrame, ndarray]

The dataset to be scored.

 required

Returns:

Name Type Description
dict

The response JSON from the model serving endpoint.

Raises:

Type Description
HTTPError

If the request to the model serving endpoint fails.
Source code in src/fleming/discovery/model_query.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def score_model(self, dataset):
    """
    Scores the model using the provided dataset.

    Parameters:
        dataset (Union[pd.DataFrame, np.ndarray]): The dataset to be scored.

    Returns:
        dict: The response JSON from the model serving endpoint.

    Raises:
        requests.exceptions.HTTPError: If the request to the model serving endpoint fails.
    """
    headers = {
        "Authorization": f"Bearer {self.token}",
        "Content-Type": "application/json",
    }
    ds_dict = (
        {"dataframe_split": dataset.to_dict(orient="split")}
        if isinstance(dataset, pd.DataFrame)
        else self.create_tf_serving_json(dataset)
    )
    data_json = json.dumps(ds_dict, allow_nan=True)
    response = requests.request(
        method="POST", headers=headers, url=self.url, data=data_json
    )
    if response.status_code != 200:
        raise requests.exceptions.HTTPError(
            f"Request failed with status {response.status_code}, {response.text}"
        )
    return response.json()