End-to-end example of a scenario training script

from typing import Dict
from monad.ui.config import TrainingParams
from monad.ui.module import load_from_foundation_model, BinaryClassificationTask
from monad.ui.target_function import Attributes, Events

import torch
import os

# specifying the target
def target_fn(history: Events, future: Events, _entity: Attributes, _ctx: Dict) -> np.ndarray:
    if history["product.buy"].count() == 0:
        return None

    # churn definition:
    # 1 - churned
    # 0 - not churned
    
    churn = 0 if future["product.buy"].count() > 0 else 1

    # output should be a float32 numpy array
    return np.array([churn], dtype=np.float32) 

# selecting the foundation model
fm_path = "/path/to/your/models/pretrain/fm" # parameter for: load_from_foundation_model
checkpoint_dir = "path/to/your/models/downstream/model_name" # parameter for: training_params

# adapting training parameters
training_params = TrainingParams(
    checkpoint_dir=checkpoint_dir, # location to save your scenario model
    epochs=1,
    learning_rate=0.001,
    devices=[0]
)

# instantiating the trainer & training the model
if __name__ == "__main__":
    trainer = load_from_foundation_model(
        checkpoint_path=fm_path, # location of foundation model
        downstream_task=BinaryClassificationTask(), # task aligned with business scenario
        target_fn=target_fn, 
        num_outputs=1 # task-specific number of outputs
    )
    trainer.fit(training_params=training_params)