441 lines
20 KiB
Python
441 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""Sample from a trained hybrid diffusion model with routing-aware export fixes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import gzip
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from data_utils import inverse_quantile_transform, load_split, normalize_cont, quantile_calibrate_to_real
|
|
from export_samples import build_inverse_vocab, load_stats, load_torch_state, read_header
|
|
from hybrid_diffusion import (
|
|
HybridDiffusionModel,
|
|
TemporalGRUGenerator,
|
|
TemporalTransformerGenerator,
|
|
cosine_beta_schedule,
|
|
)
|
|
from platform_utils import resolve_device, resolve_path
|
|
from submission_type_utils import (
|
|
denormalize_cont_tensor,
|
|
resolve_routing_features,
|
|
resolve_taxonomy_features,
|
|
)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Sample and export HAI feature sequences with routing-aware fixes.")
|
|
base_dir = Path(__file__).resolve().parent
|
|
repo_dir = base_dir.parent.parent
|
|
parser.add_argument("--data-path", default=str(repo_dir / "dataset" / "hai" / "hai-21.03" / "train1.csv.gz"))
|
|
parser.add_argument("--data-glob", default=str(repo_dir / "dataset" / "hai" / "hai-21.03" / "train*.csv.gz"))
|
|
parser.add_argument("--split-path", default=str(base_dir / "feature_split.json"))
|
|
parser.add_argument("--stats-path", default=str(base_dir / "results" / "cont_stats.json"))
|
|
parser.add_argument("--vocab-path", default=str(base_dir / "results" / "disc_vocab.json"))
|
|
parser.add_argument("--model-path", default=str(base_dir / "results" / "model.pt"))
|
|
parser.add_argument("--out", default=str(base_dir / "results" / "generated.csv"))
|
|
parser.add_argument("--timesteps", type=int, default=200)
|
|
parser.add_argument("--seq-len", type=int, default=64)
|
|
parser.add_argument("--batch-size", type=int, default=2)
|
|
parser.add_argument("--device", default="auto", help="cpu, cuda, or auto")
|
|
parser.add_argument("--include-time", action="store_true", help="Include time column as a simple index")
|
|
parser.add_argument("--clip-k", type=float, default=5.0, help="Clip continuous values to mean±k*std")
|
|
parser.add_argument("--use-ema", action="store_true", help="Use EMA weights if available")
|
|
parser.add_argument("--config", default=None, help="Optional config_used.json to infer conditioning")
|
|
parser.add_argument("--condition-id", type=int, default=-1, help="Condition file id (0..N-1), -1=random")
|
|
parser.add_argument("--include-condition", action="store_true", help="Include condition id column in CSV")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
base_dir = Path(__file__).resolve().parent
|
|
args.data_path = str(resolve_path(base_dir, args.data_path))
|
|
args.data_glob = str(resolve_path(base_dir, args.data_glob)) if args.data_glob else ""
|
|
args.split_path = str(resolve_path(base_dir, args.split_path))
|
|
args.stats_path = str(resolve_path(base_dir, args.stats_path))
|
|
args.vocab_path = str(resolve_path(base_dir, args.vocab_path))
|
|
args.model_path = str(resolve_path(base_dir, args.model_path))
|
|
args.out = str(resolve_path(base_dir, args.out))
|
|
|
|
if not os.path.exists(args.model_path):
|
|
raise SystemExit(f"missing model file: {args.model_path}")
|
|
|
|
data_path = args.data_path
|
|
if args.data_glob:
|
|
base = Path(args.data_glob).parent
|
|
pat = Path(args.data_glob).name
|
|
matches = sorted(base.glob(pat))
|
|
if matches:
|
|
data_path = str(matches[0])
|
|
|
|
split = load_split(args.split_path)
|
|
time_col = split.get("time_column", "time")
|
|
cont_cols = [c for c in split["continuous"] if c != time_col]
|
|
disc_cols = [c for c in split["discrete"] if not c.startswith("attack") and c != time_col]
|
|
|
|
stats = load_stats(args.stats_path)
|
|
mean = stats["mean"]
|
|
std = stats["std"]
|
|
vmin = stats.get("min", {})
|
|
vmax = stats.get("max", {})
|
|
int_like = stats.get("int_like", {})
|
|
max_decimals = stats.get("max_decimals", {})
|
|
transforms = stats.get("transform", {})
|
|
quantile_probs = stats.get("quantile_probs")
|
|
quantile_values = stats.get("quantile_values")
|
|
quantile_raw_values = stats.get("quantile_raw_values")
|
|
|
|
vocab_json = json.load(open(args.vocab_path, "r", encoding="utf-8"))
|
|
vocab = vocab_json["vocab"]
|
|
top_token = vocab_json.get("top_token", {})
|
|
inv_vocab = build_inverse_vocab(vocab)
|
|
vocab_sizes = [len(vocab[c]) for c in disc_cols]
|
|
|
|
device = resolve_device(args.device)
|
|
cfg = {}
|
|
use_condition = False
|
|
cond_vocab_size = 0
|
|
if args.config:
|
|
args.config = str(resolve_path(base_dir, args.config))
|
|
if args.config and os.path.exists(args.config):
|
|
with open(args.config, "r", encoding="utf-8") as f:
|
|
cfg = json.load(f)
|
|
use_condition = bool(cfg.get("use_condition")) and cfg.get("condition_type") == "file_id"
|
|
if use_condition:
|
|
cfg_base = Path(args.config).resolve().parent
|
|
cfg_glob = cfg.get("data_glob", args.data_glob)
|
|
cfg_glob = str(resolve_path(cfg_base, cfg_glob))
|
|
base = Path(cfg_glob).parent
|
|
pat = Path(cfg_glob).name
|
|
cond_vocab_size = len(sorted(base.glob(pat)))
|
|
if cond_vocab_size <= 0:
|
|
raise SystemExit(f"use_condition enabled but no files matched data_glob: {cfg_glob}")
|
|
|
|
cont_target = str(cfg.get("cont_target", "eps"))
|
|
cont_clamp_x0 = float(cfg.get("cont_clamp_x0", 0.0))
|
|
use_quantile = bool(cfg.get("use_quantile_transform", False))
|
|
cont_bound_mode = str(cfg.get("cont_bound_mode", "clamp"))
|
|
cont_bound_strength = float(cfg.get("cont_bound_strength", 1.0))
|
|
cont_post_scale = cfg.get("cont_post_scale", {}) if isinstance(cfg.get("cont_post_scale", {}), dict) else {}
|
|
cont_post_calibrate = bool(cfg.get("cont_post_calibrate", False))
|
|
|
|
route_type1_cols = resolve_routing_features(cfg, cont_cols, "type1_features")
|
|
route_type5_cols = resolve_routing_features(cfg, cont_cols, "type5_features")
|
|
type4_cols = resolve_taxonomy_features(cfg, cont_cols, "type4_features")
|
|
model_cont_cols = [c for c in cont_cols if c not in route_type1_cols and c not in route_type5_cols]
|
|
|
|
use_temporal_stage1 = bool(cfg.get("use_temporal_stage1", False))
|
|
temporal_use_type1_cond = bool(cfg.get("temporal_use_type1_cond", False))
|
|
temporal_focus_type4 = bool(cfg.get("temporal_focus_type4", False))
|
|
temporal_exclude_type4 = bool(cfg.get("temporal_exclude_type4", False))
|
|
temporal_backbone = str(cfg.get("temporal_backbone", "gru"))
|
|
temporal_hidden_dim = int(cfg.get("temporal_hidden_dim", 256))
|
|
temporal_num_layers = int(cfg.get("temporal_num_layers", 1))
|
|
temporal_dropout = float(cfg.get("temporal_dropout", 0.0))
|
|
temporal_pos_dim = int(cfg.get("temporal_pos_dim", 64))
|
|
temporal_use_pos_embed = bool(cfg.get("temporal_use_pos_embed", True))
|
|
temporal_transformer_num_layers = int(cfg.get("temporal_transformer_num_layers", 2))
|
|
temporal_transformer_nhead = int(cfg.get("temporal_transformer_nhead", 4))
|
|
temporal_transformer_ff_dim = int(cfg.get("temporal_transformer_ff_dim", 512))
|
|
temporal_transformer_dropout = float(cfg.get("temporal_transformer_dropout", 0.1))
|
|
backbone_type = str(cfg.get("backbone_type", "gru"))
|
|
transformer_num_layers = int(cfg.get("transformer_num_layers", 2))
|
|
transformer_nhead = int(cfg.get("transformer_nhead", 4))
|
|
transformer_ff_dim = int(cfg.get("transformer_ff_dim", 512))
|
|
transformer_dropout = float(cfg.get("transformer_dropout", 0.1))
|
|
|
|
model = HybridDiffusionModel(
|
|
cont_dim=len(model_cont_cols),
|
|
disc_vocab_sizes=vocab_sizes,
|
|
time_dim=int(cfg.get("model_time_dim", 64)),
|
|
hidden_dim=int(cfg.get("model_hidden_dim", 256)),
|
|
num_layers=int(cfg.get("model_num_layers", 1)),
|
|
dropout=float(cfg.get("model_dropout", 0.0)),
|
|
ff_mult=int(cfg.get("model_ff_mult", 2)),
|
|
pos_dim=int(cfg.get("model_pos_dim", 64)),
|
|
use_pos_embed=bool(cfg.get("model_use_pos_embed", True)),
|
|
backbone_type=backbone_type,
|
|
transformer_num_layers=transformer_num_layers,
|
|
transformer_nhead=transformer_nhead,
|
|
transformer_ff_dim=transformer_ff_dim,
|
|
transformer_dropout=transformer_dropout,
|
|
cond_cont_dim=len(route_type1_cols),
|
|
cond_vocab_size=cond_vocab_size if use_condition else 0,
|
|
cond_dim=int(cfg.get("cond_dim", 32)),
|
|
use_tanh_eps=bool(cfg.get("use_tanh_eps", False)),
|
|
eps_scale=float(cfg.get("eps_scale", 1.0)),
|
|
).to(device)
|
|
if args.use_ema and os.path.exists(args.model_path.replace("model.pt", "model_ema.pt")):
|
|
ema_path = args.model_path.replace("model.pt", "model_ema.pt")
|
|
model.load_state_dict(load_torch_state(ema_path, device))
|
|
else:
|
|
model.load_state_dict(load_torch_state(args.model_path, device))
|
|
model.eval()
|
|
|
|
temporal_model = None
|
|
if use_temporal_stage1:
|
|
temporal_path = Path(args.model_path).with_name("temporal.pt")
|
|
if not temporal_path.exists():
|
|
raise SystemExit(f"missing temporal model file: {temporal_path}")
|
|
temporal_state = load_torch_state(str(temporal_path), device)
|
|
temporal_cond_dim = len(route_type1_cols) if (temporal_use_type1_cond and route_type1_cols) else 0
|
|
if isinstance(temporal_state, dict):
|
|
if "in_proj.weight" in temporal_state:
|
|
try:
|
|
temporal_cond_dim = max(0, int(temporal_state["in_proj.weight"].shape[1]) - len(model_cont_cols))
|
|
except Exception:
|
|
pass
|
|
elif "gru.weight_ih_l0" in temporal_state:
|
|
try:
|
|
temporal_cond_dim = max(0, int(temporal_state["gru.weight_ih_l0"].shape[1]) - len(model_cont_cols))
|
|
except Exception:
|
|
pass
|
|
if temporal_backbone == "transformer":
|
|
temporal_model = TemporalTransformerGenerator(
|
|
input_dim=len(model_cont_cols),
|
|
hidden_dim=temporal_hidden_dim,
|
|
num_layers=temporal_transformer_num_layers,
|
|
nhead=temporal_transformer_nhead,
|
|
ff_dim=temporal_transformer_ff_dim,
|
|
dropout=temporal_transformer_dropout,
|
|
pos_dim=temporal_pos_dim,
|
|
use_pos_embed=temporal_use_pos_embed,
|
|
cond_dim=temporal_cond_dim,
|
|
).to(device)
|
|
else:
|
|
temporal_model = TemporalGRUGenerator(
|
|
input_dim=len(model_cont_cols),
|
|
hidden_dim=temporal_hidden_dim,
|
|
num_layers=temporal_num_layers,
|
|
dropout=temporal_dropout,
|
|
cond_dim=temporal_cond_dim,
|
|
).to(device)
|
|
temporal_model.load_state_dict(temporal_state)
|
|
temporal_model.eval()
|
|
|
|
betas = cosine_beta_schedule(args.timesteps).to(device)
|
|
alphas = 1.0 - betas
|
|
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
|
|
|
x_cont = torch.randn(args.batch_size, args.seq_len, len(model_cont_cols), device=device)
|
|
x_disc = torch.full((args.batch_size, args.seq_len, len(disc_cols)), 0, device=device, dtype=torch.long)
|
|
mask_tokens = torch.tensor(vocab_sizes, device=device)
|
|
for i in range(len(disc_cols)):
|
|
x_disc[:, :, i] = mask_tokens[i]
|
|
|
|
cond = None
|
|
if use_condition:
|
|
if cond_vocab_size <= 0:
|
|
raise SystemExit("use_condition enabled but no files matched data_glob")
|
|
if args.condition_id < 0:
|
|
cond_id = torch.randint(0, cond_vocab_size, (args.batch_size,), device=device)
|
|
else:
|
|
cond_id = torch.full((args.batch_size,), int(args.condition_id), device=device, dtype=torch.long)
|
|
cond = cond_id
|
|
|
|
cond_cont = None
|
|
if route_type1_cols:
|
|
ref_glob = cfg.get("data_glob") or args.data_glob
|
|
if ref_glob:
|
|
ref_glob = str(resolve_path(Path(args.config).parent, ref_glob)) if args.config else ref_glob
|
|
base = Path(ref_glob).parent
|
|
pat = Path(ref_glob).name
|
|
refs = sorted(base.glob(pat))
|
|
if refs:
|
|
ref_path = refs[0]
|
|
ref_rows = []
|
|
with gzip.open(ref_path, "rt", newline="") as fh:
|
|
reader = csv.DictReader(fh)
|
|
for row in reader:
|
|
ref_rows.append(row)
|
|
if len(ref_rows) >= args.seq_len:
|
|
seq = ref_rows[: args.seq_len]
|
|
cond_cont = torch.zeros(args.batch_size, args.seq_len, len(route_type1_cols), device=device)
|
|
for t, row in enumerate(seq):
|
|
for i, c in enumerate(route_type1_cols):
|
|
cond_cont[:, t, i] = float(row[c])
|
|
cond_cont = normalize_cont(
|
|
cond_cont,
|
|
route_type1_cols,
|
|
mean,
|
|
std,
|
|
transforms=transforms,
|
|
quantile_probs=quantile_probs,
|
|
quantile_values=quantile_values,
|
|
use_quantile=use_quantile,
|
|
)
|
|
|
|
trend = None
|
|
if temporal_model is not None:
|
|
trend = temporal_model.generate(args.batch_size, args.seq_len, device, cond_cont=cond_cont)
|
|
if temporal_focus_type4 and type4_cols:
|
|
type4_model_idx = [model_cont_cols.index(c) for c in type4_cols if c in model_cont_cols]
|
|
if type4_model_idx:
|
|
trend_mask = torch.zeros(1, 1, len(model_cont_cols), device=device, dtype=trend.dtype)
|
|
trend_mask[:, :, type4_model_idx] = 1.0
|
|
trend = trend * trend_mask
|
|
elif temporal_exclude_type4 and type4_cols:
|
|
type4_model_idx = [model_cont_cols.index(c) for c in type4_cols if c in model_cont_cols]
|
|
if type4_model_idx:
|
|
trend_mask = torch.ones(1, 1, len(model_cont_cols), device=device, dtype=trend.dtype)
|
|
trend_mask[:, :, type4_model_idx] = 0.0
|
|
trend = trend * trend_mask
|
|
|
|
for t in reversed(range(args.timesteps)):
|
|
t_batch = torch.full((args.batch_size,), t, device=device, dtype=torch.long)
|
|
eps_pred, logits = model(x_cont, x_disc, t_batch, cond, cond_cont=cond_cont)
|
|
|
|
a_t = alphas[t]
|
|
a_bar_t = alphas_cumprod[t]
|
|
|
|
if cont_target == "x0":
|
|
x0_pred = eps_pred
|
|
if cont_clamp_x0 > 0:
|
|
x0_pred = torch.clamp(x0_pred, -cont_clamp_x0, cont_clamp_x0)
|
|
eps_pred = (x_cont - torch.sqrt(a_bar_t) * x0_pred) / torch.sqrt(1.0 - a_bar_t)
|
|
coef1 = 1.0 / torch.sqrt(a_t)
|
|
coef2 = (1 - a_t) / torch.sqrt(1 - a_bar_t)
|
|
mean_x = coef1 * (x_cont - coef2 * eps_pred)
|
|
if t > 0:
|
|
noise = torch.randn_like(x_cont)
|
|
x_cont = mean_x + torch.sqrt(betas[t]) * noise
|
|
else:
|
|
x_cont = mean_x
|
|
if args.clip_k > 0:
|
|
x_cont = torch.clamp(x_cont, -args.clip_k, args.clip_k)
|
|
|
|
for i, logit in enumerate(logits):
|
|
if t == 0:
|
|
probs = F.softmax(logit, dim=-1)
|
|
x_disc[:, :, i] = torch.argmax(probs, dim=-1)
|
|
else:
|
|
mask = x_disc[:, :, i] == mask_tokens[i]
|
|
if mask.any():
|
|
probs = F.softmax(logit, dim=-1)
|
|
sampled = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(
|
|
args.batch_size, args.seq_len
|
|
)
|
|
x_disc[:, :, i][mask] = sampled[mask]
|
|
|
|
if trend is not None:
|
|
x_cont = x_cont + trend
|
|
x_cont = x_cont.cpu()
|
|
x_disc = x_disc.cpu()
|
|
|
|
if args.clip_k > 0:
|
|
x_cont = torch.clamp(x_cont, -args.clip_k, args.clip_k)
|
|
|
|
if use_quantile:
|
|
q_vals = {c: quantile_values[c] for c in model_cont_cols}
|
|
x_cont = inverse_quantile_transform(x_cont, model_cont_cols, quantile_probs, q_vals)
|
|
else:
|
|
mean_vec = torch.tensor([mean[c] for c in model_cont_cols], dtype=x_cont.dtype)
|
|
std_vec = torch.tensor([std[c] for c in model_cont_cols], dtype=x_cont.dtype)
|
|
x_cont = x_cont * std_vec + mean_vec
|
|
for i, c in enumerate(model_cont_cols):
|
|
if transforms.get(c) == "log1p":
|
|
x_cont[:, :, i] = torch.expm1(x_cont[:, :, i])
|
|
if cont_post_calibrate and quantile_raw_values and quantile_probs:
|
|
q_raw = {c: quantile_raw_values[c] for c in model_cont_cols}
|
|
x_cont = quantile_calibrate_to_real(x_cont, model_cont_cols, quantile_probs, q_raw)
|
|
if vmin and vmax:
|
|
for i, c in enumerate(model_cont_cols):
|
|
lo = vmin.get(c, None)
|
|
hi = vmax.get(c, None)
|
|
if lo is None or hi is None:
|
|
continue
|
|
lo = float(lo)
|
|
hi = float(hi)
|
|
if cont_bound_mode == "none":
|
|
continue
|
|
if cont_bound_mode == "sigmoid":
|
|
x_cont[:, :, i] = lo + (hi - lo) * torch.sigmoid(x_cont[:, :, i])
|
|
elif cont_bound_mode == "soft_tanh":
|
|
mid = 0.5 * (lo + hi)
|
|
half = 0.5 * (hi - lo)
|
|
denom = cont_bound_strength if cont_bound_strength > 0 else 1.0
|
|
x_cont[:, :, i] = mid + half * torch.tanh(x_cont[:, :, i] / denom)
|
|
else:
|
|
x_cont[:, :, i] = torch.clamp(x_cont[:, :, i], lo, hi)
|
|
|
|
if cont_post_scale:
|
|
for i, c in enumerate(model_cont_cols):
|
|
if c in cont_post_scale:
|
|
try:
|
|
scale = float(cont_post_scale[c])
|
|
except Exception:
|
|
scale = 1.0
|
|
x_cont[:, :, i] = x_cont[:, :, i] * scale
|
|
|
|
full_cont = torch.zeros(args.batch_size, args.seq_len, len(cont_cols), dtype=x_cont.dtype)
|
|
for i, c in enumerate(model_cont_cols):
|
|
full_idx = cont_cols.index(c)
|
|
full_cont[:, :, full_idx] = x_cont[:, :, i]
|
|
if cond_cont is not None and route_type1_cols:
|
|
cond_denorm = denormalize_cont_tensor(
|
|
cond_cont.cpu(),
|
|
route_type1_cols,
|
|
mean,
|
|
std,
|
|
transforms=transforms,
|
|
quantile_probs=quantile_probs,
|
|
quantile_values=quantile_values,
|
|
use_quantile=use_quantile,
|
|
)
|
|
for i, c in enumerate(route_type1_cols):
|
|
full_idx = cont_cols.index(c)
|
|
full_cont[:, :, full_idx] = cond_denorm[:, :, i]
|
|
for c in route_type5_cols:
|
|
if c.endswith("Z"):
|
|
base = c[:-1]
|
|
if base in cont_cols:
|
|
bidx = cont_cols.index(base)
|
|
cidx = cont_cols.index(c)
|
|
full_cont[:, :, cidx] = full_cont[:, :, bidx]
|
|
|
|
header = read_header(data_path)
|
|
out_cols = [c for c in header if c != time_col or args.include_time]
|
|
if args.include_condition and use_condition:
|
|
out_cols = ["__cond_file_id"] + out_cols
|
|
|
|
os.makedirs(os.path.dirname(args.out), exist_ok=True)
|
|
with open(args.out, "w", newline="", encoding="utf-8") as f:
|
|
writer = csv.DictWriter(f, fieldnames=out_cols)
|
|
writer.writeheader()
|
|
|
|
row_index = 0
|
|
for b in range(args.batch_size):
|
|
for t in range(args.seq_len):
|
|
row = {}
|
|
if args.include_condition and use_condition:
|
|
row["__cond_file_id"] = str(int(cond[b].item())) if cond is not None else "-1"
|
|
if args.include_time and time_col in header:
|
|
row[time_col] = str(row_index)
|
|
for i, c in enumerate(cont_cols):
|
|
val = float(full_cont[b, t, i])
|
|
if int_like.get(c, False):
|
|
row[c] = str(int(round(val)))
|
|
else:
|
|
dec = int(max_decimals.get(c, 6))
|
|
fmt = ("%%.%df" % dec) if dec > 0 else "%.0f"
|
|
row[c] = fmt % val
|
|
for i, c in enumerate(disc_cols):
|
|
tok_idx = int(x_disc[b, t, i])
|
|
tok = inv_vocab[c][tok_idx] if tok_idx < len(inv_vocab[c]) else "<UNK>"
|
|
if tok == "<UNK>" and c in top_token:
|
|
tok = top_token[c]
|
|
row[c] = tok
|
|
writer.writerow(row)
|
|
row_index += 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|