fix: use st max_seq_length (#167)

Author: OlivierDehaene

core/src/download.rs

@@ -2,6 +2,17 @@ use hf_hub::api::tokio::{ApiError, ApiRepo};
 use std::path::PathBuf;
 use tracing::instrument;
 
+// Old classes used other config names than 'sentence_bert_config.json'
+pub const ST_CONFIG_NAMES: [&str; 7] = [
+    "sentence_bert_config.json",
+    "sentence_roberta_config.json",
+    "sentence_distilbert_config.json",
+    "sentence_camembert_config.json",
+    "sentence_albert_config.json",
+    "sentence_xlm-roberta_config.json",
+    "sentence_xlnet_config.json",
+];
+
 #[instrument(skip_all)]
 pub async fn download_artifacts(api: &ApiRepo) -> Result<PathBuf, ApiError> {
     let start = std::time::Instant::now();
@@ -32,3 +43,20 @@ pub async fn download_pool_config(api: &ApiRepo) -> Result<PathBuf, ApiError> {
     let pool_config_path = api.get("1_Pooling/config.json").await?;
     Ok(pool_config_path)
 }
+
+#[instrument(skip_all)]
+pub async fn download_st_config(api: &ApiRepo) -> Result<PathBuf, ApiError> {
+    // Try default path
+    let err = match api.get(ST_CONFIG_NAMES[0]).await {
+        Ok(st_config_path) => return Ok(st_config_path),
+        Err(err) => err,
+    };
+
+    for name in &ST_CONFIG_NAMES[1..] {
+        if let Ok(st_config_path) = api.get(name).await {
+            return Ok(st_config_path);
+        }
+    }
+
+    Err(err)
+}

router/src/http/server.rs

@@ -9,10 +9,10 @@ use crate::{
     shutdown, ClassifierModel, EmbeddingModel, ErrorResponse, ErrorType, Info, ModelType,
     ResponseMetadata,
 };
+use ::http::HeaderMap;
 use anyhow::Context;
 use axum::extract::Extension;
 use axum::http::HeaderValue;
-use ::http::HeaderMap;
 use axum::http::{Method, StatusCode};
 use axum::routing::{get, post};
 use axum::{http, Json, Router};

router/src/lib.rs

@@ -4,11 +4,13 @@ mod prometheus;
 
 #[cfg(feature = "http")]
 mod http;
+
 #[cfg(feature = "http")]
 use ::http::HeaderMap;
 
 #[cfg(feature = "grpc")]
 mod grpc;
+
 #[cfg(feature = "grpc")]
 use tonic::codegen::http::HeaderMap;
 
@@ -25,14 +27,14 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::path::Path;
 use std::time::{Duration, Instant};
 use text_embeddings_backend::DType;
-use text_embeddings_core::download::{download_artifacts, download_pool_config};
+use text_embeddings_core::download::{
+    download_artifacts, download_pool_config, download_st_config, ST_CONFIG_NAMES,
+};
 use text_embeddings_core::infer::Infer;
 use text_embeddings_core::queue::Queue;
 use text_embeddings_core::tokenization::Tokenization;
 use text_embeddings_core::TextEmbeddingsError;
-use tokenizers::decoders::metaspace::PrependScheme;
-use tokenizers::pre_tokenizers::sequence::Sequence;
-use tokenizers::{PreTokenizerWrapper, Tokenizer};
+use tokenizers::Tokenizer;
 use tracing::Span;
 
 pub use logging::init_logging;
@@ -83,6 +85,9 @@ pub async fn run(
             let _ = download_pool_config(&api_repo).await;
         }
 
+        // Download sentence transformers config
+        let _ = download_st_config(&api_repo).await;
+
         // Download model from the Hub
         download_artifacts(&api_repo)
             .await
@@ -178,7 +183,25 @@ pub async fn run(
     } else {
         0
     };
-    let max_input_length = config.max_position_embeddings - position_offset;
+
+    // Try to load ST Config
+    let mut st_config: Option<STConfig> = None;
+    for name in ST_CONFIG_NAMES {
+        let config_path = model_root.join(name);
+        if let Ok(config) = fs::read_to_string(config_path) {
+            st_config =
+                Some(serde_json::from_str(&config).context(format!("Failed to parse `{}`", name))?);
+            break;
+        }
+    }
+    let max_input_length = match st_config {
+        Some(config) => config.max_seq_length,
+        None => {
+            tracing::warn!("Could not find a Sentence Transformers config");
+            config.max_position_embeddings - position_offset
+        }
+    };
+    tracing::info!("Maximum number of tokens per request: {max_input_length}");
 
     let tokenization_workers = tokenization_workers.unwrap_or_else(num_cpus::get_physical);
 
@@ -311,6 +334,11 @@ pub struct PoolConfig {
     pooling_mode_mean_sqrt_len_tokens: bool,
 }
 
+#[derive(Debug, Deserialize)]
+pub struct STConfig {
+    pub max_seq_length: usize,
+}
+
 #[derive(Clone, Debug, Serialize)]
 #[cfg_attr(feature = "http", derive(utoipa::ToSchema))]
 pub struct EmbeddingModel {

router/src/main.rs

@@ -1,9 +1,9 @@
 use anyhow::Result;
 use clap::Parser;
+use mimalloc::MiMalloc;
 use opentelemetry::global;
 use text_embeddings_backend::DType;
 use veil::Redact;
-use mimalloc::MiMalloc;
 
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;