Generalization of input types for pipelines

benches/squad_benchmark.rs

@@ -78,7 +78,7 @@ fn bench_squad(c: &mut Criterion) {
     }
     //    Define input
     let mut squad_path = PathBuf::from(env::var("squad_dataset")
-.expect("Please set the \"squad_dataset\" environment variable pointing to the SQuAD dataset folder"));
+        .expect("Please set the \"squad_dataset\" environment variable pointing to the SQuAD dataset folder"));
     squad_path.push("dev-v2.0.json");
     let mut qa_inputs = squad_processor(squad_path);
     qa_inputs.truncate(1000);

examples/albert.rs

@@ -52,8 +52,7 @@ fn main() -> anyhow::Result<()> {
         "Looks like one [MASK] is missing",
         "It was a very nice and [MASK] day",
     ];
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(&input, 128, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

examples/bart.rs

@@ -54,8 +54,7 @@ fn main() -> anyhow::Result<()> {
 
     //    Credits: WikiNews, CC BY 2.5 license (https://en.wikinews.org/wiki/Astronomers_find_water_vapour_in_atmosphere_of_exoplanet_K2-18b)
 
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 1024, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(&input, 1024, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

examples/bert.rs

@@ -47,8 +47,7 @@ fn main() -> anyhow::Result<()> {
         "Looks like one [MASK] is missing",
         "It was a very nice and [MASK] day",
     ];
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(&input, 128, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

examples/distilbert_masked_lm.rs

@@ -50,8 +50,7 @@ fn main() -> anyhow::Result<()> {
         "Looks like one thing is missing",
         "It\'s like comparing oranges to apples",
     ];
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(input, 128, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

examples/electra_discriminator.rs

@@ -51,7 +51,7 @@ fn main() -> anyhow::Result<()> {
     let input = ["One Two Three Ten Five Six Seven Eight"];
     let tokenized_input = MultiThreadedTokenizer::encode_list(
         &tokenizer,
-        input.to_vec(),
+        &input,
         128,
         &TruncationStrategy::LongestFirst,
         0,

examples/electra_masked_lm.rs

@@ -51,8 +51,7 @@ fn main() -> anyhow::Result<()> {
         "Looks like one [MASK] is missing",
         "It was a very nice and [MASK] day",
     ];
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(&input, 128, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

examples/gpt2.rs

@@ -51,8 +51,7 @@ fn main() -> anyhow::Result<()> {
 
     //    Define input
     let input = ["One two three four five six seven eight nine ten eleven"];
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(&input, 128, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

examples/openai_gpt.rs

@@ -56,8 +56,7 @@ fn main() -> anyhow::Result<()> {
 
     //    Define input
     let input = ["Wondering what the next word will"];
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(&input, 128, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

examples/roberta.rs

@@ -60,8 +60,7 @@ fn main() -> anyhow::Result<()> {
         "<pad> Looks like one thing is missing",
         "It\'s like comparing oranges to apples",
     ];
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(&input, 128, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

examples/xlnet.rs

@@ -49,8 +49,7 @@ fn main() -> anyhow::Result<()> {
 
     //    Define input
     let input = ["One two three four"];
-    let tokenized_input =
-        tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
+    let tokenized_input = tokenizer.encode_list(&input, 128, &TruncationStrategy::LongestFirst, 0);
     let max_len = tokenized_input
         .iter()
         .map(|input| input.token_ids.len())

src/pipelines/common.rs

@@ -32,7 +32,7 @@ use rust_tokenizers::tokenizer::{
 use rust_tokenizers::vocab::{
     AlbertVocab, BertVocab, MarianVocab, RobertaVocab, T5Vocab, XLMRobertaVocab, XLNetVocab,
 };
-use rust_tokenizers::{Mask, Offset, OffsetSize, TokenizedInput};
+use rust_tokenizers::{TokenIdsWithOffsets, TokenizedInput};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::path::Path;
@@ -273,7 +273,7 @@ impl TokenizerOption {
     /// Interface method
     pub fn encode_list(
         &self,
-        text_list: Vec<&str>,
+        text_list: &[&str],
         max_len: usize,
         truncation_strategy: &TruncationStrategy,
         stride: usize,
@@ -330,7 +330,7 @@ impl TokenizerOption {
     /// Interface method for pair encoding
     pub fn encode_pair_list(
         &self,
-        text_pair_list: Vec<(&str, &str)>,
+        text_pair_list: &[(&str, &str)],
         max_len: usize,
         truncation_strategy: &TruncationStrategy,
         stride: usize,
@@ -400,110 +400,61 @@ impl TokenizerOption {
     /// Interface method to build input with special tokens
     pub fn build_input_with_special_tokens(
         &self,
-        tokens_1: Vec<i64>,
-        tokens_2: Option<Vec<i64>>,
-        offsets_1: Vec<Option<Offset>>,
-        offsets_2: Option<Vec<Option<Offset>>>,
-        original_offsets_1: Vec<Vec<OffsetSize>>,
-        original_offsets_2: Option<Vec<Vec<OffsetSize>>>,
-        mask_1: Vec<Mask>,
-        mask_2: Option<Vec<Mask>>,
+        token_ids_with_offsets_1: TokenIdsWithOffsets,
+        token_ids_with_offsets_2: Option<TokenIdsWithOffsets>,
     ) -> TokenizedInput {
-        let (token_ids, segment_ids, special_tokens_mask, token_offsets, reference_offsets, mask) =
-            match *self {
-                Self::Bert(ref tokenizer) => tokenizer.build_input_with_special_tokens(
-                    tokens_1,
-                    tokens_2,
-                    offsets_1,
-                    offsets_2,
-                    original_offsets_1,
-                    original_offsets_2,
-                    mask_1,
-                    mask_2,
-                ),
-                Self::Roberta(ref tokenizer) => tokenizer.build_input_with_special_tokens(
-                    tokens_1,
-                    tokens_2,
-                    offsets_1,
-                    offsets_2,
-                    original_offsets_1,
-                    original_offsets_2,
-                    mask_1,
-                    mask_2,
-                ),
-                Self::XLMRoberta(ref tokenizer) => tokenizer.build_input_with_special_tokens(
-                    tokens_1,
-                    tokens_2,
-                    offsets_1,
-                    offsets_2,
-                    original_offsets_1,
-                    original_offsets_2,
-                    mask_1,
-                    mask_2,
-                ),
-                Self::Marian(ref tokenizer) => tokenizer.build_input_with_special_tokens(
-                    tokens_1,
-                    tokens_2,
-                    offsets_1,
-                    offsets_2,
-                    original_offsets_1,
-                    original_offsets_2,
-                    mask_1,
-                    mask_2,
-                ),
-                Self::T5(ref tokenizer) => tokenizer.build_input_with_special_tokens(
-                    tokens_1,
-                    tokens_2,
-                    offsets_1,
-                    offsets_2,
-                    original_offsets_1,
-                    original_offsets_2,
-                    mask_1,
-                    mask_2,
-                ),
-                Self::Albert(ref tokenizer) => tokenizer.build_input_with_special_tokens(
-                    tokens_1,
-                    tokens_2,
-                    offsets_1,
-                    offsets_2,
-                    original_offsets_1,
-                    original_offsets_2,
-                    mask_1,
-                    mask_2,
-                ),
-                Self::XLNet(ref tokenizer) => tokenizer.build_input_with_special_tokens(
-                    tokens_1,
-                    tokens_2,
-                    offsets_1,
-                    offsets_2,
-                    original_offsets_1,
-                    original_offsets_2,
-                    mask_1,
-                    mask_2,
-                ),
-            };
+        let token_ids_with_special_tokens = match *self {
+            Self::Bert(ref tokenizer) => tokenizer.build_input_with_special_tokens(
+                token_ids_with_offsets_1,
+                token_ids_with_offsets_2,
+            ),
+            Self::Roberta(ref tokenizer) => tokenizer.build_input_with_special_tokens(
+                token_ids_with_offsets_1,
+                token_ids_with_offsets_2,
+            ),
+            Self::XLMRoberta(ref tokenizer) => tokenizer.build_input_with_special_tokens(
+                token_ids_with_offsets_1,
+                token_ids_with_offsets_2,
+            ),
+            Self::Marian(ref tokenizer) => tokenizer.build_input_with_special_tokens(
+                token_ids_with_offsets_1,
+                token_ids_with_offsets_2,
+            ),
+            Self::T5(ref tokenizer) => tokenizer.build_input_with_special_tokens(
+                token_ids_with_offsets_1,
+                token_ids_with_offsets_2,
+            ),
+            Self::Albert(ref tokenizer) => tokenizer.build_input_with_special_tokens(
+                token_ids_with_offsets_1,
+                token_ids_with_offsets_2,
+            ),
+            Self::XLNet(ref tokenizer) => tokenizer.build_input_with_special_tokens(
+                token_ids_with_offsets_1,
+                token_ids_with_offsets_2,
+            ),
+        };
         TokenizedInput {
-            token_ids,
-            segment_ids,
-            special_tokens_mask,
+            token_ids: token_ids_with_special_tokens.token_ids,
+            segment_ids: token_ids_with_special_tokens.segment_ids,
+            special_tokens_mask: token_ids_with_special_tokens.special_tokens_mask,
             overflowing_tokens: vec![],
             num_truncated_tokens: 0,
-            token_offsets,
-            reference_offsets,
-            mask,
+            token_offsets: token_ids_with_special_tokens.token_offsets,
+            reference_offsets: token_ids_with_special_tokens.reference_offsets,
+            mask: token_ids_with_special_tokens.mask,
         }
     }
 
     /// Interface method to convert tokens to ids
     pub fn convert_tokens_to_ids(&self, tokens: &[String]) -> Vec<i64> {
         match *self {
-            Self::Bert(ref tokenizer) => tokenizer.convert_tokens_to_ids(&tokens.into()),
-            Self::Roberta(ref tokenizer) => tokenizer.convert_tokens_to_ids(&tokens.into()),
-            Self::Marian(ref tokenizer) => tokenizer.convert_tokens_to_ids(&tokens.into()),
-            Self::T5(ref tokenizer) => tokenizer.convert_tokens_to_ids(&tokens.into()),
-            Self::XLMRoberta(ref tokenizer) => tokenizer.convert_tokens_to_ids(&tokens.into()),
-            Self::Albert(ref tokenizer) => tokenizer.convert_tokens_to_ids(&tokens.into()),
-            Self::XLNet(ref tokenizer) => tokenizer.convert_tokens_to_ids(&tokens.into()),
+            Self::Bert(ref tokenizer) => tokenizer.convert_tokens_to_ids(tokens.into()),
+            Self::Roberta(ref tokenizer) => tokenizer.convert_tokens_to_ids(tokens.into()),
+            Self::Marian(ref tokenizer) => tokenizer.convert_tokens_to_ids(tokens.into()),
+            Self::T5(ref tokenizer) => tokenizer.convert_tokens_to_ids(tokens.into()),
+            Self::XLMRoberta(ref tokenizer) => tokenizer.convert_tokens_to_ids(tokens.into()),
+            Self::Albert(ref tokenizer) => tokenizer.convert_tokens_to_ids(tokens.into()),
+            Self::XLNet(ref tokenizer) => tokenizer.convert_tokens_to_ids(tokens.into()),
         }
     }
 

src/pipelines/conversation.rs

@@ -671,7 +671,7 @@ impl ConversationModel {
                 .map(|c| &c.history)
                 .collect_vec();
 
-            let prompt_ids = self.encode_prompts(texts.as_slice());
+            let prompt_ids = self.encode_prompts(texts.as_ref());
             let input_tensor = self.concat_input_history(prompt_ids, history);
             let input_length = *input_tensor.size().last().unwrap() as usize;
             let mut generated = self.model.generate_from_ids_and_past(input_tensor, None);
@@ -791,7 +791,7 @@ impl ConversationModel {
 
     fn encode_prompts(&self, texts: &[&str]) -> Vec<Vec<i64>> {
         // Encode the user prompt into token ids
-        let tokens = self.model.get_tokenizer().tokenize_list(texts.to_vec());
+        let tokens = self.model.get_tokenizer().tokenize_list(texts);
 
         tokens
             .into_iter()