Merge pull request #205 from arnaudbriche/fix-tokenizer-config-with-addedd-token

fix: support AddedToken objects in tokenizer config special tokens

Author: akshayballal95

rust/src/embeddings/local/bert.rs

@@ -30,11 +30,29 @@ pub trait BertEmbed {
         late_chunking: Option<bool>,
     ) -> Result<Vec<EmbeddingResult>, anyhow::Error>;
 }
+// HuggingFace tokenizer configs represent special tokens either as a plain
+// string or as an AddedToken object. Both forms must be accepted.
+#[derive(Debug, Deserialize, Clone)]
+#[serde(untagged)]
+pub enum SpecialToken {
+    String(String),
+    Added { content: String },
+}
+
+impl SpecialToken {
+    pub fn content(&self) -> &str {
+        match self {
+            SpecialToken::String(s) => s,
+            SpecialToken::Added { content } => content,
+        }
+    }
+}
+
 #[derive(Debug, Deserialize, Clone)]
 pub struct TokenizerConfig {
     pub max_length: Option<usize>,
     pub model_max_length: Option<usize>,
-    pub mask_token: Option<String>,
+    pub mask_token: Option<SpecialToken>,
     pub added_tokens_decoder: Option<HashMap<String, AddedToken>>,
 }
 

rust/src/embeddings/local/colbert.rs

@@ -100,9 +100,12 @@ impl OrtColbertEmbedder {
         };
 
         let mut tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
-        let mask_token = tokenizer_config.clone().mask_token;
-        let pad_id = match mask_token.clone() {
-            Some(mask_token) => tokenizer_config.get_token_id_from_token(&mask_token),
+        let mask_token: Option<String> = tokenizer_config
+            .mask_token
+            .as_ref()
+            .map(|t| t.content().to_owned());
+        let pad_id = match mask_token.as_deref() {
+            Some(mask_token) => tokenizer_config.get_token_id_from_token(mask_token),
             None => None,
         };
         let document_marker_token_id = tokenizer_config.get_token_id_from_token("[DocumentMarker]");