feat: Improve loading, support more formats

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "model2vec-rs"
-version = "0.1.4"
+version = "0.2.0"
 edition = "2021"
 description = "Official Rust Implementation of Model2Vec"
 readme = "README.md"
@@ -15,7 +15,7 @@ exclude = ["tests/*"]
 
 [features]
 default      = ["onig", "hf-hub"]
-hf-hub       = ["dep:hf-hub"]
+hf-hub       = ["dep:hf-hub", "dep:ureq"]
 local-only   = []
 onig         = ["tokenizers/onig",
                 "tokenizers/progressbar",
@@ -32,11 +32,14 @@ tokenizers = { version = "0.21", default-features = false }
 safetensors = "0.5"
 ndarray = "0.15"
 hf-hub = { version = "0.4", default-features = false, features = ["ureq"], optional = true }
+ureq = { version = "2", optional = true }
 clap = { version = "4.0", features = ["derive"] }
 anyhow = "1.0"
 serde_json = "1.0"
+serde = "1.0"
 half = "2.0"
 
 [dev-dependencies]
 approx = "0.5"
 serde_json = "1.0"
+tempfile = "3"

README.md

@@ -91,6 +91,25 @@ fn main() -> Result<()> {
 }
 ```
 
+**Alternative: Loading from in-memory bytes (`wasm` / embedded use cases):**
+
+When filesystem access is unavailable, use `from_bytes` instead of `from_pretrained`:
+
+```rust
+use model2vec_rs::model::StaticModel;
+
+let model = StaticModel::from_bytes(
+    include_bytes!("path/to/tokenizer.json").as_ref(),
+    include_bytes!("path/to/model.safetensors").as_ref(),
+    include_bytes!("path/to/config.json").as_ref(),
+    None, // normalize: None reads the value from config.json
+)?;
+
+let embeddings = model.encode(&["Hello world".to_string()]);
+```
+
+See the [feature flags](#feature-flags) section for `wasm` and `local-only` build configurations.
+
 ---
 
 ### 2. Using the `model2vec-rs` CLI
@@ -117,31 +136,13 @@ The compiled binary installed via `cargo install` is significantly faster (often
     model2vec-rs encode my_texts.txt "minishlab/potion-base-8M" --output embeddings_output.json
     ```
 
-**c. (Alternative for Developers) Running CLI from a cloned repository:**
-
-```shell
-# Clone and navigate to the repository directory
-git clone https://github.com/MinishLab/model2vec-rs.git
-cd model2vec-rs
-
-# Build and run with release optimizations (recommended for better performance):
-cargo run --release -- encode "Hello world" "minishlab/potion-base-8M"
-
-# For quicker development cycles instead (slower execution):
-cargo run -- encode "Hello world" "minishlab/potion-base-8M"
-
-# Alternatively, build the executable first:
-cargo build --release
-
-# Then run with:
-./target/release/model2vec-rs encode "Hello world" "minishlab/potion-base-8M"
-```
 
 ## Features
 
 *   **Fast Inference:** Optimized Rust implementation for fast embedding generation.
 *   **Hugging Face Hub Integration:** Load pre-trained Model2Vec models directly from the Hugging Face Hub using model IDs, or use models from local paths.
 *   **Model Formats:** Supports models with f32, f16, and i8 weight types stored in `safetensors` files.
+*   **Flexible Loading:** Load models from the Hub (`from_pretrained`), raw bytes (`from_bytes`), or owned/borrowed data (`from_owned`, `from_borrowed`) for embedded and WASM use cases.
 *   **Batch Processing:** Encodes multiple sentences in batches.
 *   **Configurable Encoding:** Allows customization of maximum sequence length and batch size during encoding.
 
@@ -154,23 +155,13 @@ The crate exposes a few feature combinations for different runtimes:
 * `local-only`: disable remote model downloads and restrict loading to local paths or `from_bytes(...)`
 * `wasm`: minimal WebAssembly-oriented feature set for in-memory loading via `from_bytes(...)`
 
-Typical invocations are:
-
-* native local-only build:
-  `cargo build --no-default-features --features onig,local-only`
-* wasm check:
-  `RUSTFLAGS='--cfg getrandom_backend="wasm_js"' cargo check --no-default-features --features wasm --target wasm32-unknown-unknown`
-
-The `wasm` feature is intended for `wasm32-unknown-unknown` builds that load models
-from in-memory bytes, for example after fetching assets over HTTP or embedding them
-into the binary. Direct filesystem access is usually not available in browser-style
-WebAssembly environments, so callers should pass file contents through `from_bytes(...)`.
-Remote Hugging Face downloads are not available in this mode.
-
-For `wasm32-unknown-unknown`, `getrandom` also requires a target-specific backend
-configuration. The minimal check command is:
+Typical invocations:
 
 ```bash
+# native local-only build
+cargo build --no-default-features --features onig,local-only
+
+# wasm (requires getrandom backend config)
 RUSTFLAGS='--cfg getrandom_backend="wasm_js"' \
 cargo check --no-default-features --features wasm --target wasm32-unknown-unknown
 ```
@@ -190,6 +181,7 @@ A variety of pre-trained Model2Vec models are available on the [HuggingFace Hub
 | [potion-base-32M](https://huggingface.co/minishlab/potion-base-32M)   | English    | [bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) | 32.3M   | General   |
 | [potion-multilingual-128M](https://huggingface.co/minishlab/potion-multilingual-128M) | Multilingual | [bge-m3](https://huggingface.co/BAAI/bge-m3)      | 128M    | General   |
 | [potion-retrieval-32M](https://huggingface.co/minishlab/potion-retrieval-32M) | English    | [bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) | 32.3M   | Retrieval |
+| [potion-code-16M](https://huggingface.co/minishlab/potion-code-16M)           | Code       | [CodeRankEmbed](https://huggingface.co/nomic-ai/CodeRankEmbed)    | 16M     | Code      |
 | [potion-base-8M](https://huggingface.co/minishlab/potion-base-8M)     | English    | [bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) | 7.5M    | General   |
 | [potion-base-4M](https://huggingface.co/minishlab/potion-base-4M)     | English    | [bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) | 3.7M    | General   |
 | [potion-base-2M](https://huggingface.co/minishlab/potion-base-2M)     | English    | [bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) | 1.8M    | General   |

src/main.rs

@@ -7,6 +7,19 @@ use std::path::Path;
 mod model;
 use model::StaticModel;
 
+fn write_output<T: serde::Serialize + std::fmt::Debug>(data: &T, path: Option<String>) -> Result<()> {
+    match path {
+        Some(p) => {
+            let file = File::create(&p).context("failed to create output file")?;
+            serde_json::to_writer(BufWriter::new(file), data).context("failed to write JSON")
+        }
+        None => {
+            println!("{data:#?}");
+            Ok(())
+        }
+    }
+}
+
 #[derive(Parser)]
 #[command(author, version, about = "Model2Vec Rust CLI")]
 struct Cli {
@@ -40,40 +53,22 @@ enum Commands {
 fn main() -> Result<()> {
     let cli = Cli::parse();
     match cli.cmd {
-        // Encode multiple sentences from a file or input string
         Commands::Encode { input, model, output } => {
             let texts = if Path::new(&input).exists() {
                 std::fs::read_to_string(&input)?.lines().map(str::to_string).collect()
             } else {
                 vec![input]
             };
-
-            let m = StaticModel::from_pretrained(&model, None, None, None)?;
-            let embs = m.encode(&texts);
-
-            if let Some(output) = output {
-                let file = File::create(&output).context("failed to create output file")?;
-                let writer = BufWriter::new(file);
-                serde_json::to_writer(writer, &embs).context("failed to write embeddings to JSON")?;
-            } else {
-                println!("{:?}", embs);
-            }
+            let embs = StaticModel::from_pretrained(&model, None, None, None)?.encode(&texts);
+            write_output(&embs, output)?;
         }
-        // Encode a single sentence
         Commands::EncodeSingle {
             sentence,
             model,
             output,
         } => {
-            let m = StaticModel::from_pretrained(&model, None, None, None)?;
-            let embedding = m.encode_single(&sentence);
-
-            if let Some(path) = output {
-                let file = File::create(path).context("creating output file failed")?;
-                serde_json::to_writer(BufWriter::new(file), &embedding).context("writing JSON failed")?;
-            } else {
-                println!("{embedding:#?}");
-            }
+            let embedding = StaticModel::from_pretrained(&model, None, None, None)?.encode_single(&sentence);
+            write_output(&embedding, output)?;
         }
     }
     Ok(())