kumoai.experimental.rfm.KumoRFM

kumoai.experimental.rfm.KumoRFM#

class kumoai.experimental.rfm.KumoRFM[source]#

Bases: object

The KumoRFM class is an interface to the Kumo Relational Foundation model (RFM), see KumoRFM.

KumoRFM is a relational foundation model, which can make generate predicitons in context for any relational dataset. The model is pretrained and the class provides an interface to query the model.

The class is constructed from a LocalGraph object.

Example

# raw data
df_users = pd.DataFrame(...)
df_items = pd.DataFrame(...)
df_transactions = pd.DataFrame(...)

# construct LocalGraph from raw data
graph = LocalGraph.from_data(
    {
        'users': df_users,
        'items': df_items,
        'transactions': df_transactions,
    }
)

# start KumoRFM with this graph
rfm = KumoRFM(graph)

# query the model
query_str = ("PREDICT COUNT(transactions.*, 0, 30, days) > 0 "
             "FOR users.user_id=1")
result = rfm.query(query_str)
# Result is a pandas DataFrame with prediction probabilities
print(result)  # user_id  COUNT(transactions.*, 0, 30, days) > 0
               # 1        0.85
__init__(graph, endpoint=None)[source]#
query(query, wait_result_timeout=10.0)[source]#

Query the RFM model with a query string

Parameters:

  • query (str) – The RFM query string

  • (e.g. (orders.*, 0, 30, days) –

  • COUNT ("PREDICT) –

  • users.user_id=1") ("FOR) –

  • wait_result_timeout (float) – Timeout in seconds to wait for the result

Return type:

Dict[str, Any]

Returns:

Dictionary containing the prediction result

shutdown()[source]#

Clean up resources by destroying the RFM endpoint.

This method attempts to destroy the RFM endpoint if one exists.

Return type:

None

Example

rfm = KumoRFM(graph)
rfm.wait_until_ready()
# ... use RFM ...
rfm.shutdown()  # Clean up endpoint when done
poll_status()[source]#

Poll the current status of the RFM endpoint.

Return type:

Optional[OnlineServingStatusCode]

Returns:

The current status code of the endpoint, or None if no endpoint exists.

Possible values:

  • OnlineServingStatusCode.IN_PROGRESS: Endpoint is being created/updated

  • OnlineServingStatusCode.READY: Endpoint is ready for queries

  • OnlineServingStatusCode.FAILED: Endpoint creation/update failed

is_ready()[source]#

Check if the RFM endpoint is ready for queries.

Return type:

bool

Returns:

True if the endpoint is ready, False otherwise.

wait_until_ready(timeout=None, sleep_interval=10.0)[source]#

Wait until the RFM endpoint is ready.

Parameters:

  • timeout (Optional[float]) – Maximum time to wait in seconds. If None, waits indefinitely.

  • sleep_interval (float) – Time to sleep between polls in seconds.

Raises:

  • TimeoutError – If the endpoint is not ready within the timeout period.

  • RuntimeError – If the endpoint failed to start.

Return type:

None

Contents