Skip to navigation
User Code GuideCode ExamplesLinux Deployment

Train a Classification Model on Linux

View as Markdown

Train a model to identify transactions labeled as money laundering in IBM’s synthetic AML dataset. This example prepares account and transaction tables, connects them in a graph, trains a binary classifier, and saves the model for later predictions. It uses Kumo Engine’s local Python API on a Linux GPU machine.

Before You Begin

Complete Linux Machine Deployment Setup and activate the installed Python environment. Keep KUMO_SCRATCH set to a writable directory on local NVMe SSD storage. Allow space for the downloaded data, a prepared copy, graph indexes, training files, and saved model. See GPU Sizing for compute guidance.

Download and extract the IBM Transactions for Anti Money Laundering dataset. Use these two files:

FileContents
LI-Medium_accounts.csvAccount attributes for transaction senders and recipients
LI-Medium_Trans.csvTimestamped transactions and the Is Laundering label

Review IBM’s data description and license to determine whether the dataset is suitable for your use. The archive also contains LI-Medium_Patterns.txt, which describes the generated laundering attempts. Use it as a reference only: its outcome annotations would reveal the answers if added as model features.

Point IBM_AML_SOURCE to the extracted files. Choose new directories for the prepared tables and saved model, separate from the source files and scratch:

export IBM_AML_SOURCE=/path/to/archive
export IBM_AML_PREPARED=/path/to/local-storage/aml-prepared-01
export KUMO_MODEL_OUTPUT=/path/to/retained-models/aml-run-01

Run the Python blocks below in order in the same process. Configure scratch for that process before building the graph:

import csv
import json
import os
from datetime import datetime
from pathlib import Path
from kumo.engine import configure
source_dir = Path(os.environ["IBM_AML_SOURCE"]).expanduser().resolve()
prepared_dir = Path(os.environ["IBM_AML_PREPARED"]).expanduser().resolve()
output_dir = Path(os.environ["KUMO_MODEL_OUTPUT"]).expanduser().resolve()
scratch_dir = Path(os.environ["KUMO_SCRATCH"]).expanduser().resolve()
configure(scratch_dir=scratch_dir)
for name in ("LI-Medium_accounts.csv", "LI-Medium_Trans.csv"):
    if not (source_dir / name).is_file():
        raise FileNotFoundError(source_dir / name)

Prepare the Input Tables

Create an account key from the bank ID and account number so each transaction can connect to its sender and recipient. The source files pad bank IDs differently, so normalize them before creating the keys. The transaction file also has two columns named Account; read those columns by position to distinguish the two roles.

The following code writes accounts.csv and transactions.csv without changing the downloaded files. It streams transactions and keeps account keys in memory to check duplicates and missing accounts. Each transaction receives an ID based on its position in the source file.

prepared_dir.mkdir(parents=True, exist_ok=False)
def account_key(bank_id, account_number):
    return f"{int(bank_id)}:{account_number.strip()}"
account_ids = set()
accounts_tmp = prepared_dir / "accounts.csv.tmp"
with (source_dir / "LI-Medium_accounts.csv").open(
    encoding="utf-8-sig", newline=""
) as source, accounts_tmp.open("w", encoding="utf-8", newline="") as target:
    reader = csv.DictReader(source)
    expected = ["Bank Name", "Bank ID", "Account Number", "Entity ID", "Entity Name"]
    if reader.fieldnames != expected:
        raise ValueError(f"Unexpected accounts header: {reader.fieldnames}")
    writer = csv.writer(target)
    writer.writerow(["AccountID", "BankName", "EntityType"])
    for row in reader:
        key = account_key(row["Bank ID"], row["Account Number"])
        if key in account_ids:
            raise ValueError(f"Duplicate normalized account key: {key}")
        account_ids.add(key)
        entity_type = row["Entity Name"].split("#", 1)[0].strip()
        writer.writerow([key, row["Bank Name"], entity_type])
accounts_tmp.replace(prepared_dir / "accounts.csv")
transaction_header = [
    "Timestamp", "From Bank", "Account", "To Bank", "Account",
    "Amount Received", "Receiving Currency", "Amount Paid",
    "Payment Currency", "Payment Format", "Is Laundering",
]
transactions_tmp = prepared_dir / "transactions.csv.tmp"
positive_count = missing_account_endpoints = 0
with (source_dir / "LI-Medium_Trans.csv").open(
    encoding="utf-8-sig", newline=""
) as source, transactions_tmp.open(
    "w", encoding="utf-8", newline=""
) as target:
    reader = csv.reader(source)
    if next(reader) != transaction_header:
        raise ValueError("Unexpected transaction header or column order")
    writer = csv.writer(target)
    writer.writerow([
        "TransactionID", "Timestamp", "SenderAccountID",
        "RecipientAccountID", "AmountReceived", "ReceivingCurrency",
        "AmountPaid", "PaymentCurrency", "PaymentFormat", "IsLaundering",
    ])
    for transaction_id, row in enumerate(reader, start=1):
        if len(row) != len(transaction_header) or row[10] not in {"0", "1"}:
            raise ValueError(f"Invalid transaction row {transaction_id}")
        sender = account_key(row[1], row[2])
        recipient = account_key(row[3], row[4])
        missing_account_endpoints += (sender not in account_ids)
        missing_account_endpoints += (recipient not in account_ids)
        positive_count += row[10] == "1"
        timestamp = datetime.strptime(
            row[0], "%Y/%m/%d %H:%M"
        ).strftime("%Y-%m-%d %H:%M:%S")
        writer.writerow([
            transaction_id, timestamp, sender, recipient,
            row[5], row[6], row[7], row[8], row[9], row[10],
        ])
