v0.2: mmap, SIMD, DistilBERT, rayon & CI

Bump crate to v0.2 and introduce several performance, compatibility, and CI improvements.

Key changes:
- Version bump to 0.2.0 and add dependencies: memmap2, rayon, unicode-normalization.
- Add GitHub automation: dependabot config, CI workflow (check/test/clippy/fmt/doc/bench/msrv) and release workflow.
- Memory-mapped SafeTensors support: SafeTensorsFile::load_mmap, DataStore enum (Owned/Mapped), data_offset handling, and Embedder now loads weights with ModelWeights::load_mmap to avoid copying large model files.
- Extended SafeTensors parsing: header parsing refactor, BF16 support (bf16_to_f32), F16/F32 handling, more robust error checks, and unit tests.
- Multi-model weight support: ModelConfig.model_type, DistilBERT-aware loading path (no token_type_embeddings, different tensor names) and ModelWeights.from_safetensors abstraction; token_type_embeddings made optional.
- SIMD-accelerated tensor primitives: new tensor/simd module with AVX2 implementations and scalar fallbacks; used in matmul and ops for faster matmul/elementwise/scalar operations; includes tests.
- Parallel tokenization: tokenizer.encode_batch now uses rayon for parallel encoding.
- Small fixes and cleanups: clamp GELU constant precision, lib crate-level clippy allows, example doc comment formatting, and test/fixture tooling (scripts/generate_golden.py and tests/golden_test.rs).

These changes reduce memory usage and improve performance for large models, add broader dtype support, and improve CI/automation.

Author: Tuntii

.github/dependabot.yml

@@ -0,0 +1,15 @@
+version: 2
+updates:
+  - package-ecosystem: "cargo"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    commit-message:
+      prefix: "deps"
+
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    commit-message:
+      prefix: "ci"

.github/workflows/ci.yml

