model.nu

packages/embed/src/model.nu — the XLM-RoBERTa-family text encoder, pure NURL on gpukit's dev-layer kernel library.

This is the model class behind BGE-M3, multilingual-e5 and friends: word/position/type embeddings + LayerNorm, N transformer blocks (bidirectional self-attention, exact-erf GELU FFN, post-LN residuals), then CLS or mean pooling and an optional L2 normalize. Every operator is a gkd_* kernel — the same dtype-generic library tensor and onnx run on — so this file contains NO kernel sources, only the wiring.

A model is a DIRECTORY (the Hugging Face layout): config.json hidden_size / num_hidden_layers / … (read here) tokenizer.json Unigram vocabulary (packages/tokenizer) model.safetensors f32 weights (packages/safetensor, mmap-backed)

( embed_open dir ) → !*Embed String ( embed_encode e text out ) → b out = the embedding vector ( embed_dim e ) ( embed_close e )

Inference is per text (batch 1): no padding, no attention mask, and the numbers match sentence-transformers row for row. A batch endpoint loops — at embedding sizes the forward pass dominates regardless.

Numerics are true float32 on the device. Verified against sentence-transformers (BGE-M3): cosine ≥ 0.9999 on a multilingual corpus (tests/embed_test.sh).

API

: i EM_POOL_CLS 0

pooling modes

: i EM_POOL_MEAN 1

: EmbedCfg

: EmbedCfg {
    i layers
    i heads
    i dim
    i ffn
    i vocab
    i maxpos  // position table rows (XLM-R: max_seq + pad offset 2)
    i padid
    f eps
    i pool
    b normalize
    i maxseq  // token cap per text (specials included)
}

: EmbedLayer

: EmbedLayer {
    GkBuf qw GkBuf qb
    GkBuf kw GkBuf kb
    GkBuf vw GkBuf vb
    GkBuf ow GkBuf ob
    GkBuf ln1w GkBuf ln1b
    GkBuf iw GkBuf ib
    GkBuf dw GkBuf db
    GkBuf ln2w GkBuf ln2b
}

: Embed

: Embed {
    * GpuKit kit
    EmbedCfg cfg
    GkBuf wemb
    GkBuf pemb
    GkBuf temb
    GkBuf elnw
    GkBuf elnb
    ( Vec EmbedLayer ) layers
    * Unigram tok
    b has_tok
    b ok
}

@ embed_open s dir → !*Embed String

Open a model directory. Pooling defaults to CLS + normalize (the BGE convention); callers can override with embed_set_pooling.

@ embed_set_pooling * Embed e i mode b normalize → v

Pooling override: mode EM_POOL_CLS | EM_POOL_MEAN, normalize on/off. (cfg is an inline struct; a field write through two levels is not an lvalue in NURL, so the setters rebuild the struct.)

@ embed_set_maxseq * Embed e i n → v

Cap on tokens per text (specials included); clamped to the model's position table.

@ embed_dim * Embed e → i

@ embed_ok * Embed e → b

@ embed_backend * Embed e → s

@ embed_close * Embed e → v

@ embed_tokenize * Embed e s text ( Vec i ) out → b

Tokenize (Unigram, <s>…</s>) with truncation to cfg.maxseq: the head of the sequence is kept and </s> re-appended, sentence-transformers style.

@ embed_encode * Embed e s text ( Vec f ) out → b

The embedding for one text. out receives embed_dim floats.

@ embed_encode_ids * Embed e ( Vec i ) ids ( Vec f ) out → b

Forward over an already-tokenized id sequence.

---

serve.nu

packages/embed/src/serve.nu — the embedding model as a service.

embed serve <model-dir> [--addr 0.0.0.0:8000] [--token T] [--maxseq N]

The HTTP surface mirrors the reference FastAPI embedding service, so existing clients work unchanged:

POST /create_embedding {"text": "..." | ["...", …], "normalize": true} → {"embeddings": [[…]], "model": "…", "dimension": N} GET /create_embedding?text=…&normalize=true (single text) GET /health {"status":"healthy", "model", "model_loaded", "device": "cuda"|"cpu", …} GET / the same (a browser poking the port should learn something, not get a 404)

Auth: no --token → open server (bind loopback!). With a token, requests must carry Authorization: Bearer <t> (or ?token=<t> for clients that cannot set headers); the compare is constant-time over the configured token. A batch loops single-text forwards — exact per-row numerics, and the forward pass dominates at embedding sizes anyway.

Handlers are top-level functions over module globals, not closures — closure environments are manual in NURL and a server's handlers live for the process (same idiom as whisper/nurllama).

API

: ~ i g_em 0 // *Embed as an address (0 = not serving)

: ~ s g_em_token ``

: ~ s g_em_name ``

: ~ i g_em_reqs 0

@ embed_serve * Embed e s name s host i port s token → i

Serve e (borrowed for the server's lifetime). Blocks until stopped.

---

main.nu

packages/embed — the embedding server. CLI:

embed serve <model-dir> [--addr HOST:PORT] [--token T] [--maxseq N] [--pool cls|mean] [--no-normalize] embed text <model-dir> <text…> one-shot: print the vector (CSV)

A model directory is the Hugging Face layout: config.json + tokenizer.json + model.safetensors (f32). BGE-M3 works as-is; any XLM-RoBERTa-family encoder does (multilingual-e5: --pool mean). The model argument can be a local directory OR a Hugging Face ref (embed serve BAAI/bge-m3): a ref is fetched into the shared ~/.nurl cache via the hub package and the downloaded directory is used.

API

@ main → i