transactions_tmp.replace(prepared_dir / "transactions.csv")
print("Accounts:", len(account_ids))
print("Transactions:", transaction_id)
print("Laundering-labeled transactions:", positive_count)
print("Account endpoints absent from accounts file:", missing_account_endpoints)

The printed counts summarize the prepared data and any transaction endpoints missing from the account table. A missing endpoint has no account row to connect to; check the source files and key formatting if the count is unexpected. If preparation fails, correct the reported input problem and use a new prepared directory.

Build the Graph

Add the account table twice, as senders and recipients, to represent the two roles in a transaction. Connect each transaction to those accounts and use its timestamp as the event time.

from kumo.engine import FileConnector, Graph, PredictiveQuery, Table, Trainer
from kumo.engine.trainer.output_config import ModelOutputConfig
from kumoapi.graph import Edge
from kumoapi.model_plan import RunMode, TrainingTableGenerationPlan
from kumoapi.typing import Dtype, Stype
connector = FileConnector(root_dir=str(prepared_dir))
accounts = Table.from_source_table(
    connector["accounts"],
    ["AccountID", "BankName", "EntityType"],
    primary_key="AccountID",
)
transactions = Table.from_source_table(
    connector["transactions"],
    [
        "TransactionID", "Timestamp", "SenderAccountID",
        "RecipientAccountID", "AmountReceived", "ReceivingCurrency",
        "AmountPaid", "PaymentCurrency", "PaymentFormat", "IsLaundering",
    ],
    primary_key="TransactionID",
    time_column="Timestamp",
)
transactions = transactions.with_column(
    "Timestamp", dtype=Dtype.date,
    timestamp_format="%Y-%m-%d %H:%M:%S",
)
for column in ("SenderAccountID", "RecipientAccountID"):
    transactions = transactions.with_column(column, stype=Stype.ID)
for column in ("ReceivingCurrency", "PaymentCurrency", "PaymentFormat"):
    transactions = transactions.with_column(column, stype=Stype.categorical)
for column in ("AmountReceived", "AmountPaid"):
    transactions = transactions.with_column(column, stype=Stype.numerical)
graph = Graph(
    tables={
        "senders": accounts,
        "recipients": accounts,
        "transactions": transactions,
    },
    edges=[
        Edge(src_table="transactions", fkey="SenderAccountID", dst_table="senders"),
        Edge(src_table="transactions", fkey="RecipientAccountID", dst_table="recipients"),
    ],
)
graph.validate()
snapshot = graph.snapshot()
print("Snapshot:", snapshot.id)

snapshot.id identifies the materialized graph used for training. Kumo stores its graph indexes in the configured scratch directory.

Define the Classification Task

The predictive query asks whether each transaction’s IsLaundering value equals 1. The comparison defines a binary classification target. Kumo treats that target as the value to predict, rather than an input feature for the transaction being scored.

Use a chronological split with training, validation, and test ratios of 0.8, 0.1, and 0.1. Validation and test examples follow the training period:

query = PredictiveQuery(
    graph,
    "PREDICT transactions.IsLaundering = 1 "
    "FOR EACH transactions.TransactionID",
)
query.validate()
print("Task type:", query.get_task_type())
generation_plan = TrainingTableGenerationPlan(
    split="TemporalSplit([0.8, 0.1, 0.1])",
)
train_table = query.generate_training_table(
    snapshot, plan=generation_plan
)
print("Training rows:", train_table.train_table_path)

Preview ten generated rows without loading the full training table into memory:

import pyarrow.parquet as pq
parquet = pq.ParquetFile(train_table.train_table_path)
print("Generated rows:", parquet.metadata.num_rows)
first_batch = next(parquet.iter_batches(batch_size=10))
print(first_batch.to_pandas())

ENTITY1 identifies the transaction, TIMESTAMP is its event time, and TARGET1 is the binary label. For more on training populations and splits, see Predictive Queries.

Train and Save the Model

Start with one experiment and up to five epochs. Use area under the precision-recall curve (AUPRC) to select the model because laundering labels are rare. The plan also records AUROC, precision, and recall. See Training to adjust the training budget.

The following code saves the plan as model-plan.json and publishes the trained model under model/ in your output directory:

output_dir.mkdir(parents=True, exist_ok=False)
plan = query.suggest_model_plan(snapshot, run_mode=RunMode.FAST)
plan.training_job.num_experiments = 1
plan.training_job.refit_trainval = False
plan.training_job.refit_full = False
plan.training_job.metrics = ["auprc", "auroc", "precision", "recall"]
plan.training_job.tune_metric = "auprc"
plan.optimization.max_epochs = 5
(output_dir / "model-plan.json").write_text(
    json.dumps(plan.as_json(), indent=2) + "\n"
)
trainer = Trainer(plan)
trained = trainer.fit(
    train_table,
    snapshot=snapshot,
    export_config=ModelOutputConfig(
        connector=FileConnector(root_dir=str(output_dir)),
        path="model",
    ),
)
print("Metrics:", json.dumps(trained.metrics, indent=2))
print("Published model bundle:", output_dir / "model")

Review the Results

Use the validation metrics printed by the fit to adjust the training budget and prediction threshold. Reserve the test split for the final evaluation after choosing those settings. For this imbalanced task, use AUPRC, precision, and recall to assess how well the model finds positive examples; accuracy alone can hide missed laundering cases. Choose a prediction threshold based on the tradeoff between missed events and false alerts. The data is synthetic, so evaluate the model on representative data before using it for real transactions.

Keep the complete model/ directory on storage that survives the machine’s lifetime. The published bundle can be loaded in a later Python process; detailed trial histories and explanation files remain in scratch. Use Training Diagnostics to inspect and retain those files.

Next Steps