Fix QA pipelines for tokenizers not providing segment ids

CHANGELOG.md

@@ -10,6 +10,7 @@ All notable changes to this project will be documented in this file. The format
 - (BREAKING) Fixed the keyword extraction pipeline for n-gram sizes > 2. Add new configuration option `tokenizer_forbidden_ngram_chars` to specify characters that should be excluded from n-grams (allows filtering m-grams spanning multiple sentences).
 - Improved MPS device compatibility setting the `sparse_grad` flag to false for `gather` operations
 - Updated ONNX runtime backend version to 1.15.x
+- Issue with incorrect results for QA models with a tokenizer not using segment ids
 
 ## [0.21.0] - 2023-06-03
 ## Added

src/models/deberta_v2/encoder.rs

@@ -235,7 +235,7 @@ impl DebertaV2Encoder {
                         .unsqueeze(-1)
                         .to_kind(Kind::Uint8)
             }
-            value if value == 3 => attention_mask.unsqueeze(1),
+            3 => attention_mask.unsqueeze(1),
             _ => attention_mask.shallow_clone(),
         }
     }

src/models/mbart/mbart_model.rs

@@ -450,7 +450,7 @@ impl MBartForConditionalGeneration {
     {
         let p = p.borrow();
 
-        let base_model = MBartModel::new(p.borrow() / "model", config);
+        let base_model = MBartModel::new(p / "model", config);
         let final_logits_bias = p.var(
             "final_logits_bias",
             &[1, config.vocab_size],

src/pipelines/question_answering.rs

@@ -64,14 +64,14 @@ use std::cmp::min;
 use std::collections::HashMap;
 use std::fs;
 use std::path::PathBuf;
-use tch::kind::Kind::Float;
 use tch::nn::VarStore;
 use tch::{no_grad, Device, Kind, Tensor};
 
 use crate::deberta_v2::DebertaV2ForQuestionAnswering;
 #[cfg(feature = "onnx")]
 use crate::pipelines::onnx::{config::ONNXEnvironmentConfig, ONNXEncoder};
 
+use crate::common::kind::get_min;
 #[cfg(feature = "remote")]
 use crate::{
     distilbert::{DistilBertConfigResources, DistilBertModelResources, DistilBertVocabResources},
@@ -830,11 +830,15 @@ impl QuestionAnsweringModel {
                             .to_device(start_logits.device())
                             .eq(0);
 
-                        let start = start_logits.get(feature_idx).masked_fill(&p_mask, -10000);
-                        let end = end_logits.get(feature_idx).masked_fill(&p_mask, -10000);
+                        let start = start_logits
+                            .get(feature_idx)
+                            .masked_fill(&p_mask, get_min(start_logits.kind()).unwrap());
+                        let end = end_logits
+                            .get(feature_idx)
+                            .masked_fill(&p_mask, get_min(start_logits.kind()).unwrap());
 
-                        let start = start.exp() / start.exp().sum(Float);
-                        let end = end.exp() / end.exp().sum(Float);
+                        let start = start.softmax(0, start.kind());
+                        let end = end.softmax(0, end.kind());
 
                         let (starts, ends, scores) = self.decode(&start, &end, top_k);
 
@@ -861,9 +865,7 @@ impl QuestionAnsweringModel {
                         }
                     }
                     feature_id_start = max_feature_id;
-                    let example_answers = example_top_k_answers_map
-                        .entry(example_id)
-                        .or_insert_with(Vec::new);
+                    let example_answers = example_top_k_answers_map.entry(example_id).or_default();
                     example_answers.extend(answers);
                 }
             });
@@ -928,6 +930,20 @@ impl QuestionAnsweringModel {
             masks: encoded_query.masks,
         };
 
+        let sequence_added_tokens = self
+            .tokenizer
+            .build_input_with_special_tokens(
+                TokenIdsWithOffsets {
+                    ids: vec![],
+                    offsets: vec![],
+                    reference_offsets: vec![],
+                    masks: vec![],
+                },
+                None,
+            )
+            .token_ids
+            .len();
+
         let sequence_pair_added_tokens = self
             .tokenizer
             .build_input_with_special_tokens(
@@ -975,7 +991,10 @@ impl QuestionAnsweringModel {
             let encoded_span = self
                 .tokenizer
                 .build_input_with_special_tokens(encoded_query.clone(), Some(sub_encoded_context));
-            let p_mask = self.get_mask(&encoded_span);
+            let p_mask = self.get_mask(
+                &encoded_span,
+                encoded_query.ids.len() + sequence_added_tokens,
+            );
             let qa_feature = QaFeature {
                 input_ids: encoded_span.token_ids,
                 offsets: encoded_span.token_offsets,
@@ -1038,7 +1057,7 @@ impl QuestionAnsweringModel {
         (input_ids, attention_masks, token_type_ids)
     }
 
-    fn get_mask(&self, encoded_span: &TokenizedInput) -> Vec<i8> {
+    fn get_mask(&self, encoded_span: &TokenizedInput, question_length: usize) -> Vec<i8> {
         let sep_indices: Vec<usize> = encoded_span
             .token_ids
             .iter()
@@ -1047,12 +1066,9 @@ impl QuestionAnsweringModel {
             .map(|(position, _)| position)
             .collect();
 
-        let mut p_mask: Vec<i8> = encoded_span
-            .segment_ids
-            .iter()
-            .map(|v| min(v, &1i8))
-            .map(|&v| 1i8 - v)
-            .collect();
+        let mut p_mask: Vec<i8> = Vec::with_capacity(encoded_span.token_ids.len());
+        p_mask.extend(vec![1; question_length]);
+        p_mask.extend(vec![0; encoded_span.token_ids.len() - question_length]);
         for sep_position in sep_indices {
             p_mask[sep_position] = 1;
         }

src/pipelines/translation/translation_builder.rs

@@ -336,12 +336,10 @@ impl TranslationModelBuilder {
         ) {
             (Some(ModelType::M2M100), source_languages, target_languages) => {
                 match self.model_size {
-                    Some(value) if value == ModelSize::XLarge => {
-                        model_fetchers::get_m2m100_xlarge_resources(
-                            source_languages.as_ref(),
-                            target_languages.as_ref(),
-                        )?
-                    }
+                    Some(ModelSize::XLarge) => model_fetchers::get_m2m100_xlarge_resources(
+                        source_languages.as_ref(),
+                        target_languages.as_ref(),
+                    )?,
                     _ => model_fetchers::get_m2m100_large_resources(
                         source_languages.as_ref(),
                         target_languages.as_ref(),
@@ -447,7 +445,7 @@ mod model_fetchers {
         Ok(match get_marian_model(source_languages, target_languages) {
             Ok(marian_resources) => marian_resources,
             Err(_) => match model_size {
-                Some(value) if value == &ModelSize::XLarge => {
+                Some(ModelSize::XLarge) => {
                     get_m2m100_xlarge_resources(source_languages, target_languages)?
                 }
                 _ => get_m2m100_large_resources(source_languages, target_languages)?,