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

Project.toml

@@ -5,14 +5,22 @@ authors = ["wardjm@gmail.com"]
 
 [deps]
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Oxygen = "df9a0d86-3283-4920-82dc-4555fc0d1d8b"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 HTTP = "1.11.0"
+JLD2 = "0.6.4"
 JSON3 = "1.14.3"
 Logging = "1.11.0"
+Lux = "1.31.4"
+Optimisers = "0.4.7"
 Oxygen = "1.10.2"
 UUIDs = "1.11.0"
+Zygote = "0.7.11"

README.md

@@ -5,8 +5,9 @@ pool of worker threads for processing. The HTTP endpoint does no real work: it
 spools each uploaded file to disk, pushes a lightweight reference onto a work
 queue, and responds immediately — staying free to accept the next upload.
 
-Right now the "processing" is just logging the received filename, to prove the
+The per-file "processing" runs each file through a small neural-network
-flow. That's the seam where real heavy-lifting goes later.
+classifier that labels it **known** (a file type resembling the training set) or
+**unknown**, and logs the result. See "File classifier" below.
 
 ## Architecture
 
@@ -82,6 +83,49 @@ Both SIGINT and SIGTERM trigger the same idempotent graceful drain
   `STOPSIGNAL SIGINT`). Give the stop timeout enough headroom to drain
   in-flight work (systemd: `TimeoutStopSec`).
 
+## File classifier
+
+Each file is scored by a fixed-structure neural network (Lux.jl) that answers a
+single binary question: is this file **known** (like the types in the training
+set) or **unknown**? It's novelty detection, not exact file-typing — it won't
+tell you "PDF", just "this looks like something I was trained on, or not".
+
+- **Features:** the first 16 bytes + last 16 bytes of the file, each scaled
+  0–255 → `[0,1]`, giving a 32-dim input. Files under 32 bytes can't form that
+  window and are classified `unknown` without touching the model.
+- **Architecture:** `Dense(32→64,relu) → Dense(64→16,relu) → Dense(16→2)`,
+  raw logits; decision is `argmax` (class 1 = known, class 2 = unknown).
+- **Artifact:** trained weights live in `model/classifier.jld2` (committed), so
+  the server just loads them at startup. Missing/unreadable ⇒ the server fails
+  fast rather than run without classification.
+- **Effect today:** *annotate-only*. The class is logged
+  (`classification=known|unknown`) but every file still moves to `done/`; the
+  classifier can't misroute real files while it's unproven.
+
+The architecture and byte→feature mapping are defined once in `src/model.jl` and
+shared by the trainer and the server, so they can't drift apart.
+
+### Training
+
+Training is a separate, offline script — it never runs in the request path:
+
+```bash
+julia --project=. bin/train.jl <positives_dir> [negatives_dir]
+```
+
+- **positives_dir** — every file in it (≥32 bytes) is a "known" example.
+- **negatives_dir** *(optional)* — a grab-bag of *other* real file types used as
+  "unknown" examples. Negatives are generated ~1:1 with positives, split 50/50
+  between uniform-random byte vectors and grab-bag files. With no grab-bag dir,
+  negatives are all random (weaker: the net may just learn "high entropy =
+  unknown" rather than your actual types, so a grab-bag of real off-distribution
+  files is recommended).
+
+The script uses an 80/20 seeded split, reports validation accuracy, and writes
+`model/classifier.jld2` (path overridable via `FS_MODEL_PATH`). A fixed seed
+(`FS_TRAIN_SEED`, default 42) drives negative generation, the split, and weight
+init, so the artifact is exactly regenerable from the same inputs.
+
 ## Configuration (environment variables)
 
 | Variable            | Default        | Meaning                                  |
@@ -93,6 +137,7 @@ Both SIGINT and SIGTERM trigger the same idempotent graceful drain
 | `FS_SPOOL_DIR`      | `data/spool`   | Incoming files (pending)                 |
 | `FS_DONE_DIR`       | `data/done`    | Files after successful processing        |
 | `FS_FAILED_DIR`     | `data/failed`  | Files whose processing threw             |
+| `FS_MODEL_PATH`     | `model/classifier.jld2` | Classifier artifact loaded at startup |
 
 > To get real parallelism, start Julia with enough threads (`-t N`) to match
 > `FS_WORKERS`. If `FS_WORKERS` exceeds available threads you'll get a warning
@@ -126,8 +171,13 @@ src/
   job.jl          Job (the queue reference)
   queue.jl        JobQueue seam + in-process ChannelQueue
   spool.jl        filename sanitizing, spool/move, startup recovery
