kumoai.experimental.rfm.KumoRFM

kumoai.experimental.rfm.KumoRFM#

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

Bases: object

The Kumo Relational Foundation model (RFM) from the KumoRFM: A Foundation Model for In-Context Learning on Relational Data paper.

KumoRFM is a foundation model to generate predictions for any relational dataset without training. The model is pre-trained and the class provides an interface to query the model from a LocalGraph object.

from kumoai.experimental.rfm import LocalGraph, KumoRFM

df_users = pd.DataFrame(...)
df_items = pd.DataFrame(...)
df_orders = pd.DataFrame(...)

graph = LocalGraph.from_data({
    'users': df_users,
    'items': df_items,
    'orders': df_orders,
})

rfm = KumoRFM(graph)

query = ("PREDICT COUNT(transactions.*, 0, 30, days)>0 "
         "FOR users.user_id=0")
result = rfm.query(query)

print(result)  # user_id  COUNT(transactions.*, 0, 30, days) > 0
               # 1        0.85
Parameters:

  • graph (LocalGraph) – The graph.

  • preprocess (bool) – Whether to pre-process the data in advance during graph materialization. This is a runtime trade-off between graph materialization and model processing speed. It can be benefical to preprocess your data once and then run many queries on top to achieve maximum model speed. However, if activiated, graph materialization can take potentially much longer, especially on graphs with many large text columns. Best to tune this option manually.

  • verbose (bool) – Whether to print verbose output.

__init__(graph, preprocess=False, verbose=True)[source]#
predict(query, *, anchor_time=None, run_mode=fast, num_neighbors=None, num_hops=2, max_pq_iterations=20, random_seed=42, verbose=True)[source]#

Returns predictions for a predictive query.

Parameters:

  • query (str) – The predictive query.

  • anchor_time (Union[Timestamp, Literal['entity'], None]) – The anchor timestamp for the query. If set to None, will use the maximum timestamp in the data. If set to :”entity”, will use the timestamp of the entity.

  • run_mode (Union[RunMode, str]) – The RunMode for the query.

  • num_neighbors (Optional[List[int]]) – The number of neighbors to sample for each hop. If specified, the num_hops option will be ignored.

  • num_hops (int) – The number of hops to sample when generating the context.

  • max_pq_iterations (int) – The maximum number of iterations to perform to collect valid labels. It is advised to increase the number of iterations in case the predictive query has strict entity filters, in which case, KumoRFM needs to sample more entities to find valid labels.

  • random_seed (Optional[int]) – A manual seed for generating pseudo-random numbers.

  • verbose (bool) – Whether to print verbose output.

Return type:

DataFrame

Returns:

The predictions as a pandas.DataFrame

evaluate(query, *, metrics=None, anchor_time=None, run_mode=fast, num_neighbors=None, num_hops=2, max_pq_iterations=20, random_seed=42, verbose=True)[source]#

Evaluates a predictive query.

Parameters:

  • query (str) – The predictive query.

  • metrics (Optional[List[str]]) – The metrics to use.

  • anchor_time (Union[Timestamp, Literal['entity'], None]) – The anchor timestamp for the query. If set to None, will use the maximum timestamp in the data. If set to :”entity”, will use the timestamp of the entity.

  • run_mode (Union[RunMode, str]) – The RunMode for the query.

  • num_neighbors (Optional[List[int]]) – The number of neighbors to sample for each hop. If specified, the num_hops option will be ignored.

  • num_hops (int) – The number of hops to sample when generating the context.

  • max_pq_iterations (int) – The maximum number of iterations to perform to collect valid labels. It is advised to increase the number of iterations in case the predictive query has strict entity filters, in which case, KumoRFM needs to sample more entities to find valid labels.

  • random_seed (Optional[int]) – A manual seed for generating pseudo-random numbers.

  • verbose (bool) – Whether to print verbose output.

Return type:

DataFrame

Returns:

The metrics as a pandas.DataFrame

get_train_table(query, size, *, anchor_time=None, random_seed=42, max_iterations=20)[source]#

Returns the labels of a predictive query for a specified anchor time.

Parameters:

  • query (str) – The predictive query.

  • size (int) – The maximum number of entities to generate labels for.

  • anchor_time (Union[Timestamp, Literal['entity'], None]) – The anchor timestamp for the query. If set to None, will use the maximum timestamp in the data. If set to :”entity”, will use the timestamp of the entity.

  • random_seed (Optional[int]) – A manual seed for generating pseudo-random numbers.

  • max_iterations (int) – The number of steps to run before aborting.

Return type:

DataFrame

Returns:

The labels as a pandas.DataFrame.

Contents