Implement gpt2 (BPE) GGUF tokenizer conversion

mistralrs-core/src/pipeline/gguf_tokenizer.rs

@@ -1,10 +1,10 @@
-use std::sync::atomic::Ordering;
+use std::{collections::HashMap, sync::atomic::Ordering};
 
 use anyhow::Result;
 use candle_core::quantized::gguf_file::Content;
 use tokenizers::{
     decoders::{self, byte_fallback::ByteFallback, fuse::Fuse, strip::Strip},
-    models::unigram::Unigram,
+    models::{bpe::BpeBuilder, unigram::Unigram},
     normalizers::{self, Prepend, Replace},
     AddedToken, DecoderWrapper, ModelWrapper, NormalizerWrapper, Tokenizer,
 };
@@ -77,7 +77,7 @@
 
     let bos_str = tokens[bos as usize].clone();
     let eos_str = tokens[eos as usize].clone();
-    let unk_str;
+    let mut unk_str = None;
 
     let (tokenizer, ty) = match model.as_str() {
         "llama" | "replit" => {
@@ -92,7 +92,7 @@
 
             // Unigram (sentencepiece) default UNK is 0
             let unk = unk.map(|x| x as usize).unwrap_or(0);
-            unk_str = tokens[unk].clone();
+            unk_str = Some(tokens[unk].clone());
 
             let unigram = Unigram::from(vocab, Some(unk), true).map_err(anyhow::Error::msg)?;
             let mut tokenizer = Tokenizer::new(ModelWrapper::Unigram(unigram));
@@ -113,6 +113,34 @@
 
             (tokenizer, "unigram")
         }
+        "gpt2" => {
+            // This is a `bpe` tokenizer
+            let merges = merges
+                .as_ref()
+                .expect("Expect `tokenizer.ggml.merges` for `llama` unigram tokeizer.")
+                .into_iter()
+                .map(|merges| {
+                    let res = merges.splitn(2, ' ').collect::<Vec<_>>();
+                    (res[0].to_string(), res[1].to_string())
+                })
+                .collect::<Vec<_>>();
+            let mut vocab = HashMap::new();
+            for (i, token) in tokens.iter().enumerate() {
+                vocab.insert(token.clone(), i as u32);
+            }
+
+            let bpe = BpeBuilder::new()
+                .vocab_and_merges(vocab, merges)
+                .build()
+                .map_err(anyhow::Error::msg)?;
+            let mut tokenizer = Tokenizer::new(ModelWrapper::BPE(bpe));
+            tokenizer.with_decoder(decoders::byte_level::ByteLevel::new(true, true, true));
+
+            tokenizer.add_special_tokens(&[AddedToken::from(tokens[bos as usize].clone(), true)]);
+            tokenizer.add_special_tokens(&[AddedToken::from(tokens[eos as usize].clone(), true)]);
+
+            (tokenizer, "bpe")
+        }
         other => {
             anyhow::bail!("Tokenizer model `{other}` not supported.");
         }
@@ -132,7 +160,7 @@
         tokenizer,
         bos: Some(bos_str),
         eos: Some(eos_str),
-        unk: Some(unk_str),
+        unk: unk_str,
     })
 }