Hugging Face's logo

jsantillana
/
vx-one

Other
English
decision-heads
calibration
cybersecurity
perimeter-triage
reproducibility
acm
Model card Files
xet
 Community
1
vx-one / models /data.py
jsantillana's picture
jsantillana
Add VX-One source code
908fd1a verified
Raw History Blame Contribute Delete
18.6 kB
"""Datasets y loaders del motor VX-One (spec F2 Β§0.1, Β§6.2, D12).
- UNSW / CIC: tensores completos en RAM (UNSW 149KΓ—40Γ—4B β‰ˆ 24 MB;
CIC train 1.98MΓ—78Γ—4B β‰ˆ 620 MB).
- EMBER: streaming por shards (46 train parquets de 60K filas; ~616 MB
decodificados c/u): por epoch se permutan los shards (seed 42+epoch),
shuffle intra-shard, batches con drop_last; prefetch de 1 shard en thread.
Shuffle de dos niveles (shard + intra) β€” aproximaciΓ³n documentada (Β§6.2).
Candidate set por corrida (D12): familias nativas con β‰₯1 fila en TRAIN bajo
el esquema de la corrida, en el ORDEN CONGELADO del registro. Pesos Ο‰
(Β§3.5) derivados de los conteos de train DE LA CORRIDA (LOAO los recalcula).
Ο† (feature stage) se ajusta SOLO con train:
- UNSW/CIC: train completo en RAM (float64).
- EMBER: submuestra uniforme de 60,000 filas de train (1,304 por shard,
rng seed 42, una pasada por los 46 shards) β€” la SE de la mediana con
n=60K es despreciable y evita 28 GB en RAM; decisiΓ³n documentada.
"""
from __future__ import annotations
import threading
from dataclasses import dataclass, field
from pathlib import Path
import numpy as np
import pyarrow.parquet as pq
import torch
from src.state.state import SOURCE_REGISTRY
from .feature_stage import FeatureStage
from .losses import class_weights
from .network import HeadSpec
from .vocab import TextVocab
ROOT = Path(__file__).resolve().parents[2]
DATA = ROOT / "data" / "processed"
REGISTRY = SOURCE_REGISTRY
ALIAS = {"unsw": "unsw_nb15_v1", "cic": "cicids2017_v1", "ember": "ember2024_v1"}
def _registry(source: str) -> dict:
"""Alias corto (unsw/cic/ember) o clave de registro β†’ entrada del registro."""
return REGISTRY[ALIAS.get(source, source)]
DDOS_DAYS = {
"Wednesday-workingHours.pcap_ISCX",
"Friday-WorkingHours-Afternoon-DDos.pcap_ISCX",
}
DDOS_FULL = {"BENIGN", "DDoS", "DoS Hulk", "DoS GoldenEye", "DoS slowloris", "DoS Slowhttptest"}
DDOS_CANDIDATES = ("normal", "tcp_flood", "http_dos", "slowloris") # orden congelado F2 D8
_DDOS_MAP = {"BENIGN": "normal", "DDoS": "tcp_flood", "DoS Hulk": "http_dos",
"DoS GoldenEye": "http_dos", "DoS slowloris": "slowloris",
"DoS Slowhttptest": "slowloris"}
EMBER_STATS_ROWS_PER_SHARD = 1_304 # Γ— 46 shards = 60,024 filas de stats
@dataclass
class Labels:
y_bin: torch.Tensor # int64 [N] 1[family != normal] (o label_binary en EMBER)
y_fam: torch.Tensor # int64 [N] Γ­ndice en candidates
y_sev: torch.Tensor | None # int64 [N] 0..3 (None en EMBER)
ddos_mask: torch.Tensor | None # bool [N] filas del subconjunto (solo CIC strat)
y_ddos: torch.Tensor | None # int64 [N] Γ­ndice en DDOS_CANDIDATES
@dataclass
class RunData:
"""Todo lo que una corrida necesita, materializado y verificado."""
source: str
scale: str
split_scheme: str # "oficial" (UNSW) | "strat" | "time" (CIC) | "temporal" (EMBER)
candidates: list[str]
ddos_candidates: list[str] | None
omega: torch.Tensor # float32 [K]
omega_ddos: torch.Tensor | None
vocab: dict[str, TextVocab] | None
feature_stage: FeatureStage
heads: HeadSpec
train: torch.Tensor # u_r float32 [Ntr, d+1]
train_txt: dict[str, torch.Tensor] | None
train_labels: Labels
val: torch.Tensor
val_txt: dict[str, torch.Tensor] | None
val_labels: Labels
test: torch.Tensor | None # None en EMBER (streaming, no se materializa)
test_txt: dict[str, torch.Tensor] | None
test_labels: Labels | None
val_fam_names: np.ndarray | None = None # familia nativa por fila (E2; -1 = fuera de candidatos)
test_fam_names: np.ndarray | None = None
ember_train_shards: list[Path] | None = None
ember_val_shards: list[Path] | None = None
ember_test_shards: list[Path] | None = None
meta: dict = field(default_factory=dict)
# ── utilidades de lectura ──────────────────────────────────────────────
def _read_x_num(path: Path) -> np.ndarray:
tbl = pq.read_table(path, columns=["x_num"])
arr = tbl.column("x_num").combine_chunks()
flat = arr.flatten().to_numpy(zero_copy_only=False)
d = arr.type.list_size
return np.asarray(flat, dtype=np.float64).reshape(-1, d)
def _frozen_candidates(source: str, fam_counts: dict[str, int]) -> list[str]:
"""D12: familias nativas con β‰₯1 fila en train, en orden congelado."""
native = _registry(source)["native_families"]
cands = [f for f in native if fam_counts.get(f, 0) > 0]
faltan = sorted(set(native) - set(cands))
if faltan:
print(f"[d12] {source}: sin filas en train β†’ fuera del candidate set: {faltan}")
return cands
def _idx_map(candidates: list[str]) -> dict[str, int]:
return {c: i for i, c in enumerate(candidates)}
# ── UNSW / CIC: datasets en RAM ────────────────────────────────────────
def _load_network_split(path: Path, source: str, split_col: str, split: str,
candidates: list[str], need_ddos: bool,
exclude_family: str | None = None
) -> tuple[np.ndarray, dict[str, list[str | None]], Labels]:
txt_keys = ("proto", "service", "state") if source == "unsw" else ("protocol",)
cols = ["x_num", split_col, "label_family", "label_severity"]
cols += [f"x_txt_{k}" for k in txt_keys]
if need_ddos:
cols += ["day", "label_full"] # el subset usa etiquetas crudas (label_full)
tbl = pq.read_table(path, columns=cols)
mask = np.asarray(tbl.column(split_col).to_pylist()) == split
x = _read_x_num(path)[mask]
fam = np.asarray(tbl.column("label_family").to_pylist())[mask]
sev = np.asarray(tbl.column("label_severity").to_pylist())[mask]
txt = {k: list(np.asarray(tbl.column(f"x_txt_{k}").to_pylist())[mask])
for k in txt_keys}
# LOAO: la familia excluida sale de train Y val (Β§6.7); test se evalΓΊa completo
if exclude_family is not None and split != "test":
keep = fam != exclude_family
x, fam, sev = x[keep], fam[keep], sev[keep]
txt = {k: [v for v, kp in zip(vals, keep) if kp] for k, vals in txt.items()}
imap = _idx_map(candidates)
# filas con familia fuera del candidate set (D12: excluida LOAO en test,
# o recon/botnet en val/test de split_time) β†’ sentinela -1; los nombres
# se conservan para E2 (evaluaciΓ³n de la familia excluida, Β§6.7)
y_fam = torch.tensor([imap.get(f, -1) for f in fam], dtype=torch.int64)
y_bin = torch.tensor([f != "normal" for f in fam], dtype=torch.int64)
y_sev = torch.tensor(sev, dtype=torch.int64)
if need_ddos:
# subset congelado (SCHEMA Β§3.3): dΓ­as Wednesday+Friday-DDoS; solo
# mapean a las 4 clases las filas con etiqueta nativa en DDOS_FULL
# (642,862 de las 918,437; el resto del dΓ­a no entrena este head)
day = np.asarray(tbl.column("day").to_pylist())[mask]
full = np.asarray(tbl.column("label_full").to_pylist())[mask]
if exclude_family is not None and split != "test":
day, full = day[keep], full[keep] # misma poda LOAO que x/fam
in_subset = np.asarray([d in DDOS_DAYS and f in DDOS_FULL
for d, f in zip(day, full)])
ddos_mask = torch.tensor(in_subset)
y_ddos = torch.zeros(len(fam), dtype=torch.int64) # alineado a la fila
y_ddos[ddos_mask] = torch.tensor(
[_idx_map(DDOS_CANDIDATES)[_DDOS_MAP[f]] for f in np.asarray(full)[in_subset]],
dtype=torch.int64)
else:
ddos_mask = y_ddos = None
return x, txt, Labels(y_bin, y_fam, y_sev, ddos_mask, y_ddos), fam
def build_network_run(source: str, scale: str, split_scheme: str,
exclude_family: str | None = None) -> RunData:
path = DATA / ("unsw_nb15/unsw_nb15.parquet" if source == "unsw"
else "cicids2017/cicids2017.parquet")
split_col = "split" if source == "unsw" else f"split_{split_scheme}"
need_ddos = source == "cic" and split_scheme == "strat"
# 1) conteos de train de la corrida β†’ candidates (D12) + Ο‰ (Β§3.5)
count_cols = [split_col, "label_family"] + (["day"] if need_ddos else [])
tbl = pq.read_table(path, columns=count_cols)
tr = np.asarray(tbl.column(split_col).to_pylist()) == "train"
fam_all = np.asarray(tbl.column("label_family").to_pylist())
if exclude_family is not None:
tr = tr & (fam_all != exclude_family)
counts = {f: int((fam_all[tr] == f).sum()) for f in _registry(source)["native_families"]}
candidates = _frozen_candidates(source, counts)
omega = torch.tensor(class_weights([counts[c] for c in candidates]))
# 2) splits (LOAO: la familia sale de train Y val; test completo Β§6.7)
x_tr, txt_tr, lab_tr, _ = _load_network_split(path, source, split_col, "train", candidates, need_ddos, exclude_family)
ddos_counts = None
if need_ddos:
sub = lab_tr.ddos_mask
ddos_counts = [int((lab_tr.y_ddos[sub] == i).sum())
for i in range(len(DDOS_CANDIDATES))]
if any(c == 0 for c in ddos_counts):
# fold LOAO dos_ddos: las clases de ataque del vector ddos viven
# en la familia excluida (el subset conserva solo BENIGN) β†’ sin
# head ddos en esta corrida (E2 evalΓΊa noul/choice en la familia)
need_ddos = False
ddos_counts = None
lab_tr = Labels(lab_tr.y_bin, lab_tr.y_fam, lab_tr.y_sev, None, None)
x_va, txt_va, lab_va, va_fams = _load_network_split(path, source, split_col, "val", candidates, need_ddos, exclude_family)
x_te, txt_te, lab_te, te_fams = _load_network_split(path, source, split_col, "test", candidates, need_ddos, None)
# 3) vocab (train de la corrida, Β§4.2) + Ο† (train)
vocab = None
if source == "unsw":
vocab = {k: TextVocab.fit(k, txt_tr[k]) for k in ("proto", "service", "state")}
stage = FeatureStage.fit(x_tr.astype(np.float64))
u_tr, _ = stage.transform(x_tr)
u_va, _ = stage.transform(x_va)
u_te, _ = stage.transform(x_te)
train_txt = {k: torch.tensor(vocab[k].encode(txt_tr[k]), dtype=torch.int64) for k in vocab} if vocab else None
val_txt = {k: torch.tensor(vocab[k].encode(txt_va[k]), dtype=torch.int64) for k in vocab} if vocab else None
test_txt = {k: torch.tensor(vocab[k].encode(txt_te[k]), dtype=torch.int64) for k in vocab} if vocab else None
ddos_candidates = list(DDOS_CANDIDATES) if need_ddos else None
omega_ddos = None
if need_ddos:
assert all(c > 0 for c in ddos_counts), f"ddos class sin filas: {ddos_counts}"
omega_ddos = torch.tensor(class_weights(ddos_counts))
heads = HeadSpec(n_candidates=len(candidates), ddos_candidates=4 if need_ddos else 0)
return RunData(
source=source, scale=scale, split_scheme=split_scheme,
candidates=candidates, ddos_candidates=ddos_candidates,
omega=omega, omega_ddos=omega_ddos, vocab=vocab, feature_stage=stage,
heads=heads,
train=torch.tensor(u_tr), train_txt=train_txt, train_labels=lab_tr,
val=torch.tensor(u_va), val_txt=val_txt, val_labels=lab_va,
test=torch.tensor(u_te), test_txt=test_txt, test_labels=lab_te,
val_fam_names=va_fams, test_fam_names=te_fams,
meta={"train_counts": {c: counts[c] for c in candidates},
"exclusion": {"exclude": exclude_family} if exclude_family else {}},
)
# ── EMBER: streaming por shards ────────────────────────────────────────
def ember_shards(subset: str) -> list[Path]:
d = DATA / "ember2024"
if subset == "train":
return [d / f"part_{i:02d}_train.parquet" for i in range(46)]
if subset == "val":
return [d / f"part_{i:02d}_val.parquet" for i in range(46, 52)]
return [d / f"part_{i:02d}_test.parquet" for i in range(52, 64)]
def _ember_labels(candidates: list[str]):
imap = _idx_map(candidates)
def decode(tbl) -> tuple[np.ndarray, torch.Tensor, torch.Tensor]:
fam = tbl.column("label_family").to_pylist()
# A5-EMBER: 'unknown' fuera del candidate set β†’ sentinela -1 (la
# mΓ‘scara de _batch_losses la saca del CE de choice)
y_fam = torch.tensor([imap.get(f, -1) for f in fam], dtype=torch.int64)
y_bin = torch.tensor(tbl.column("label_binary").to_pylist(), dtype=torch.int64)
return y_fam, y_bin
return decode
def build_ember_run(scale: str, exclude_family: str | None = None) -> RunData:
tr_shards, va_shards, te_shards = ember_shards("train"), ember_shards("val"), ember_shards("test")
# 1) conteos de train (una pasada ligera por shard, solo label_family)
counts = {f: 0 for f in _registry("ember")["native_families"]}
for p in tr_shards:
for f in pq.read_table(p, columns=["label_family"]).column("label_family").to_pylist():
counts[f] += 1
if exclude_family is not None:
counts[exclude_family] = 0
candidates = _frozen_candidates("ember2024_v1", counts)
omega = torch.tensor(class_weights([counts[c] for c in candidates]))
# 2) Ο† con submuestra uniforme de train (decisiΓ³n documentada arriba)
rng = np.random.default_rng(42)
sample_chunks: list[np.ndarray] = []
for p in tr_shards:
x = _read_x_num(p).astype(np.float32)
take = rng.choice(x.shape[0], size=EMBER_STATS_ROWS_PER_SHARD, replace=False)
sample_chunks.append(x[take].astype(np.float64))
del x
stage = FeatureStage.fit(np.concatenate(sample_chunks))
del sample_chunks
return RunData(
source="ember2024_v1", scale=scale, split_scheme="temporal",
candidates=candidates, ddos_candidates=None,
omega=omega, omega_ddos=None, vocab=None, feature_stage=stage,
heads=HeadSpec(n_candidates=len(candidates), has_score=False),
train=torch.empty(0), train_txt=None,
train_labels=Labels(torch.empty(0, dtype=torch.int64), torch.empty(0, dtype=torch.int64), None, None, None),
val=torch.empty(0), val_txt=None,
val_labels=Labels(torch.empty(0, dtype=torch.int64), torch.empty(0, dtype=torch.int64), None, None, None),
test=None, test_txt=None, test_labels=None,
ember_train_shards=tr_shards, ember_val_shards=va_shards, ember_test_shards=te_shards,
meta={"train_counts": {c: counts[c] for c in candidates}, "exclusion": {"exclude": exclude_family} if exclude_family else {}},
)
class EmberBatcher:
"""Itera batches de train (o evalΓΊa val) con streaming por shards.
Train: permutaciΓ³n de shards seed 42+epoch β†’ decode β†’ shuffle intra-shard
β†’ batches con drop_last. Prefetch de 1 shard en thread (Β§6.2).
Eval (val/test/challenge): shards en orden fijo, batches sin drop_last,
shuffle off.
"""
def __init__(self, run: RunData) -> None:
self.run = run
self.stage = run.feature_stage
self._decode = _ember_labels(run.candidates)
def _load_shard(self, path: Path) -> tuple[np.ndarray, torch.Tensor, torch.Tensor]:
tbl = pq.read_table(path, columns=["x_num", "label_family", "label_binary"])
x = _read_x_num(path) # float64 para Ο†
u, _ = self.stage.transform(x)
del x
y_fam, y_bin = self._decode(tbl)
return u, y_fam, y_bin
def _iter_shard_batches(self, path: Path, batch_size: int, rng: np.random.Generator | None):
u, y_fam, y_bin = self._load_shard(path)
n = u.shape[0]
order = rng.permutation(n) if rng is not None else np.arange(n)
n_batches = n // batch_size if rng is not None else int(np.ceil(n / batch_size))
for b in range(n_batches):
idx = order[b * batch_size:(b + 1) * batch_size]
yield (torch.tensor(u[idx]), y_fam[idx], y_bin[idx])
def _stream(self, shards: list[Path], batch_size: int, rng: np.random.Generator | None):
"""Genera batches; con prefetch de 1 shard (thread) si rng (train)."""
if rng is None: # eval: secuencial simple
for p in shards:
yield from self._iter_shard_batches(p, batch_size, None)
return
# train: productor en thread descarga el shard i+1 mientras se
# consumen los batches del shard i (prefetch=1, Β§6.2).
import queue
q: queue.Queue = queue.Queue(maxsize=1)
exc: list[BaseException] = []
def producer() -> None:
try:
for p in shards:
q.put(self._load_shard(p))
except BaseException as e:
exc.append(e)
finally:
q.put(None)
t = threading.Thread(target=producer, daemon=True)
t.start()
while True:
item = q.get()
if item is None:
break
u, y_fam, y_bin = item
n = u.shape[0]
order = rng.permutation(n)
for b in range(n // batch_size): # drop_last
idx = order[b * batch_size:(b + 1) * batch_size]
yield (torch.tensor(u[idx]), y_fam[idx], y_bin[idx])
if exc:
raise exc[0]
def train_batches(self, epoch: int, batch_size: int):
rng = np.random.default_rng(42 + epoch)
return self._stream(self.run.ember_train_shards, batch_size, rng)
def eval_batches(self, subset: str, batch_size: int):
# "train" = streaming en orden (sin shuffle/drop_last): para
# materializar features en los baselines tabulares
shards = {"train": self.run.ember_train_shards,
"val": self.run.ember_val_shards,
"test": self.run.ember_test_shards}[subset]
return self._stream(shards, batch_size, None)
# ── fΓ‘brica ────────────────────────────────────────────────────────────
def build_run(source: str, scale: str, split_scheme: str,
exclude_family: str | None = None) -> RunData:
if source == "unsw":
return build_network_run("unsw", scale, "oficial", exclude_family)
if source == "cic":
return build_network_run("cic", scale, split_scheme, exclude_family)
if source == "ember":
return build_ember_run(scale, exclude_family)
raise ValueError(source)