Merge pull request huggingface#5 from huggingface/roberta

feat: add support for XLM-RoBERTa

Author: OlivierDehaene

README.md

@@ -48,21 +48,25 @@ Benchmark for [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1
 
 ### Supported Models
 
-You can use any BERT model with absolute positions in `text-embeddings-inference`. If the model does not have `safetensors` weights
-you can convert it using [this space](https://huggingface.co/spaces/safetensors/convert).
+You can use any BERT or XLM-RoBERTa model with absolute positions in `text-embeddings-inference`. 
+If the model does not have `safetensors` weights you can convert it using [this space](https://huggingface.co/spaces/safetensors/convert).
 
 **Support for other model types will be added in the future.**
 
-| MTEB Rank | Model Type | Model ID               | Specific Revision                                                        |
-|-----------|------------|------------------------|--------------------------------------------------------------------------|
-| 1         | Bert       | BAAI/bge-large-en-v1.5 | [refs/pr/5](https://huggingface.co/BAAI/bge-large-en-v1.5/discussions/5) |
-| 2         |            | BAAI/bge-base-en-v1.5  | [refs/pr/1](https://huggingface.co/BAAI/bge-base-en-v1.5/discussions/1)  |
-| 3         |            | llmrails/ember-v1      |                                                                          |
-| 4         |            | thenlper/gte-large     |                                                                          |
-| 5         |            | thenlper/gte-base      |                                                                          |
-| 6         |            | intfloat/e5-large-v2   |                                                                          |
-| 7         |            | BAAI/bge-small-en-v1.5 | [refs/pr/3](https://huggingface.co/BAAI/bge-small-en-v1.5/discussions/3) |
-| 10        |            | intfloat/e5-base-v2    |                                                                          |
+Examples of supported models:
+
+| MTEB Rank | Model Type   | Model ID                       | Specific Revision                                                        |
+|-----------|--------------|--------------------------------|--------------------------------------------------------------------------|
+| 1         | Bert         | BAAI/bge-large-en-v1.5         | [refs/pr/5](https://huggingface.co/BAAI/bge-large-en-v1.5/discussions/5) |
+| 2         |              | BAAI/bge-base-en-v1.5          | [refs/pr/1](https://huggingface.co/BAAI/bge-base-en-v1.5/discussions/1)  |
+| 3         |              | llmrails/ember-v1              |                                                                          |
+| 4         |              | thenlper/gte-large             |                                                                          |
+| 5         |              | thenlper/gte-base              |                                                                          |
+| 6         |              | intfloat/e5-large-v2           |                                                                          |
+| 7         |              | BAAI/bge-small-en-v1.5         | [refs/pr/3](https://huggingface.co/BAAI/bge-small-en-v1.5/discussions/3) |
+| 10        |              | intfloat/e5-base-v2            |                                                                          |
+| 11        | XLM-RoBERTa  | intfloat/multilingual-e5-large |                                                                          |
+
 
 You can explore the list of best performing text embeddings models [here](https://huggingface.co/spaces/mteb/leaderboard).
 

backends/candle/src/lib.rs

@@ -39,7 +39,9 @@ impl CandleBackend {
         };
 
         // Check model type
-        if config.model_type != Some("bert".to_string()) {
+        if config.model_type != Some("bert".to_string())
+            && config.model_type != Some("xlm-roberta".to_string())
+        {
             return Err(BackendError::Start(format!(
                 "Model {:?} is not supported",
                 config.model_type

core/src/tokenization.rs

@@ -13,7 +13,12 @@ pub struct Tokenization {
 }
 
 impl Tokenization {
-    pub fn new(workers: usize, tokenizer: Tokenizer, max_input_length: usize) -> Self {
+    pub fn new(
+        workers: usize,
+        tokenizer: Tokenizer,
+        max_input_length: usize,
+        position_offset: usize,
+    ) -> Self {
         // Create channel
         let (sender, receiver) = flume::unbounded();
 
@@ -24,7 +29,12 @@ impl Tokenization {
 
             // Spawn worker
             tokio::task::spawn_blocking(move || {
-                tokenizer_worker(tokenizer_clone, max_input_length, receiver_clone)
+                tokenizer_worker(
+                    tokenizer_clone,
+                    max_input_length,
+                    position_offset,
+                    receiver_clone,
+                )
             });
         }
 
@@ -66,6 +76,7 @@ impl Tokenization {
 fn tokenizer_worker(
     tokenizer: Tokenizer,
     max_input_length: usize,
+    position_offset: usize,
     receiver: flume::Receiver<TokenizerRequest>,
 ) {
     // Loop over requests
@@ -74,8 +85,13 @@ fn tokenizer_worker(
             if !response_tx.is_closed() {
                 // It's possible that the user dropped its request resulting in a send error.
                 // We just discard the error
-                let _ =
-                    response_tx.send(encode_input(inputs, truncate, max_input_length, &tokenizer));
+                let _ = response_tx.send(encode_input(
+                    inputs,
+                    truncate,
+                    max_input_length,
+                    position_offset,
+                    &tokenizer,
+                ));
             }
         })
     }
@@ -86,6 +102,7 @@ fn encode_input(
     inputs: String,
     truncate: bool,
     max_input_length: usize,
+    position_offset: usize,
     tokenizer: &Tokenizer,
 ) -> Result<Encoding, TextEmbeddingsError> {
     // Get the number of tokens in the input
@@ -109,7 +126,8 @@ fn encode_input(
     Ok(Encoding {
         input_ids: encoding.get_ids().to_vec(),
         token_type_ids: encoding.get_type_ids().to_vec(),
-        position_ids: (0..seq_len as u32).collect::<Vec<_>>(),
+        position_ids: (position_offset as u32..(seq_len + position_offset) as u32)
+            .collect::<Vec<_>>(),
     })
 }
 

router/src/main.rs

@@ -116,6 +116,7 @@ pub struct ModelConfig {
     pub model_type: String,
     #[serde(alias = "n_positions")]
     pub max_position_embeddings: usize,
+    pub pad_token_id: usize,
 }
 
 #[tokio::main]
@@ -167,11 +168,19 @@ async fn main() -> Result<()> {
     );
     tokenizer.with_padding(None);
 
+    // Position IDs offset. Used for Roberta.
+    let position_offset = if config.pad_token_id == 0 {
+        0
+    } else {
+        config.pad_token_id + 1
+    };
+
     // Tokenization logic
     let tokenization = Tokenization::new(
         args.tokenization_workers,
         tokenizer,
         config.max_position_embeddings,
+        position_offset,
     );
 
     // Create backend