Predict with Custom Context

Use predict_task() when you want to provide the labeled examples that Kumo Relational uses as context. Unlike predict(), which derives context examples from a Predictive Query Language (PQL) query, predict_task() accepts the context and prediction entities as pandas DataFrames.

You still bind a relational graph to the client. The graph provides the tables, relationships, and metadata that Kumo Relational uses to construct relational context.

Before you begin

You need:

• A configured and validated Kumo Relational graph.
• A RelationalClient connected to a Kumo Relational NIM.
• A context DataFrame containing labeled examples.
• A prediction DataFrame containing the entities to score.

The table named by entity_table must exist in the graph.

1. Prepare the input tables

By default, predict_task() expects the following columns:

ENTITY values reference entities from entity_table. TARGET contains the known label or value for each context example.

To use different column names, set entity_column, target_column, or time_column in the predict_task() call.

Time-column behavior

Kumo Relational selects the task time column in the following order:

1. The column passed through time_column.
2. ANCHOR_TIMESTAMP, when that column appears in either input DataFrame.
3. The time column configured for the entity table.

When you explicitly set time_column, the context DataFrame must contain that column.

2. Select the task type

Set task_type to one of the task types supported by Kumo Relational:

• binary_classification
• multiclass_classification
• regression
• forecasting
• temporal_link_prediction

You can inspect the supported values through client.capabilities("kumo-relational").tasks.

For classification, regression, and forecasting tasks, set entity_table to the name of the entity table in the graph. For temporal link prediction, pass the source and target entity-table names as a pair.

Forecasting requests require step_size and also accept num_forecasts. step_size is an integer number of nanoseconds. For example, calculate a 30-day step with pandas:

import pandas as pd
step_size = int(pd.Timedelta(days=30).value)

See Forecasting for background on Kumo Relational forecasting tasks.

For temporal link prediction, entity_table must be a (source_table, target_table) pair and each context target must be a list of target IDs.

3. Run the prediction

Bind the graph to the client and call predict_task():

from kumo_relational_client import RelationalClient
with RelationalClient(url="http://localhost:8000") as client:
    result = client.relational(graph).predict_task(
        context=train_df,      # ENTITY, TARGET, [ANCHOR_TIMESTAMP]
        predict=predict_df,    # ENTITY, [ANCHOR_TIMESTAMP]
        task_type="multiclass_classification",
        entity_table="users",
    )

This example uses rows in train_df as labeled context and scores the entities in predict_df.

Configure the prediction

predict_task() supports the same prediction controls used by the PQL workflow, including:

• run_mode
• num_neighbors
• num_hops
• inference_config
• use_prediction_time
• return_embeddings
• random_seed
• explain
• batch_size and num_retries
• top_k
• exclude_cols_dict
• num_forecasts and step_size

return_embeddings=True is supported for binary classification and regression. It appends an EMBEDDINGS column to each returned row without changing the task’s row count.

See Configuration for inference and neighborhood settings, Batch Prediction for batching behavior, and Prediction Explainability to request an explanation.

Work with the result

By default, predict_task() returns a pandas DataFrame. The result follows the same task-specific schema as a PQL-based prediction. When explain=True, it returns an Explanation object whose prediction attribute contains the prediction DataFrame.

See Prediction Results for the fields returned by each task.

Troubleshoot input validation

The SDK rejects the request before prediction when:

• task_type is not supported.
• entity_table does not exist in the graph.
• The context DataFrame is missing the entity or target column.
• The prediction DataFrame is missing the entity column.
• An explicitly selected time column is missing from the context DataFrame.

Use the column names reported in the error to correct the DataFrame or the corresponding entity_column, target_column, or time_column argument.