feat: optimize inference postprocessing

Author: besaleli

encoderfile/src/inference/embedding.rs

@@ -33,39 +33,26 @@ pub fn embedding<'a>(
             false => embs.into_owned(),
         };
 
-        let mut token_ids = encoding.get_ids().iter();
-        let mut tokens = encoding.get_tokens().iter();
-        let mut special_tokens_mask = encoding.get_special_tokens_mask().iter();
-        let mut offsets = encoding.get_offsets().iter();
-        let mut embeddings_iter = transformed.axis_iter(Axis(0));
-
         let mut results = Vec::new();
 
-        while let (Some(token_id), Some(token), Some(special_tokens_mask), Some(offset), Some(e)) = (
-            token_ids.next(),
-            tokens.next(),
-            special_tokens_mask.next(),
-            offsets.next(),
-            embeddings_iter.next(),
-        ) {
-            if *special_tokens_mask == 1 {
+        for i in 0..encoding.len() {
+            if encoding.get_special_tokens_mask()[i] == 1 {
                 continue;
             }
 
-            let (start, end) = *offset;
-            let embedding: Vec<f32> = e.iter().map(|i| *i).collect();
-
+            let (start, end) = encoding.get_offsets()[i];
             let token_info = TokenInfo {
-                token: token.clone(),
-                token_id: *token_id,
+                token: encoding.get_tokens()[i].clone(),
+                token_id: encoding.get_ids()[i],
                 start,
                 end,
             };
 
+            let e = transformed.index_axis(Axis(0), i);
             results.push(TokenEmbedding {
-                embedding,
+                embedding: e.to_owned().into_raw_vec_and_offset().0,
                 token_info: Some(token_info),
-            })
+            });
         }
 
         embeddings.push(TokenEmbeddingSequence {
@@ -74,6 +61,4 @@ pub fn embedding<'a>(
     }
 
     Ok(embeddings)
-
-    // Err(ApiError::InternalError("Not Implemented"))
 }

encoderfile/src/inference/sequence_classification.rs

@@ -1,6 +1,7 @@
 use crate::{common::SequenceClassificationResult, config::ModelConfig, error::ApiError};
 use ndarray::{Axis, Ix2};
 use ndarray_stats::QuantileExt;
+use ort::tensor::ArrayExtensions;
 use tokenizers::Encoding;
 
 pub fn sequence_classification<'a>(
@@ -23,16 +24,16 @@ pub fn sequence_classification<'a>(
         .expect("Model does not return tensor of shape [n_batch, n_labels]")
         .into_owned();
 
-    let probabilities = super::utils::softmax(&outputs, Axis(1));
+    let probabilities = outputs.softmax(Axis(1));
 
     let results = outputs
         .axis_iter(Axis(0))
         .zip(probabilities.axis_iter(Axis(0)))
         .map(|(logs, probs)| {
             let predicted_index = probs.argmax().expect("Model has 0 labels");
             SequenceClassificationResult {
-                logits: logs.iter().map(|i| *i).collect(),
-                scores: probs.iter().map(|i| *i).collect(),
+                logits: logs.to_owned().into_raw_vec_and_offset().0,
+                scores: probs.to_owned().into_raw_vec_and_offset().0,
                 predicted_index: (predicted_index as u32),
                 predicted_label: config
                     .id2label(predicted_index as u32)

encoderfile/src/inference/token_classification.rs

@@ -2,10 +2,10 @@ use crate::{
     common::{TokenClassification, TokenClassificationResult, TokenInfo},
     config::ModelConfig,
     error::ApiError,
-    inference::utils::softmax,
 };
 use ndarray::{Axis, Ix3};
 use ndarray_stats::QuantileExt;
+use ort::tensor::ArrayExtensions;
 use tokenizers::Encoding;
 
 pub fn token_classification<'a>(
@@ -29,35 +29,16 @@ pub fn token_classification<'a>(
 
     for (encoding, logits) in encodings.iter().zip(outputs.axis_iter(Axis(0))) {
         let logits = logits.to_owned();
-
-        let scores = softmax(&logits, Axis(1));
-
-        let mut token_ids = encoding.get_ids().iter();
-        let mut tokens = encoding.get_tokens().iter();
-        let mut special_tokens_mask = encoding.get_special_tokens_mask().iter();
-        let mut offsets = encoding.get_offsets().iter();
-        let mut logs_iter = logits.axis_iter(Axis(0));
-        let mut scores_iter = scores.axis_iter(Axis(0));
+        let scores = logits.softmax(Axis(1));
 
         let mut results = Vec::new();
 
-        while let (
-            Some(token_id),
-            Some(token),
-            Some(special_tokens_mask),
-            Some(offset),
-            Some(logs),
-            Some(scores),
-        ) = (
-            token_ids.next(),
-            tokens.next(),
-            special_tokens_mask.next(),
-            offsets.next(),
-            logs_iter.next(),
-            scores_iter.next(),
-        ) {
-            let argmax = scores.argmax().expect("Model has 0 labels");
-            let score = scores[argmax];
+        for i in 0..encoding.len() {
+            let argmax = scores
+                .index_axis(Axis(0), i)
+                .argmax()
+                .expect("Model has 0 labels");
+            let score = scores.index_axis(Axis(0), i)[argmax];
             let label = match config.id2label(argmax as u32) {
                 Some(l) => l.to_string(),
                 None => {
@@ -67,24 +48,31 @@ pub fn token_classification<'a>(
                     )
                 }
             };
+            let (start, end) = encoding.get_offsets()[i];
 
-            let (start, end) = *offset;
-
-            if *special_tokens_mask == 1 {
+            if encoding.get_special_tokens_mask()[i] == 1 {
                 continue;
             }
 
             results.push(TokenClassification {
                 token_info: TokenInfo {
-                    token_id: *token_id,
-                    token: token.clone(),
+                    token_id: encoding.get_ids()[i],
+                    token: encoding.get_tokens()[i].clone(),
                     start,
                     end,
                 },
                 score: score,
                 label,
-                logits: logs.iter().map(|i| *i).collect(),
-                scores: scores.iter().map(|i| *i).collect(),
+                logits: logits
+                    .index_axis(Axis(0), i)
+                    .to_owned()
+                    .into_raw_vec_and_offset()
+                    .0,
+                scores: scores
+                    .index_axis(Axis(0), i)
+                    .to_owned()
+                    .into_raw_vec_and_offset()
+                    .0,
             })
         }