fix: add_pooling_layer for bert classification (#190)

OlivierDehaene · web-flow · commit 2b8ad5f4f12c · 2024-03-07T11:15:10.000+01:00
diff --git a/backends/candle/src/models/bert.rs b/backends/candle/src/models/bert.rs
@@ -365,6 +365,7 @@ pub trait ClassificationHead {
 }
 
 pub struct BertClassificationHead {
+    pooler: Option<Linear>,
     output: Linear,
     span: tracing::Span,
 }
@@ -376,11 +377,24 @@ impl BertClassificationHead {
             Some(id2label) => id2label.len(),
         };
 
-        let output_weight = vb.get((n_classes, config.hidden_size), "weight")?;
-        let output_bias = vb.get(n_classes, "bias")?;
+        let pooler = if let Ok(pooler_weight) = vb
+            .pp("bert.pooler.dense")
+            .get((config.hidden_size, config.hidden_size), "weight")
+        {
+            let pooler_bias = vb.pp("bert.pooler.dense").get(config.hidden_size, "bias")?;
+            Some(Linear::new(pooler_weight, Some(pooler_bias), None))
+        } else {
+            None
+        };
+
+        let output_weight = vb
+            .pp("classifier")
+            .get((n_classes, config.hidden_size), "weight")?;
+        let output_bias = vb.pp("classifier").get(n_classes, "bias")?;
         let output = Linear::new(output_weight, Some(output_bias), None);
 
         Ok(Self {
+            pooler,
             output,
             span: tracing::span!(tracing::Level::TRACE, "classifier"),
         })
@@ -390,7 +404,14 @@ impl BertClassificationHead {
 impl ClassificationHead for BertClassificationHead {
     fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
         let _enter = self.span.enter();
-        let hidden_states = self.output.forward(hidden_states)?;
+
+        let mut hidden_states = hidden_states.clone();
+        if let Some(pooler) = self.pooler.as_ref() {
+            hidden_states = pooler.forward(&hidden_states)?;
+            hidden_states = hidden_states.tanh()?;
+        }
+
+        let hidden_states = self.output.forward(&hidden_states)?;
         Ok(hidden_states)
     }
 }
@@ -551,7 +572,7 @@ impl BertModel {
                 let pool = Pool::Cls;
 
                 let classifier: Box<dyn ClassificationHead + Send> =
-                    Box::new(BertClassificationHead::load(vb.pp("classifier"), config)?);
+                    Box::new(BertClassificationHead::load(vb.clone(), config)?);
                 (pool, Some(classifier), None)
             }
             ModelType::Embedding(pool) => {
diff --git a/backends/candle/src/models/flash_bert.rs b/backends/candle/src/models/flash_bert.rs
@@ -246,7 +246,7 @@ impl FlashBertModel {
                 let pool = Pool::Cls;
 
                 let classifier: Box<dyn ClassificationHead + Send> =
-                    Box::new(BertClassificationHead::load(vb.pp("classifier"), config)?);
+                    Box::new(BertClassificationHead::load(vb.clone(), config)?);
                 (pool, Some(classifier), None)
             }
             ModelType::Embedding(pool) => {
diff --git a/backends/candle/tests/common.rs b/backends/candle/tests/common.rs
@@ -65,11 +65,22 @@ pub fn sort_embeddings(embeddings: Embeddings) -> (Vec<Vec<f32>>, Vec<Vec<f32>>)
     (pooled_embeddings, raw_embeddings)
 }
 
-pub fn download_artifacts(model_id: &'static str) -> Result<PathBuf> {
+pub fn download_artifacts(
+    model_id: &'static str,
+    revision: Option<&'static str>,
+) -> Result<PathBuf> {
     let builder = ApiBuilder::new().with_progress(false);
 
     let api = builder.build().unwrap();
-    let api_repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model));
+    let api_repo = if let Some(revision) = revision {
+        api.repo(Repo::with_revision(
+            model_id.to_string(),
+            RepoType::Model,
+            revision.to_string(),
+        ))
+    } else {
+        api.repo(Repo::new(model_id.to_string(), RepoType::Model))
+    };
 
     api_repo.get("config.json")?;
     api_repo.get("tokenizer.json")?;
diff --git a/backends/candle/tests/snapshots/test_bert__bert_classification_single.snap b/backends/candle/tests/snapshots/test_bert__bert_classification_single.snap
@@ -0,0 +1,7 @@
+---
+source: backends/candle/tests/test_bert.rs
+assertion_line: 211
+expression: predictions_single
+---
+- - 2.8580017
+  - -2.9722357
diff --git a/backends/candle/tests/snapshots/test_flash_bert__bert_classification_single.snap b/backends/candle/tests/snapshots/test_flash_bert__bert_classification_single.snap
@@ -0,0 +1,6 @@
+---
+source: backends/candle/tests/test_flash_bert.rs
+expression: predictions_single
+---
+- - 2.8574219
+  - -2.9726563
diff --git a/backends/candle/tests/test_bert.rs b/backends/candle/tests/test_bert.rs
@@ -9,7 +9,7 @@ use text_embeddings_backend_core::{Backend, ModelType, Pool};
 #[test]
 #[serial_test::serial]
 fn test_mini() -> Result<()> {
-    let model_root = download_artifacts("sentence-transformers/all-MiniLM-L6-v2")?;
+    let model_root = download_artifacts("sentence-transformers/all-MiniLM-L6-v2", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
@@ -69,7 +69,7 @@ fn test_mini() -> Result<()> {
 #[test]
 #[serial_test::serial]
 fn test_mini_pooled_raw() -> Result<()> {
-    let model_root = download_artifacts("sentence-transformers/all-MiniLM-L6-v2")?;
+    let model_root = download_artifacts("sentence-transformers/all-MiniLM-L6-v2", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
@@ -139,7 +139,7 @@ fn test_mini_pooled_raw() -> Result<()> {
 #[test]
 #[serial_test::serial]
 fn test_emotions() -> Result<()> {
-    let model_root = download_artifacts("SamLowe/roberta-base-go_emotions")?;
+    let model_root = download_artifacts("SamLowe/roberta-base-go_emotions", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(model_root, "float32".to_string(), ModelType::Classifier)?;
@@ -185,3 +185,38 @@ fn test_emotions() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[serial_test::serial]
+fn test_bert_classification() -> Result<()> {
+    let model_root = download_artifacts("ibm/re2g-reranker-nq", Some("refs/pr/3"))?;
+    let tokenizer = load_tokenizer(&model_root)?;
+
+    let backend = CandleBackend::new(model_root, "float32".to_string(), ModelType::Classifier)?;
+
+    let input_single = batch(
+        vec![tokenizer
+            .encode(
+                (
+                    "PrimeTime is a timing signoff tool",
+                    "PrimeTime can perform most accurate timing analysis",
+                ),
+                true,
+            )
+            .unwrap()],
+        [0].to_vec(),
+        vec![],
+    );
+
+    let predictions: Vec<Vec<f32>> = backend
+        .predict(input_single)?
+        .into_iter()
+        .map(|(_, v)| v)
+        .collect();
+    let predictions_single = SnapshotScores::from(predictions);
+
+    let matcher = relative_matcher();
+    insta::assert_yaml_snapshot!("bert_classification_single", predictions_single, &matcher);
+
+    Ok(())
+}
diff --git a/backends/candle/tests/test_flash_bert.rs b/backends/candle/tests/test_flash_bert.rs
@@ -15,7 +15,7 @@ use text_embeddings_backend_core::{Backend, ModelType, Pool};
     any(feature = "flash-attn", feature = "flash-attn-v1")
 ))]
 fn test_flash_mini() -> Result<()> {
-    let model_root = download_artifacts("sentence-transformers/all-MiniLM-L6-v2")?;
+    let model_root = download_artifacts("sentence-transformers/all-MiniLM-L6-v2", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
@@ -79,7 +79,7 @@ fn test_flash_mini() -> Result<()> {
     any(feature = "flash-attn", feature = "flash-attn-v1")
 ))]
 fn test_flash_mini_pooled_raw() -> Result<()> {
-    let model_root = download_artifacts("sentence-transformers/all-MiniLM-L6-v2")?;
+    let model_root = download_artifacts("sentence-transformers/all-MiniLM-L6-v2", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
@@ -153,7 +153,7 @@ fn test_flash_mini_pooled_raw() -> Result<()> {
     any(feature = "flash-attn", feature = "flash-attn-v1")
 ))]
 fn test_flash_emotions() -> Result<()> {
-    let model_root = download_artifacts("SamLowe/roberta-base-go_emotions")?;
+    let model_root = download_artifacts("SamLowe/roberta-base-go_emotions", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(model_root, "float16".to_string(), ModelType::Classifier)?;
@@ -199,3 +199,42 @@ fn test_flash_emotions() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[serial_test::serial]
+#[cfg(all(
+    feature = "cuda",
+    any(feature = "flash-attn", feature = "flash-attn-v1")
+))]
+fn test_flash_bert_classification() -> Result<()> {
+    let model_root = download_artifacts("ibm/re2g-reranker-nq", Some("refs/pr/3"))?;
+    let tokenizer = load_tokenizer(&model_root)?;
+
+    let backend = CandleBackend::new(model_root, "float16".to_string(), ModelType::Classifier)?;
+
+    let input_single = batch(
+        vec![tokenizer
+            .encode(
+                (
+                    "PrimeTime is a timing signoff tool",
+                    "PrimeTime can perform most accurate timing analysis",
+                ),
+                true,
+            )
+            .unwrap()],
+        [0].to_vec(),
+        vec![],
+    );
+
+    let predictions: Vec<Vec<f32>> = backend
+        .predict(input_single)?
+        .into_iter()
+        .map(|(_, v)| v)
+        .collect();
+    let predictions_single = SnapshotScores::from(predictions);
+
+    let matcher = relative_matcher();
+    insta::assert_yaml_snapshot!("bert_classification_single", predictions_single, &matcher);
+
+    Ok(())
+}
diff --git a/backends/candle/tests/test_flash_jina.rs b/backends/candle/tests/test_flash_jina.rs
@@ -11,7 +11,7 @@ use text_embeddings_backend_core::{Backend, ModelType, Pool};
 #[serial_test::serial]
 #[cfg(all(feature = "cuda", feature = "flash-attn"))]
 fn test_flash_jina_small() -> Result<()> {
-    let model_root = download_artifacts("jinaai/jina-embeddings-v2-small-en")?;
+    let model_root = download_artifacts("jinaai/jina-embeddings-v2-small-en", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
diff --git a/backends/candle/tests/test_flash_nomic.rs b/backends/candle/tests/test_flash_nomic.rs
@@ -11,7 +11,7 @@ use text_embeddings_backend_core::{Backend, ModelType, Pool};
 #[serial_test::serial]
 #[cfg(all(feature = "cuda", feature = "flash-attn"))]
 fn test_flash_nomic_small() -> Result<()> {
-    let model_root = download_artifacts("nomic-ai/nomic-embed-text-v1.5")?;
+    let model_root = download_artifacts("nomic-ai/nomic-embed-text-v1.5", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
diff --git a/backends/candle/tests/test_jina.rs b/backends/candle/tests/test_jina.rs
@@ -8,7 +8,7 @@ use text_embeddings_backend_core::{Backend, ModelType, Pool};
 
 #[test]
 fn test_jina_small() -> Result<()> {
-    let model_root = download_artifacts("jinaai/jina-embeddings-v2-small-en")?;
+    let model_root = download_artifacts("jinaai/jina-embeddings-v2-small-en", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
diff --git a/backends/candle/tests/test_nomic.rs b/backends/candle/tests/test_nomic.rs
@@ -8,7 +8,7 @@ use text_embeddings_backend_core::{Backend, ModelType, Pool};
 
 #[test]
 fn test_nomic_small() -> Result<()> {
-    let model_root = download_artifacts("nomic-ai/nomic-embed-text-v1.5")?;
+    let model_root = download_artifacts("nomic-ai/nomic-embed-text-v1.5", None)?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
diff --git a/backends/python/src/logging.rs b/backends/python/src/logging.rs
@@ -52,7 +52,7 @@ impl TryFrom<&String> for PythonLogMessage {
 }
 
 pub(crate) fn log_lines<S: Sized + BufRead>(lines: Lines<S>) {
-    for line in lines.flatten() {
+    for line in lines.map_while(Result::ok) {
         match PythonLogMessage::try_from(&line) {
             Ok(log) => log.trace(),
             Err(_) => tracing::debug!("{line}"),
diff --git a/backends/python/src/management.rs b/backends/python/src/management.rs
@@ -89,7 +89,7 @@ impl BackendProcess {
                 // We read stderr in another thread as it seems that lines() can block in some cases
                 let (err_sender, err_receiver) = mpsc::channel();
                 thread::spawn(move || {
-                    for line in stderr_reader.lines().flatten() {
+                    for line in stderr_reader.lines().map_while(Result::ok) {
                         err_sender.send(line).unwrap_or(());
                     }
                 });

Original file line number	Diff line number	Diff line change
`@@ -246,7 +246,7 @@ impl FlashBertModel {`
`246`	`246`	`let pool = Pool::Cls;`
`247`	`247`
`248`	`248`	`let classifier: Box<dyn ClassificationHead + Send> =`
`249`		`- Box::new(BertClassificationHead::load(vb.pp("classifier"), config)?);`
	`249`	`+ Box::new(BertClassificationHead::load(vb.clone(), config)?);`
`250`	`250`	`(pool, Some(classifier), None)`
`251`	`251`	`}`
`252`	`252`	`ModelType::Embedding(pool) => {`