@@ -0,0 +1,86 @@
+name: CI
+
+on:
+  push:
+    branches: [main, master]
+  pull_request:
+    branches: [main, master]
+
+env:
+  CARGO_TERM_COLOR: always
+  RUSTFLAGS: -Dwarnings
+
+jobs:
+  check:
+    name: Check
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo check --all-targets
+
+  test:
+    name: Test (${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo test --verbose
+
+  clippy:
+    name: Clippy
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+        with:
+          components: clippy
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo clippy --all-targets -- -D warnings
+
+  fmt:
+    name: Formatting
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+        with:
+          components: rustfmt
+      - run: cargo fmt --all -- --check
+
+  doc:
+    name: Documentation
+    runs-on: ubuntu-latest
+    env:
+      RUSTDOCFLAGS: -Dwarnings
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo doc --no-deps
+
+  bench:
+    name: Benchmarks (compile check)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo bench --no-run
+
+  msrv:
+    name: MSRV (1.70)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@master
+        with:
+          toolchain: "1.70"
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo check

.github/workflows/release.yml

@@ -0,0 +1,45 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - "v*"
+
+permissions:
+  contents: write
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  # Run full CI first
+  ci:
+    uses: ./.github/workflows/ci.yml
+
+  publish-check:
+    name: Publish dry-run
+    needs: ci
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo publish --dry-run
+
+  release:
+    name: Create GitHub Release
+    needs: [ci, publish-check]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Extract version from tag
+        id: version
+        run: echo "VERSION=${GITHUB_REF#refs/tags/v}" >> $GITHUB_OUTPUT
+
+      - name: Create Release
+        uses: softprops/action-gh-release@v2
+        with:
+          name: v${{ steps.version.outputs.VERSION }}
+          generate_release_notes: true
+          draft: false
+          prerelease: ${{ contains(github.ref, '-') }}

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hypembed"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2021"
 rust-version = "1.70"
 description = "A pure-Rust, local-first embedding inference library. No external ML runtimes."
@@ -13,6 +13,9 @@ categories = ["science", "text-processing"]
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 thiserror = "2"
+memmap2 = "0.9"
+rayon = "1.10"
+unicode-normalization = "0.1"
 
 [dev-dependencies]
 criterion = { version = "0.5", features = ["html_reports"] }

examples/basic_embed.rs

@@ -1,11 +1,11 @@
-///! Basic example: load a model and generate embeddings.
-///!
-///! To run this example, download a BERT-like model (e.g., all-MiniLM-L6-v2)
-///! and place config.json, vocab.txt, and model.safetensors in a directory.
-///!
-///! ```bash
-///! cargo run --example basic_embed -- ./path/to/model
-///! ```
+//! Basic example: load a model and generate embeddings.
+//!
+//! To run this example, download a BERT-like model (e.g., all-MiniLM-L6-v2)
+//! and place config.json, vocab.txt, and model.safetensors in a directory.
+//!
+//! ```bash
+//! cargo run --example basic_embed -- ./path/to/model
+//! ```
 
 use hypembed::{Embedder, EmbeddingOptions, PoolingStrategy};
 use std::env;

scripts/generate_golden.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+"""Generate golden reference embeddings for the HypEmbed test suite.
+
+Usage:
+    python scripts/generate_golden.py \
+        --model sentence-transformers/all-MiniLM-L6-v2 \
+        --output tests/fixtures/golden.json
+
+Requires: pip install sentence-transformers
+"""
+
+import argparse
+import json
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate golden embeddings")
+    parser.add_argument("--model", default="sentence-transformers/all-MiniLM-L6-v2",
+                        help="HuggingFace model name or path")
+    parser.add_argument("--output", default="tests/fixtures/golden.json",
+                        help="Output JSON file path")
+    args = parser.parse_args()
+
+    from sentence_transformers import SentenceTransformer
+
+    # Test sentences covering various edge cases
+    sentences = [
+        "Hello world",
+        "Rust is a systems programming language.",
+        "Machine learning models generate embedding vectors.",
+        "The quick brown fox jumps over the lazy dog.",
+        "café résumé naïve",  # Accented characters
+        "你好世界",  # Chinese characters
+        "Short",
+        "",  # Empty string
+        "A " * 100,  # Long repeated text
+        "Hello, World! How are you doing today?",
+    ]
+
+    model = SentenceTransformer(args.model)
+
+    golden = {}
+    for sentence in sentences:
+        if not sentence.strip():
+            continue  # Skip empty
+        embedding = model.encode([sentence], normalize_embeddings=True)
+        golden[sentence] = embedding.tolist()
+
+    with open(args.output, "w") as f:
+        json.dump(golden, f, indent=2)
+
+    print(f"Generated golden embeddings for {len(golden)} sentences")
+    print(f"Model: {args.model}")
+    print(f"Output: {args.output}")
+    print(f"Embedding dim: {len(list(golden.values())[0][0])}")
+
+if __name__ == "__main__":
+    main()

src/lib.rs

@@ -1,3 +1,6 @@
+#![allow(clippy::too_many_arguments)]
+#![allow(clippy::empty_line_after_doc_comments)]
+#![allow(clippy::needless_range_loop)]
 //! # HypEmbed
 //!
 //! A pure-Rust, local-first embedding inference library.

src/model/config.rs

@@ -39,6 +39,11 @@ pub struct ModelConfig {
     /// Number of token type (segment) embeddings.
     #[serde(default = "default_type_vocab_size")]
     pub type_vocab_size: usize,
+
+    /// Model type (e.g., "bert", "distilbert", "xlm-roberta").
+    /// Auto-detected from config.json.
+    #[serde(default)]
+    pub model_type: Option<String>,
 }
 
 fn default_layer_norm_eps() -> f64 {
@@ -88,6 +93,16 @@ impl ModelConfig {
     pub fn ln_eps(&self) -> f32 {
         self.layer_norm_eps as f32
     }
+
+    /// Check if this is a DistilBERT model.
+    pub fn is_distilbert(&self) -> bool {
+        self.model_type.as_deref() == Some("distilbert")
+    }
+
+    /// Check if this is a standard BERT-type model (BERT, MiniLM, etc.).
+    pub fn is_bert(&self) -> bool {
+        !self.is_distilbert()
+    }
 }
 
 #[cfg(test)]

src/model/embedding.rs

@@ -16,7 +16,7 @@ use crate::tensor::layernorm;
 /// - `token_type_ids`: [batch_size, seq_len] segment IDs (0 or 1)
 /// - `word_embeddings`: [vocab_size, hidden_size]
 /// - `position_embeddings`: [max_positions, hidden_size]
-/// - `token_type_embeddings`: [type_vocab_size, hidden_size]
+/// - `token_type_embeddings`: [type_vocab_size, hidden_size] or None for models without (e.g., DistilBERT)
 /// - `ln_weight`, `ln_bias`: LayerNorm parameters [hidden_size]
 /// - `eps`: LayerNorm epsilon
 ///
@@ -27,7 +27,7 @@ pub fn compute_embeddings(
     token_type_ids: &[Vec<u32>],
     word_embeddings: &Tensor,
     position_embeddings: &Tensor,
-    token_type_embeddings: &Tensor,
+    token_type_embeddings: Option<&Tensor>,
     ln_weight: &Tensor,
     ln_bias: &Tensor,
     eps: f32,
@@ -39,24 +39,30 @@ pub fn compute_embeddings(
     let mut output = vec![0.0f32; batch_size * seq_len * hidden_size];
     let word_data = word_embeddings.data();
     let pos_data = position_embeddings.data();
-    let tt_data = token_type_embeddings.data();
+    let tt_data = token_type_embeddings.map(|t| t.data());
 
     for b in 0..batch_size {
         for s in 0..seq_len {
             let token_id = input_ids[b][s] as usize;
-            let type_id = token_type_ids[b][s] as usize;
             let out_offset = (b * seq_len + s) * hidden_size;
 
-            // word_embed[token_id] + position_embed[s] + token_type_embed[type_id]
+            // word_embed[token_id] + position_embed[s]
             let word_offset = token_id * hidden_size;
             let pos_offset = s * hidden_size;
-            let tt_offset = type_id * hidden_size;
 
             for h in 0..hidden_size {
                 output[out_offset + h] =
                     word_data[word_offset + h]
-                    + pos_data[pos_offset + h]
-                    + tt_data[tt_offset + h];
+                    + pos_data[pos_offset + h];
+            }
+
+            // + token_type_embed[type_id] (if available)
+            if let Some(tt) = tt_data {
+                let type_id = token_type_ids[b][s] as usize;
+                let tt_offset = type_id * hidden_size;
+                for h in 0..hidden_size {
+                    output[out_offset + h] += tt[tt_offset + h];
+                }
             }
         }
     }
@@ -89,7 +95,7 @@ mod tests {
 
         let result = compute_embeddings(
             &input_ids, &token_type_ids,
-            &word_emb, &pos_emb, &tt_emb,
+            &word_emb, &pos_emb, Some(&tt_emb),
             &ln_w, &ln_b, 1e-12,
         ).unwrap();
 

src/model/encoder.rs

@@ -88,6 +88,7 @@ mod tests {
             layer_norm_eps: 1e-12,
             hidden_act: "gelu".into(),
             type_vocab_size: 2,
+            model_type: None,
         }
     }