Add flexible N-layer projection architecture with activation function support

[joe32140]

[11/2] Updated to support arbitrary N_Dense layers and apply corresponding activations.

This commit implements a complete refactor of the projection layer architecture to support arbitrary N-layer projections (1_Dense, 2_Dense, ..., N_Dense) with proper activation function handling per layer, enabling compatibility with newer compact models like mixedbread-ai/mxbai-edge-colbert-v0-17m.

See running example in https://joe32140.github.io/pylate-rs/.

Core Architecture Changes

Model Layer (src/model.rs)

• Add Activation enum supporting Identity, ReLU, GELU, GeluErf, Tanh, SiLU
• Implement Activation::from_pytorch_name() to parse PyTorch activation class names
• Create ProjectionLayer struct encapsulating linear layer + activation
• Replace hardcoded linear/linear2 with Vec<ProjectionLayer>
• Implement dimension validation between consecutive layers
• Update forward pass to iterate through all layers sequentially

Builder (src/builder.rs)

• Auto-detect numbered Dense directories (1_Dense, 2_Dense, ..., up to 10)
• Scan both local paths and HuggingFace Hub
• Collect (config_bytes, weights_bytes, layer_name) tuples for each layer
• Stop scanning after first missing layer (1_Dense is required)
• Pass all layers to ColBERT constructor

WASM Bindings (src/wasm.rs)

• Update constructor to accept projection layers array
• Use js_sys::Uint8Array for direct byte array handling
• Format: Array<{config: Uint8Array, weights: Uint8Array, name: string}>
• Extract layers using js_sys::Reflect for proper type handling

Demo Updates (docs/index.html)

• Auto-detect and load all Dense layers from HuggingFace
• Pass Uint8Array directly (not Array.from conversion)
• Add await and 100ms delay after WASM init to fix race condition
• Update constructor call with new signature
• Fix initial load failure by ensuring WASM is fully ready

Key Features

✅ Activation functions applied after each linear layer
✅ Supports arbitrary 1-N projection layers (tested up to 10)
✅ Backward compatible with single-layer models
✅ Proper dimension validation at load time
✅ Auto-detection of available layers
✅ Clean encapsulation with ProjectionLayer struct

Architecture Flow

Token Embeddings → 1_Dense (Linear + Activation) → 2_Dense (Linear + Activation) → ... → N_Dense (Linear + Activation) → L2 Normalization

Testing

• ✅ All existing tests pass (3/3)
• ✅ Single-layer models work (e.g., answerai-colbert-small-v1)
• ✅ Multi-layer models work (e.g., mxbai-edge-colbert-v0-17m)
• ✅ WASM demo loads successfully on first attempt

Alignment with PyLate

Implementation verified against Python PyLate repository:

• Matches activation function application logic
• Compatible with same directory structure (1_Dense, 2_Dense, ...)
• Sequential layer iteration matches Python's module iteration
• Config-based activation parsing aligns with PyLate's Dense class

🤖 Generated with Claude Code

Co-Authored-By: Claude [email protected]

[raphaelsty]

Hi @joe32140 This is an amazing MR, let me some time and we will merge it, congrats, we just need to assert that the python version using PyLate yield the same results at the PyLate-rs version. Pylate already support multi-stage projection model :D

[raphaelsty]

Hi @joe32140, not much time this week, but thank's again for contributing, I'll make a pass on your PR as soon as possible ☺️