Apply to Your Data

from pathlib import Path

from monad.ui import pretrain


# --- Names & Paths -----------------------------------------------------------

# EDIT: provide path to your foundation model config (yaml)
# it defines data sources, splits, training params, and resource constraints
config_path = Path("/basemodel/configs/<foundation_model_config.yaml>").resolve()

# EDIT: define where foundation model outputs will be saved (this becomes your project_dir)
# the script will create /fm, /features, /lightning_checkpoints etc. under this directory
project_dir = Path("/basemodel/projects/project_dir").resolve()


# --- Pretrain ----------------------------------------------------------------

# CHECK: overwrite=True starts training from scratch and replaces existing outputs in project_dir
# replace it with resume=True to resume training if partial outputs already exist
pretrain(
    config_path=config_path,
    output_path=project_dir,
    overwrite=True,
)

# For context refer to docs:
# FM stage overview - https://docs.basemodel.ai/docs/pretrain-stage-overview
# FM training - https://docs.basemodel.ai/docs/training-foundation-model
# End-to-end tutorial - https://docs.basemodel.ai/docs/e2e-tutorial-kaggle-hm

from datetime import timedelta
from pathlib import Path
from typing import Dict

import numpy as np

from monad.batch import SPLIT_TIMESTAMP
from monad.ui.config import TrainingParams
from monad.ui.module import BinaryClassificationTask, load_from_foundation_model
from monad.ui.target_function import Attributes, Events, has_incomplete_training_window


# --- Names & Paths -----------------------------------------------------------

# EDIT: provide path to project directory, PARENT to /fm, /features, /lightning_checkpoints etc.
project_dir = Path("/basemodel/projects/project_dir").resolve()
# EDIT: define name for scenario checkpoints directory; the script will put it under the same parent directory as fm
scenario_name = "scenario_name"

# creating the relative paths
foundation_model_path = project_dir / "fm"
scenario_model_path = project_dir / "scenarios" / scenario_name


# --- Target Function ---------------------------------------------------------

# churn definition, for reference:
# 1: churned (no events in the future window)
# 0: not churned (at least one event in the future window)

# EDIT: provide target details
TARGET_EVENT_TABLE = "transactions" # the event data source to base the target logic (here: purchases of products)
TARGET_WINDOW_DAYS = 21 # the length of time to look at into the future (other time units are possible)


def target_fn(history: Events, future: Events, _entity: Attributes, _ctx: Dict) -> np.ndarray | None:

    # filters out users with too short remaining window
    if has_incomplete_training_window(_ctx, timedelta(days=TARGET_WINDOW_DAYS)):
        return None

    # filters out users with no history
    if history[TARGET_EVENT_TABLE].count() == 0:
        return None

    # trims the future to desired window
    future_window = future.interval_from(
        _ctx[SPLIT_TIMESTAMP],
        timedelta(days=TARGET_WINDOW_DAYS),
    )

    # churn label
    y = 0 if future_window[TARGET_EVENT_TABLE].count() > 0 else 1
    return np.array([y], dtype=np.float32)


# --- Training ----------------------------------------------------------------

# EDIT: metaparams - keep default unless experimenting
learning_rate = 0.0001
epochs = 3 # use 1 for smoke test

# EDIT: limited runs - use for smoke test, then comment out here and below
limit_train_batches = 5 
limit_val_batches = 5

# EDIT: parallelised training - comment out to default to a single GPU
strategy = "ddp"
devices = [0, 1] # list GPU indices

# For more options refer to docs: https://docs.basemodel.ai/reference/trainingparams

task = BinaryClassificationTask()

training_params = TrainingParams(
    checkpoint_dir=scenario_model_path,
    learning_rate=learning_rate,
    epochs=epochs,
    devices=devices,
    strategy=strategy,
    limit_train_batches=limit_train_batches, # smoke test, comment out for full runs
    limit_val_batches=limit_val_batches, # smoke test, comment out for full runs
)

trainer = load_from_foundation_model(
    checkpoint_path=foundation_model_path,
    downstream_task=task,
    target_fn=target_fn,
)

trainer.fit(training_params=training_params, overwrite=True) # replace the last param with resume=True for resumed training

from datetime import datetime
from pathlib import Path

from monad.ui.config import MetricParams, OutputType, TestingParams
from monad.ui.module import load_from_checkpoint


# --- Names & Paths -----------------------------------------------------------

# EDIT: provide path to scenario checkpoints directory saved during training
scenario_checkpoint_dir = Path("/basemodel/projects/project_dir/scenarios/scenario_name").resolve()

# EDIT: if desired, modify where test outputs will be saved
predictions_path = scenario_checkpoint_dir / "test" / "predictions_and_ground_truth.tsv"


# --- Test Set Definition -----------------------------------------------------

# EDIT: prediction_date defines the start of the temporally defined test set
prediction_date = datetime(2020, 9, 5)

# EDIT: limited runs - use for smoke test, else comment out
limit_test_batches = 100


# --- Metrics -----------------------------------------------------------------

# EDIT: define metrics to compute during testing
metrics = [
    MetricParams(alias="auroc", metric_name="AUROC", kwargs={"task": "binary"}),
    MetricParams(alias="recall", metric_name="Recall", kwargs={"task": "binary"}),
    MetricParams(alias="precision", metric_name="Precision", kwargs={"task": "binary"}),
]

# For more options refer to docs.
# Testing parameters: https://docs.basemodel.ai/reference/testingparams
# Testing metrics: https://docs.basemodel.ai/docs/customizing-testing-metrics


# --- Test --------------------------------------------------------------------

testing_module = load_from_checkpoint(checkpoint_path=scenario_checkpoint_dir)

testing_params = TestingParams(
    output_type=OutputType.DECODED,
    local_save_location=predictions_path,
    prediction_date=prediction_date,
    metrics=metrics,
    limit_test_batches=limit_test_batches,  # smoke test, comment out for full runs
)

testing_module.test(testing_params=testing_params)

from datetime import datetime
from pathlib import Path

from monad.ui.config import OutputType, TestingParams
from monad.ui.module import load_from_checkpoint


# --- Names & Paths -----------------------------------------------------------

# EDIT: provide path to scenario checkpoints directory saved during training
scenario_checkpoint_dir = Path("/basemodel/projects/project_dir/scenarios/scenario_name").resolve()

# EDIT: define where prediction outputs will be saved
predictions_path = scenario_checkpoint_dir / "predictions" / "predictions.tsv"


# --- Prediction Set Definition -----------------------------------------------

# EDIT: prediction_date defines the split point used for inference
prediction_date = datetime(2020, 9, 23)


# --- Predict -----------------------------------------------------------------

prediction_module = load_from_checkpoint(checkpoint_path=scenario_checkpoint_dir)

testing_params = TestingParams(
    output_type=OutputType.DECODED,
    local_save_location=predictions_path,
    prediction_date=prediction_date,
)

prediction_module.predict(testing_params=testing_params)