Error handling (#28)

* Error handling

* bump version

* Changed publish workflow trigger to `release`

* Added test workflow before release

Author: wdoppenberg

.github/workflows/publish.yml

@@ -1,11 +1,49 @@
 name: Publish Package
 on:
-  push:
-    branches:
-      - master
+  release:
+    types: [ published ]
+
 jobs:
+  build-and-test:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ ubuntu-latest, macos-latest, windows-latest ]
+
+    name: Test multiple workspaces on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+
+    env:
+      CARGO_TERM_COLOR: always
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: Swatinem/rust-cache@v2
+        with:
+          workspaces: |
+            crates/glowrs
+            crates/glowrs-server
+
+      - name: Build
+        run: cargo build --verbose
+
+      - name: Check formatting
+        run: cargo fmt -- --check
+
+      - name: Check clippy
+        run: cargo clippy -- -D warnings
+
+      - name: Publish dry-run
+        run: cargo publish -p glowrs --dry-run
+
+      - name: Run tests
+        run: cargo test --verbose
+
+
   publish:
     runs-on: ubuntu-20.04
+    needs: [ build-and-test ]
     steps:
       - uses: actions/checkout@v4
 
@@ -15,7 +53,7 @@ jobs:
           LAST_PUBLISHED_VERSION=$(cargo search glowrs --limit 1 | awk '{print $3}' | tr -d '"')
           LOCAL_VERSION=$(grep -e '^version\s*=\s*"' Cargo.toml | head -1 | cut -d '"' -f2)
           if [ "$LAST_PUBLISHED_VERSION" == "$LOCAL_VERSION" ]; then
-            echo "::set-output name=skip-publish::true"
+            exit 1 # Force a failure if the versions match
           fi
 
       - uses: Swatinem/rust-cache@v2
@@ -26,4 +64,3 @@ jobs:
 
       - name: Publish glowrs
         run: cargo publish -p glowrs --token ${{ secrets.CRATES_TOKEN }}
-        if: steps.check-version.outputs.skip-publish != 'true'

.github/workflows/rust.yml

@@ -1,6 +1,14 @@
 name: Build & Test
 
 on:
+  push:
+    branches: [ "master" ]
+    paths:
+      - "crates/**"
+      - ".github/workflows/rust.yml"
+      - "tests/**"
+      - "Cargo.toml"
+      - ".cargo/**"
   pull_request:
     branches: [ "master" ]
     paths:

Cargo.toml

@@ -19,5 +19,5 @@ candle-transformers = { opt-level = 3 }
 
 [workspace.package]
 license = "Apache-2.0"
-version = "0.2.2"
+version = "0.3.0"
 

crates/glowrs/Cargo.toml

@@ -22,7 +22,6 @@ tracing = "0.1.37"
 uuid = { version = "1.6.1", features = ["v4"] }
 tokenizers = "0.19.1"
 hf-hub = { version = "0.3.2", features = ["tokio"] }
-anyhow = "1.0.79"
 thiserror = "1.0.56"
 once_cell = "1.19.0"
 

crates/glowrs/src/error.rs

