Quick Start

pip install kumoai

import kumoai.experimental.rfm as rfm, os

os.environ["KUMO_API_KEY"] = "ENTER_YOUR_API_KEY_HERE"

rfm.init()

import pandas as pd

dataset_url = "s3://kumo-sdk-public/rfm-datasets/online-shopping"

users_df = pd.read_parquet(f"{dataset_url}/users.parquet")
items_df = pd.read_parquet(f"{dataset_url}/items.parquet")
orders_df = pd.read_parquet(f"{dataset_url}/orders.parquet")

graph = rfm.LocalGraph.from_data({
    "users": users_df,
    "items": items_df,
    "orders": orders_df,
})

# Inspect the graph - requires graphviz to be installed
graph.visualize()

model = rfm.KumoRFM(graph)

# Forecast 30-day product demand
query = "PREDICT SUM(orders.price, 0, 30, days) FOR items.item_id=1"
prediction_result = model.predict(query)
print(prediction_result)

# Predict customer churn
query = "PREDICT COUNT(orders.*, 0, 90, days)=0 FOR users.user_id IN (42, 123)"
prediction_result = model.predict(query)
print(prediction_result)

# Item recommendation
query = "PREDICT LIST_DISTINCT(orders.item_id, 0, 30, days) RANK TOP 10 FOR users.user_id=123"
df = model.predict(query)
display(df)

# Missing value imputation
query = "PREDICT users.age FOR users.user_id=8"
df = model.predict(query)
display(df)