Change default precision to BF16; fix the bug when input is batched (#4)

Signed-off-by: kaixuanliu <[email protected]>

backends/python/server/text_embeddings_server/models/default_model.py

@@ -51,7 +51,7 @@ def embed(self, batch: PaddedBatch) -> List[Embedding]:
 
         output = self.model(**kwargs)
         embedding = output[0][:, 0]
-        cpu_results = embedding.view(-1).tolist()
+        cpu_results = embedding.reshape(-1).tolist()
 
         return [
             Embedding(

backends/src/dtype.rs

@@ -12,6 +12,7 @@ pub enum DType {
         all(feature = "candle", not(feature = "accelerate"))
     ))]
     Float16,
+    BFloat16,
     // Float32 is not available on candle cuda
     #[cfg(any(feature = "python", feature = "candle"))]
     Float32,
@@ -28,6 +29,7 @@ impl fmt::Display for DType {
                 all(feature = "candle", not(feature = "accelerate"))
             ))]
             DType::Float16 => write!(f, "float16"),
+            DType::BFloat16 => write!(f, "bfloat16"),
             // Float32 is not available on candle cuda
             #[cfg(any(feature = "python", feature = "candle"))]
             DType::Float32 => write!(f, "float32"),

router/src/lib.rs

@@ -183,7 +183,7 @@ pub async fn run(
         }
         #[cfg(not(any(feature = "accelerate", feature = "mkl", feature = "mkl-dynamic")))]
         {
-            DType::Float16
+            DType::BFloat16
         }
     });