"""Interactive inference demo for the LFM2 Transaction Foundation Model.
Demonstrates multi-head predictions (fraud, next merchant, amount range, MCC)
on synthetic payment sequences. Includes side-by-side pretrained-vs-random-init
comparison showing the value of self-supervised pretraining.
Usage:
python -m src.demo.app --checkpoint PATH [--data-dir PATH] [--port PORT]
Runs on CPU. No GPU or internet required. Inference < 100ms per customer.
"""
from __future__ import annotations
import argparse
import time
from pathlib import Path
from typing import Any
import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from src.data.schema import SchemaConfig, load_schema
from src.demo.decode import TransactionDecoder
from src.demo.merchant_catalog import DemoMerchantCatalog
from src.demo.profile_inference import format_profile_html, infer_profile
from src.demo.render import (
format_amount_predictions,
format_fraud_score,
format_mcc_predictions,
format_merchant_predictions,
format_timeline,
render_comparison_header,
render_integration_guide,
render_production_architecture,
render_why_liquid,
)
from src.model.lfm2_small import LFM2Small, ModelConfig
from src.model.task_heads import (
AnyHead,
DownstreamHead,
HeadConfig,
MultiHeadModel,
TiedEmbeddingHead,
)
# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------
class DemoData:
"""Loads test-set sequences and labels into memory."""
def __init__(self, data_dir: Path, schema: SchemaConfig) -> None:
token_ids = np.load(data_dir / "token_ids.npy", mmap_mode="r")
seq_labels = np.load(data_dir / "sequence_labels.npy")
splits = np.load(data_dir / "split_indices.npz")
test_idx = splits["test"]
self.token_ids = np.array(token_ids[test_idx])
self.labels = seq_labels[test_idx].astype(int)
self.num_customers = len(test_idx)
self.fraud_indices = np.where(self.labels == 1)[0]
self.legit_indices = np.where(self.labels == 0)[0]
self._curated = self._find_curated_examples()
def _find_curated_examples(self) -> dict[str, int]:
"""Pick interesting examples for quick navigation.
Five legitimate profiles to show breadth of normal spending patterns,
three fraud profiles representing the main attack archetypes.
"""
examples: dict[str, int] = {}
# Legitimate customers — varied indices for behavioral variety
if len(self.legit_indices) > 0:
examples["Typical Customer"] = int(self.legit_indices[0])
if len(self.legit_indices) > 100:
examples["Frequent Shopper"] = int(self.legit_indices[100])
if len(self.legit_indices) > 200:
examples["Weekend Spender"] = int(self.legit_indices[200])
if len(self.legit_indices) > 300:
examples["High-Spend Loyalist"] = int(self.legit_indices[300])
if len(self.legit_indices) > 500:
examples["Low Activity"] = int(self.legit_indices[500])
# Fraud archetypes
if len(self.fraud_indices) > 0:
examples["Fraud: Card Testing"] = int(self.fraud_indices[0])
if len(self.fraud_indices) > 50:
examples["Fraud: Account Takeover"] = int(self.fraud_indices[50])
if len(self.fraud_indices) > 100:
examples["Fraud: High Value"] = int(self.fraud_indices[100])
return examples
@property
def curated_names(self) -> list[str]:
return list(self._curated.keys())
def get_curated_index(self, name: str) -> int:
return self._curated[name]
# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
def build_model(
model_yaml: Path,
schema: SchemaConfig,
finetune_yaml: Path,
) -> MultiHeadModel:
"""Construct MultiHeadModel with 4 downstream heads."""
import yaml
backbone = LFM2Small(ModelConfig.from_yaml(model_yaml), schema)
with open(finetune_yaml) as f:
ft_config = yaml.safe_load(f)
heads: dict[str, AnyHead] = {}
for name, hcfg in ft_config["heads"].items():
config = HeadConfig(
name=name,
output_dim=hcfg["output_dim"],
loss_type=hcfg["loss"],
pool_strategy=hcfg["pool"],
target_type=hcfg["target"],
weight=hcfg.get("weight", 1.0),
mlp_hidden=hcfg.get("mlp_hidden", 128),
dropout=hcfg.get("dropout", 0.1),
)
if hcfg.get("tied", False):
heads[name] = TiedEmbeddingHead(
config, backbone.config.hidden_size, schema.num_features,
backbone.embedding.value_tables,
)
else:
heads[name] = DownstreamHead(
config, backbone.config.hidden_size, schema.num_features,
)
return MultiHeadModel(backbone, heads)
def load_model_checkpoint(model: MultiHeadModel, checkpoint_path: Path | None) -> str:
"""Load fine-tuned weights. Returns status message."""
if checkpoint_path is None:
return "No checkpoint loaded"
if not checkpoint_path.exists():
raise FileNotFoundError(
f"Checkpoint not found: {checkpoint_path}. "
f"Run with --checkpoint PATH or --checkpoint none to skip."
)
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state_dict"], strict=True)
step = ckpt.get("step", "?")
return f"step {step}"
# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
@torch.no_grad()
def run_inference(
model: MultiHeadModel,
token_ids: np.ndarray,
) -> dict[str, np.ndarray]:
"""Run all 4 heads on a single customer sequence."""
tensor = torch.from_numpy(token_ids).unsqueeze(0).long()
predictions = model(tensor)
results: dict[str, np.ndarray] = {}
for name, logits in predictions.items():
if name == "fraud":
prob = torch.sigmoid(logits).squeeze().numpy()
results[name] = np.array([float(prob)])
else:
probs = F.softmax(logits, dim=-1).squeeze(0).numpy()
results[name] = probs
return results
# ---------------------------------------------------------------------------
# App builder
# ---------------------------------------------------------------------------
def create_app(
pretrained_model: MultiHeadModel,
random_model: MultiHeadModel,
data: DemoData,
decoder: TransactionDecoder,
merchant_catalog: DemoMerchantCatalog,
checkpoint_status: str,
) -> gr.Blocks:
"""Build the Gradio app with pretrained vs random-init comparison."""
pretrained_model.eval()
random_model.eval()
def on_customer_select(
curated_name: str | None,
customer_idx: int,
mode: str,
) -> tuple[str, str, str, str, str, str, str, str, str, str, str]:
"""Run both models, return all outputs for comparison."""
if mode == "Curated examples" and curated_name:
idx = data.get_curated_index(curated_name)
else:
idx = int(customer_idx)
idx = max(0, min(idx, data.num_customers - 1))
token_ids = data.token_ids[idx]
is_fraud = bool(data.labels[idx])
summary = decoder.summarize_customer(token_ids, is_fraud)
t0 = time.perf_counter()
pre_results = run_inference(pretrained_model, token_ids)
rand_results = run_inference(random_model, token_ids)
latency_ms = (time.perf_counter() - t0) * 1000
timeline_html = format_timeline(decoder, token_ids)
# Pretrained predictions
pre_fraud = format_fraud_score(float(pre_results["fraud"][0]))
pre_merchant = format_merchant_predictions(
pre_results["next_merchant"], merchant_catalog, k=5,
)
pre_amount = format_amount_predictions(pre_results["amount_range"], k=5)
pre_mcc = format_mcc_predictions(pre_results["mcc"], k=5)
# Random-init predictions
rand_fraud = format_fraud_score(float(rand_results["fraud"][0]))
rand_merchant = format_merchant_predictions(
rand_results["next_merchant"], merchant_catalog, k=5,
)
rand_amount = format_amount_predictions(rand_results["amount_range"], k=5)
rand_mcc = format_mcc_predictions(rand_results["mcc"], k=5)
# Behavioral profile
profile_match = infer_profile(token_ids)
profile_html = format_profile_html(profile_match)
# Combined comparison HTML for each head
fraud_compare = _side_by_side("Fraud Score", pre_fraud, rand_fraud)
merchant_compare = _side_by_side("Next Merchant", pre_merchant, rand_merchant)
amount_compare = _side_by_side("Amount Range", pre_amount, rand_amount)
mcc_compare = _side_by_side("Merchant Category", pre_mcc, rand_mcc)
latency_str = (
f"Inference: {latency_ms:.1f}ms (both models) on CPU | "
f"Customer #{idx} | Ground truth: {'FRAUD' if is_fraud else 'Legitimate'}"
)
return (
summary, timeline_html, profile_html,
fraud_compare, merchant_compare, amount_compare, mcc_compare,
latency_str,
)
with gr.Blocks(
title="LFM2 Transaction Foundation Model",
) as app:
gr.HTML("""
LFM2 Transaction Foundation Model
Liquid AI · LFM2.5 Architecture · Multi-Head Inference Demo
""")
with gr.Tabs():
# ===== Tab 1: Interactive Predictions =====
with gr.Tab("Predictions"):
gr.HTML("""
How to read this: Select a customer, see their
transaction history, then compare predictions from the pretrained model (left)
vs random initialization (right). Same architecture, same fine-tuning data.
The only difference is self-supervised pretraining on unlabeled sequences.
""")
with gr.Accordion("Reference model details", open=False):
gr.HTML(f"""
""")
# Customer Selection. Use Radio (not Checkbox) for the
# curated-vs-browse toggle: the two modes are mutually exclusive
# and Radio's visual state is more reliable across Gradio versions.
_CURATED_MODE = "Curated examples"
_BROWSE_MODE = "Browse all customers"
with gr.Row():
with gr.Column(scale=1):
gr.HTML("