Add Lux.jl file classifier (known/unknown) with offline trainer

Each uploaded file is scored by a fixed-structure neural net that labels it
known (resembling the training set) or unknown — novelty detection over the
first 16 + last 16 bytes (scaled to [0,1]), Dense(32->64->16->2), argmax.

- src/model.jl: shared architecture + byte->feature mapping (trainer + server)
- src/classify.jl: load committed artifact, classify a file at inference
- bin/train.jl: offline trainer, 1:1 blended negatives (random + grab-bag),
  seeded 80/20 split, writes model/classifier.jld2
- worker: classify (annotate-only) and log classification=known|unknown
- config: FS_MODEL_PATH; server fails fast if the artifact is missing
- deps: Lux, JLD2, Optimisers, Zygote
This commit is contained in:
2026-07-02 14:13:57 -04:00
parent 6d685cfcbb
commit e55129e3a4
10 changed files with 1145 additions and 12 deletions

137
bin/train.jl Normal file
View File

@@ -0,0 +1,137 @@
#!/usr/bin/env julia
#
# Train the file-type novelty classifier and write the artifact the server loads.
#
# julia --project=. bin/train.jl <positives_dir> [negatives_dir]
#
# positives_dir every file in it (>=32 bytes) is a "known" example (class 1).
# negatives_dir optional grab-bag of *other* real file types used as "unknown"
# (class 2). If omitted/empty, unknown examples are all random.
#
# Negatives total ~1:1 with positives, split 50/50 between uniform-random byte
# vectors and grab-bag files (falling back to all-random when no grab-bag).
# One fixed seed (FS_TRAIN_SEED, default 42) drives negatives, the train/val
# split, and weight init, so the committed artifact is exactly regenerable.
#
# The model is trained on an 80% split; validation accuracy on the held-out 20%
# is reported so you can see whether it actually learned. The 80%-trained model
# is what gets saved (no retrain-on-all for this test).
using Lux
using JLD2
using Optimisers
using Zygote
using Random
using Statistics
include(joinpath(@__DIR__, "..", "src", "model.jl")) # build_model(), read_features(), FEATURE_DIM, MIN_FILE_BYTES
const EPOCHS = 200
const LEARNING_RATE = 1f-3
const VAL_FRACTION = 0.20
# --- data assembly -----------------------------------------------------------
"List regular files directly under `dir` that are large enough to featurize."
function eligible_files(dir::AbstractString)
(isempty(dir) || !isdir(dir)) && return String[]
paths = String[]
for name in readdir(dir; join=true)
isfile(name) && filesize(name) >= MIN_FILE_BYTES && push!(paths, name)
end
return paths
end
"32×N Float32 feature matrix from a list of file paths (all assumed eligible)."
function feature_matrix(paths::Vector{String})
isempty(paths) && return Matrix{Float32}(undef, FEATURE_DIM, 0)
return reduce(hcat, (read_features(p) for p in paths))
end
# logsoftmax over the class dimension (rows), numerically stabilized. Local so
# the trainer needs no extra dependency for the loss.
function logsoftmax(x::AbstractMatrix)
m = maximum(x; dims=1)
shifted = x .- m
return shifted .- log.(sum(exp.(shifted); dims=1))
end
logitcrossentropy(logits, onehot) = mean(-sum(onehot .* logsoftmax(logits); dims=1))
accuracy(logits, onehot) = mean(vec(map(i -> i[1], argmax(logits; dims=1))) .==
vec(map(i -> i[1], argmax(onehot; dims=1))))
function main()
if isempty(ARGS)
println(stderr, "usage: julia --project=. bin/train.jl <positives_dir> [negatives_dir]")
exit(2)
end
positives_dir = ARGS[1]
negatives_dir = length(ARGS) >= 2 ? ARGS[2] : ""
seed = parse(Int, get(ENV, "FS_TRAIN_SEED", "42"))
rng = Random.MersenneTwister(seed)
pos_paths = eligible_files(positives_dir)
P = length(pos_paths)
P == 0 && error("no eligible (>= $MIN_FILE_BYTES byte) files found in positives dir: $positives_dir")
grabbag = eligible_files(negatives_dir)
# Negatives ~1:1 with positives, 50/50 random vs grab-bag (all-random if the
# grab-bag is empty). Grab-bag files are sampled with replacement so we can
# always hit the target count from a small pile.
n_neg = P
n_grab = isempty(grabbag) ? 0 : n_neg ÷ 2
n_rand = n_neg - n_grab
Xpos = feature_matrix(pos_paths)
Xrand = rand(rng, Float32, FEATURE_DIM, n_rand) # uniform [0,1) ≈ random bytes / 255
grab_paths = n_grab == 0 ? String[] : rand(rng, grabbag, n_grab)
Xgrab = feature_matrix(grab_paths)
X = hcat(Xpos, Xrand, Xgrab)
N = size(X, 2)
# One-hot targets: known = [1,0], unknown = [0,1].
Y = zeros(Float32, 2, N)
Y[1, 1:P] .= 1f0 # positives -> known
Y[2, P+1:end] .= 1f0 # all negatives -> unknown
@info "assembled dataset" positives=P negatives=n_neg random_neg=n_rand grabbag_neg=n_grab total=N grabbag_pool=length(grabbag)
# Seeded shuffle + 80/20 split.
perm = randperm(rng, N)
X, Y = X[:, perm], Y[:, perm]
n_val = round(Int, VAL_FRACTION * N)
n_train = N - n_val
Xtr, Ytr = X[:, 1:n_train], Y[:, 1:n_train]
Xval, Yval = X[:, n_train+1:end], Y[:, n_train+1:end]
# --- train (full-batch, fixed epochs) ------------------------------------
model = build_model()
ps, st = Lux.setup(rng, model)
opt_state = Optimisers.setup(Optimisers.Adam(LEARNING_RATE), ps)
for epoch in 1:EPOCHS
(loss, st), back = Zygote.pullback(ps) do p
logits, st_new = model(Xtr, p, st)
logitcrossentropy(logits, Ytr), st_new
end
grads = back((one(loss), nothing))[1]
opt_state, ps = Optimisers.update(opt_state, ps, grads)
if epoch == 1 || epoch % 20 == 0 || epoch == EPOCHS
tr_logits, _ = model(Xtr, ps, st)
@info "epoch" epoch loss=loss train_acc=accuracy(tr_logits, Ytr)
end
end
# --- report + save -------------------------------------------------------
val_logits, _ = model(Xval, ps, st)
@info "validation" n_val=n_val val_acc=(n_val == 0 ? NaN : accuracy(val_logits, Yval))
out = get(ENV, "FS_MODEL_PATH", "model/classifier.jld2")
mkpath(dirname(out))
jldsave(out; ps=ps, st=st)
@info "saved model artifact" path=out
end
main()