-  worker.jl       worker loop + per-job processing (placeholder)
+  model.jl        NN architecture + byte→feature mapping (shared with trainer)
+  classify.jl     load artifact + classify a file at inference time
+  worker.jl       worker loop + per-job processing (classify + move)
   server.jl       HTTP routes/handlers
 bin/
   server.jl       entry point
+  train.jl        offline training script → model/classifier.jld2
+model/
+  classifier.jld2 committed trained weights (loaded at startup)
 ```

bin/train.jl

@@ -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()

src/FileServer.jl

@@ -5,11 +5,15 @@ using UUIDs
 using HTTP
 using JSON3
 using Oxygen
+using Lux
+using JLD2
 
 include("config.jl")
 include("job.jl")
 include("queue.jl")
 include("spool.jl")
+include("model.jl")     # build_model() + read_features(); shared with bin/train.jl
+include("classify.jl")  # Classifier + load_classifier/classify (needs model.jl)
 include("worker.jl")
 
 # Globals the HTTP handlers read at request time. Set once in `run`, before the
@@ -17,6 +21,7 @@ include("worker.jl")
 # `Config`/`ChannelQueue` types exist.
 const CONFIG = Ref{Config}()
 const QUEUE  = Ref{ChannelQueue}()
+const CLASSIFIER = Ref{Classifier}()  # loaded once at startup, shared read-only across workers
 
 include("server.jl")   # registers routes (references CONFIG/QUEUE at call time)
 
@@ -71,6 +76,11 @@ function run(; overrides...)
     CONFIG[] = cfg
     QUEUE[]  = queue
 
+    # Load the classifier before serving. Fail fast: a server that silently
+    # doesn't classify is a worse surprise than a clear startup error.
+    CLASSIFIER[] = load_classifier(cfg.model_path)
+    @info "loaded classifier" path=cfg.model_path
+
     recovered = recover_spool!(cfg, queue)
     @info "starting file-server" host=cfg.host port=cfg.port workers=cfg.worker_count capacity=cfg.queue_capacity recovered=recovered
 

src/classify.jl

@@ -0,0 +1,49 @@
+# Runtime inference side of the classifier: load the committed artifact once at
+# startup and score files. Server-only (bin/train.jl doesn't include this). The
+# architecture and feature extraction live in model.jl, which must be included
+# first.
+
+"""
+    Classifier
+
+An in-memory, ready-to-run classifier: the (stateless) Lux `model` plus the
+learned parameters `ps` and states `st` loaded from the artifact. Immutable and
+shared read-only across all worker threads — Lux inference is a pure function
+over `ps`/`st`, so no locking is needed.
+"""
+struct Classifier{M,P,S}
+    model::M
+    ps::P
+    st::S
+end
+
+"""
+    load_classifier(path) -> Classifier
+
+Rebuild the fixed architecture and load `ps`/`st` from the JLD2 artifact at
+`path`. Throws if the file is missing or unreadable — the server fails fast at
+startup rather than run silently without classification.
+"""
+function load_classifier(path::AbstractString)
+    isfile(path) || error("model artifact not found at $path (run bin/train.jl to build it)")
+    data = JLD2.load(path)
+    (haskey(data, "ps") && haskey(data, "st")) ||
+        error("model artifact at $path is missing ps/st (was it written by bin/train.jl?)")
+    return Classifier(build_model(), data["ps"], data["st"])
+end
+
+"""
+    classify(clf, path) -> Symbol
+
+Return `:known` or `:unknown` for the file at `path`. Files shorter than
+`MIN_FILE_BYTES` short-circuit to `:unknown` without touching the model
+(they can't form a 16+16 feature window). Otherwise: scale bytes, run the net,
+argmax the two logits (class 1 = known, class 2 = unknown).
+"""
+function classify(clf::Classifier, path::AbstractString)
+    feats = read_features(path)
+    feats === nothing && return :unknown
+    x = reshape(feats, FEATURE_DIM, 1)
+    logits, _ = Lux.apply(clf.model, x, clf.ps, clf.st)
+    return argmax(vec(logits)) == 1 ? :known : :unknown
+end

src/config.jl

@@ -10,6 +10,7 @@ Base.@kwdef struct Config
     spool_dir::String   = "data/spool"   # files land here on intake (pending)
     done_dir::String    = "data/done"    # files move here after successful processing
     failed_dir::String  = "data/failed"  # files move here if a worker throws
+    model_path::String  = "model/classifier.jld2"  # committed classifier artifact, loaded at startup
 end
 
 """
@@ -21,11 +22,11 @@ and for `FileServer.run(; port=...)`).
 
 Recognised variables:
   FS_HOST, FS_PORT, FS_WORKERS, FS_QUEUE_CAPACITY,
-  FS_SPOOL_DIR, FS_DONE_DIR, FS_FAILED_DIR
+  FS_SPOOL_DIR, FS_DONE_DIR, FS_FAILED_DIR, FS_MODEL_PATH
 """
 function config_from_env(; host=nothing, port=nothing, worker_count=nothing,
                           queue_capacity=nothing, spool_dir=nothing,
-                          done_dir=nothing, failed_dir=nothing)
+                          done_dir=nothing, failed_dir=nothing, model_path=nothing)
     Config(
         host           = something(host,           get(ENV, "FS_HOST", "127.0.0.1")),
         port           = something(port,           parse(Int, get(ENV, "FS_PORT", "8080"))),
@@ -34,6 +35,7 @@ function config_from_env(; host=nothing, port=nothing, worker_count=nothing,
         spool_dir      = something(spool_dir,      get(ENV, "FS_SPOOL_DIR", "data/spool")),
         done_dir       = something(done_dir,       get(ENV, "FS_DONE_DIR", "data/done")),
         failed_dir     = something(failed_dir,     get(ENV, "FS_FAILED_DIR", "data/failed")),
+        model_path     = something(model_path,     get(ENV, "FS_MODEL_PATH", "model/classifier.jld2")),
     )
 end
 

src/model.jl

@@ -0,0 +1,51 @@
+# Shared model definition + feature extraction, used by BOTH the training script
+# (bin/train.jl) and the server (worker inference). Keeping the architecture and
+# the byte->feature mapping in one place guarantees train and inference can never
+# drift out of sync (same input layout, same scaling, same net shape).
+#
+# The including scope is expected to have already done `using Lux` and
+# `using JLD2` (FileServer.jl and bin/train.jl both do), so this file adds no
+# top-level `using` of its own and works whether it's pulled into the FileServer
+# module or a bare script.
+
+"Number of input features: first 16 bytes + last 16 bytes of a file."
+const FEATURE_DIM = 32
+
+"Minimum file size (bytes) the model can accept: needs a disjoint 16+16 window."
+const MIN_FILE_BYTES = 32
+
+"""
+    build_model()
+
+The fixed-structure classifier: 32-dim byte features → two logits
+(class 1 = known, class 2 = unknown). Raw logits out (no softmax layer); the
+loss applies log-softmax and inference just takes an argmax.
+"""
+build_model() = Chain(
+    Dense(FEATURE_DIM => 64, relu),
+    Dense(64 => 16, relu),
+    Dense(16 => 2),
+)
+
+"""
+    read_features(path) -> Union{Vector{Float32}, Nothing}
+
+Read the first 16 and last 16 bytes of the file at `path`, concatenate to a
+32-element vector, and scale each byte from 0–255 into [0,1]. Returns `nothing`
+for files shorter than `MIN_FILE_BYTES` (the caller decides what that means:
+training drops them, the server classifies them as unknown without the model).
+
+Memory stays flat: we read 16 bytes from the front and `seek` to end-16 for the
+tail rather than slurping the whole file.
+"""
+function read_features(path::AbstractString)
+    sz = filesize(path)
+    sz < MIN_FILE_BYTES && return nothing
+    bytes = open(path) do io
+        head = read(io, 16)
+        seek(io, sz - 16)
+        tail = read(io, 16)
+        vcat(head, tail)
+    end
+    return Float32.(bytes) ./ 255f0
+end

src/worker.jl

@@ -10,9 +10,11 @@ For now the "work" is just logging the received filename to prove the flow —
 this is the seam where real heavy-lifting will go later.
 """
 function handle_job(job::Job, cfg::Config, worker_id::Int)
-    # --- placeholder for real heavy-lifting work -------------------------
+    # Classify the spooled file (annotate-only for now: the result is logged but
-    @info "received file" worker=worker_id id=job.id name=job.original_name size=job.size
+    # every file still moves to done/ regardless of known/unknown). Sub-32-byte
-    # --------------------------------------------------------------------
+    # files short-circuit to :unknown inside classify without touching the model.
+    classification = classify(CLASSIFIER[], job.path)
+    @info "received file" worker=worker_id id=job.id name=job.original_name size=job.size classification=classification
     dest = move_to(cfg.done_dir, job)
     @info "completed" worker=worker_id id=job.id dest=dest
     return nothing