@@ -0,0 +1,56 @@
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("Invalid model name: {0}")]
+    InvalidModelName(&'static str),
+    #[error("Model load error: {0}")]
+    ModelLoad(&'static str),
+    #[error("Invalid model architecture: {0}")]
+    InvalidModelConfig(&'static str),
+    #[error("Candle error: {0}")]
+    Candle(#[from] candle_core::Error),
+    #[error("Tokenization error: {0}")]
+    Tokenization(#[from] tokenizers::Error),
+    #[error("Serde JSON error: {0}")]
+    Serde(#[from] serde_json::Error),
+    #[error("IO error: {0}")]
+    IO(#[from] std::io::Error),
+    #[error("HF Hub error: {0}")]
+    HFHub(#[from] hf_hub::api::sync::ApiError),
+}
+
+pub(crate) type Result<T> = std::result::Result<T, Error>;
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_error_display() {
+        let error = Error::InvalidModelName("test");
+        assert_eq!(error.to_string(), "Invalid model name: test");
+
+        let error = Error::ModelLoad("test");
+        assert_eq!(error.to_string(), "Model load error: test");
+
+        let error = Error::InvalidModelConfig("test");
+        assert_eq!(error.to_string(), "Invalid model architecture: test");
+
+        let error = Error::Candle(candle_core::Error::UnexpectedNumberOfDims {
+            shape: (32, 32).into(),
+            expected: 3,
+            got: 2,
+        });
+        assert_eq!(
+            error.to_string(),
+            "Candle error: unexpected rank, expected: 3, got: 2 ([32, 32])"
+        );
+
+        let error = Error::IO(std::io::Error::new(std::io::ErrorKind::Other, "test"));
+        assert_eq!(error.to_string(), "IO error: test");
+
+        let error = Error::HFHub(hf_hub::api::sync::ApiError::MissingHeader("test"));
+        assert_eq!(error.to_string(), "HF Hub error: Header test is missing");
+    }
+}

crates/glowrs/src/lib.rs

@@ -1,6 +1,10 @@
 #![doc = include_str!("../README.md")]
 
+mod error;
 pub mod model;
+pub use error::Error;
+pub(crate) use error::Result;
+
 pub use model::pooling::PoolingStrategy;
 pub use model::sentence_transformer::SentenceTransformer;
 use serde::Serialize;

crates/glowrs/src/model/embedder.rs

@@ -1,4 +1,3 @@
-use anyhow::{Context, Error, Result};
 use candle_core::{DType, Module, Tensor};
 use candle_nn::VarBuilder;
 use candle_transformers::models::{
@@ -16,10 +15,10 @@ pub use candle_transformers::models::{
 use serde::Deserialize;
 
 use crate::model::device::DEVICE;
+use crate::model::pooling::{pool_embeddings, PoolingStrategy};
 use crate::model::utils::normalize_l2;
-use crate::Usage;
+use crate::{Error, Result, Usage};
 
-use crate::model::pooling::{pool_embeddings, PoolingStrategy};
 #[cfg(test)]
 use candle_nn::VarMap;
 
@@ -35,16 +34,17 @@ struct BaseModelConfig {
 }
 
 pub(crate) fn parse_config(config_str: &str) -> Result<ModelConfig> {
+    use Error::*;
     let base_config: BaseModelConfig = serde_json::from_str(config_str)?;
 
     let config = match base_config.architectures {
         Some(arch) => {
             if arch.is_empty() {
-                return Err(Error::msg("No architectures found"));
+                return Err(InvalidModelConfig("No architectures found"));
             }
 
             if arch.len() > 1 {
-                return Err(Error::msg("Multiple architectures not supported"));
+                return Err(InvalidModelConfig("Multiple architectures not supported"));
             }
 
             match arch.first().map(String::as_str) {
@@ -60,10 +60,10 @@ pub(crate) fn parse_config(config_str: &str) -> Result<ModelConfig> {
                     let config: DistilBertConfig = serde_json::from_str(config_str)?;
                     ModelConfig::DistilBert(config)
                 }
-                _ => return Err(Error::msg("Invalid model architecture")),
+                _ => return Err(InvalidModelConfig("Invalid model architecture")),
             }
         }
-        None => return Err(Error::msg("Model architecture not found")),
+        None => return Err(InvalidModelConfig("Model architecture not found")),
     };
 
     Ok(config)
@@ -149,13 +149,15 @@ pub(crate) fn encode_batch_with_usage<'s, E>(
 where
     E: Into<EncodeInput<'s>> + Send,
 {
-    let tokens = tokenizer
-        .encode_batch(sentences, true)
-        .map_err(Error::msg)
-        .context("Failed to encode batch.")?;
+    let tokens = tokenizer.encode_batch(sentences, true)?;
 
     let prompt_tokens = tokens.len() as u32;
 
+    let usage = Usage {
+        prompt_tokens,
+        total_tokens: prompt_tokens,
+    };
+
     let token_ids = tokens
         .iter()
         .map(|tokens| {
@@ -170,20 +172,16 @@ where
     let embeddings = model.encode(&token_ids)?;
     tracing::trace!("generated embeddings {:?}", embeddings.shape());
 
-    // Apply some avg-pooling by taking the mean model value for all tokens (including padding)
-    let (_n_sentence, out_tokens, _hidden_size) = embeddings.dims3()?;
+    // Apply pooling
     let pooled_embeddings = pool_embeddings(&embeddings, pooling_strategy)?;
+
+    // Normalize embeddings (if required)
     let embeddings = if normalize {
         normalize_l2(&pooled_embeddings)?
     } else {
         pooled_embeddings
     };
 
-    // TODO: Incorrect usage calculation - fix
-    let usage = Usage {
-        prompt_tokens,
-        total_tokens: prompt_tokens + (out_tokens as u32),
-    };
     Ok((embeddings, usage))
 }
 

crates/glowrs/src/model/pooling.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use crate::Result;
 use candle_core::Tensor;
 use serde::Deserialize;
 
@@ -12,9 +12,9 @@ pub enum PoolingStrategy {
 
 pub fn pool_embeddings(embeddings: &Tensor, strategy: &PoolingStrategy) -> Result<Tensor> {
     match strategy {
-        PoolingStrategy::Mean => Ok(mean_pooling(embeddings)?),
-        PoolingStrategy::Max => Ok(max_pooling(embeddings)?),
-        PoolingStrategy::Sum => Ok(sum_pooling(embeddings)?),
+        PoolingStrategy::Mean => mean_pooling(embeddings),
+        PoolingStrategy::Max => max_pooling(embeddings),
+        PoolingStrategy::Sum => sum_pooling(embeddings),
     }
 }
 

crates/glowrs/src/model/sentence_transformer.rs

@@ -1,4 +1,3 @@
-use anyhow::{Context, Error, Result};
 use candle_core::Tensor;
 use hf_hub::api::sync::{Api, ApiRepo};
 use hf_hub::{Repo, RepoType};
@@ -10,7 +9,7 @@ use crate::model::embedder::{
     encode_batch, encode_batch_with_usage, load_pretrained_model, EmbedderModel,
 };
 use crate::model::utils;
-use crate::Usage;
+use crate::{Error, Result, Usage};
 
 #[cfg(test)]
 use crate::model::embedder::{load_zeros_model, parse_config};
@@ -57,8 +56,9 @@ impl SentenceTransformer {
     ///
     /// ```rust
     /// # use glowrs::SentenceTransformer;
+    /// # use std::error::Error;
     ///
-    /// # fn main() -> anyhow::Result<()> {
+    /// # fn main() -> Result<(), Box<dyn Error>> {
     /// let encoder = SentenceTransformer::from_repo_string("sentence-transformers/all-MiniLM-L6-v2")?;
     ///
     /// # Ok(())
@@ -84,26 +84,19 @@ impl SentenceTransformer {
     }
 
     pub fn from_api(api: ApiRepo) -> Result<Self> {
-        let model_path = api
-            .get("model.safetensors")
-            .context("Model repository is not available or doesn't contain `model.safetensors`.")?;
+        let model_path = api.get("model.safetensors")?;
 
-        let config_path = api
-            .get("config.json")
-            .context("Model repository doesn't contain `config.json`.")?;
+        let config_path = api.get("config.json")?;
 
-        let tokenizer_path = api
-            .get("tokenizer.json")
-            .context("Model repository doesn't contain `tokenizer.json`.")?;
+        let tokenizer_path = api.get("tokenizer.json")?;
 
         Self::from_path(&model_path, &config_path, &tokenizer_path)
     }
 
     pub fn from_path(model_path: &Path, config_path: &Path, tokenizer_path: &Path) -> Result<Self> {
-        let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(Error::msg)?;
+        let tokenizer = Tokenizer::from_file(tokenizer_path)?;
 
-        let model = load_pretrained_model(model_path, config_path)
-            .context("Something went wrong while loading the model.")?;
+        let model = load_pretrained_model(model_path, config_path)?;
 
         Ok(Self::new(model, tokenizer))
     }
@@ -119,7 +112,9 @@ impl SentenceTransformer {
     /// use glowrs::SentenceTransformer;
     /// use std::path::Path;
     ///
-    /// # fn main() -> anyhow::Result<()> {
+    /// # type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
+    ///
+    /// # fn main() -> Result<()> {
     /// let path = Path::new("path/to/folder");
     ///
     /// let encoder = SentenceTransformer::from_folder(path)?;
@@ -133,7 +128,9 @@ impl SentenceTransformer {
         let tokenizer_path = folder_path.join("tokenizer.json");
 
         if !model_path.exists() || !config_path.exists() || !tokenizer_path.exists() {
-            Err(anyhow::anyhow!("model.safetensors, config.json, or tokenizer.json does not exist in the given directory"))
+            Err(Error::ModelLoad(
+                "model.safetensors, config.json, or tokenizer.json does not exist in the given directory"
+            ))
         } else {
             Self::from_path(&model_path, &config_path, &tokenizer_path)
         }
@@ -147,7 +144,9 @@ impl SentenceTransformer {
     /// # use glowrs::SentenceTransformer;
     /// # use glowrs::PoolingStrategy;
     ///
-    /// # fn main() -> anyhow::Result<()> {
+    /// # type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
+    ///
+    /// # fn main() -> Result<()> {
     /// let encoder = SentenceTransformer::from_repo_string("sentence-transformers/all-MiniLM-L6-v2")?
     ///    .with_pooling_strategy(PoolingStrategy::Sum);
     ///
@@ -161,7 +160,7 @@ impl SentenceTransformer {
 
     #[cfg(test)]
     pub(crate) fn test_from_config_json(config_path: &Path, tokenizer_path: &Path) -> Result<Self> {
-        let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(Error::msg)?;
+        let tokenizer = Tokenizer::from_file(tokenizer_path)?;
 
         let config_str = std::fs::read_to_string(config_path)?;