rust_bert/pipelines/
question_answering.rs

1// Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
2// Copyright 2019 Guillaume Becquin
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//     http://www.apache.org/licenses/LICENSE-2.0
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS,
9// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10// See the License for the specific language governing permissions and
11// limitations under the License.
12
13//! # Question Answering pipeline
14//! Extractive question answering from a given question and context. By default, the dependencies for this
15//! model will be downloaded for a DistilBERT model finetuned on SQuAD (Stanford Question Answering Dataset).
16//! Customized DistilBERT models can be loaded by overwriting the resources in the configuration.
17//! The dependencies will be downloaded to the user's home directory, under ~/.cache/.rustbert/distilbert-qa
18//!
19//! ```no_run
20//! use rust_bert::pipelines::question_answering::{QaInput, QuestionAnsweringModel};
21//!
22//! # fn main() -> anyhow::Result<()> {
23//! let qa_model = QuestionAnsweringModel::new(Default::default())?;
24//!
25//! let question = String::from("Where does Amy live ?");
26//! let context = String::from("Amy lives in Amsterdam");
27//!
28//! let answers = qa_model.predict(&vec![QaInput { question, context }], 1, 32);
29//! # Ok(())
30//! # }
31//! ```
32//!
33//! Output: \
34//! ```no_run
35//! # use rust_bert::pipelines::question_answering::Answer;
36//! # let output =
37//! [Answer {
38//!     score: 0.9976,
39//!     start: 13,
40//!     end: 21,
41//!     answer: String::from("Amsterdam"),
42//! }]
43//! # ;
44//! ```
45
46use crate::albert::AlbertForQuestionAnswering;
47use crate::bert::BertForQuestionAnswering;
48use crate::common::error::RustBertError;
49use crate::deberta::DebertaForQuestionAnswering;
50use crate::distilbert::DistilBertForQuestionAnswering;
51use crate::fnet::FNetForQuestionAnswering;
52use crate::longformer::LongformerForQuestionAnswering;
53use crate::mobilebert::MobileBertForQuestionAnswering;
54use crate::pipelines::common::{
55    cast_var_store, get_device, ConfigOption, ModelResource, ModelType, TokenizerOption,
56};
57use crate::reformer::ReformerForQuestionAnswering;
58use crate::resources::ResourceProvider;
59use crate::roberta::RobertaForQuestionAnswering;
60use crate::xlnet::XLNetForQuestionAnswering;
61use rust_tokenizers::{Offset, TokenIdsWithOffsets, TokenizedInput};
62use serde::{Deserialize, Serialize};
63use std::cmp::min;
64use std::collections::HashMap;
65use std::fs;
66use std::path::PathBuf;
67use tch::nn::VarStore;
68use tch::{no_grad, Device, Kind, Tensor};
69
70use crate::deberta_v2::DebertaV2ForQuestionAnswering;
71#[cfg(feature = "onnx")]
72use crate::pipelines::onnx::{config::ONNXEnvironmentConfig, ONNXEncoder};
73
74use crate::common::kind::get_min;
75#[cfg(feature = "remote")]
76use crate::{
77    distilbert::{DistilBertConfigResources, DistilBertModelResources, DistilBertVocabResources},
78    resources::RemoteResource,
79};
80
81#[derive(Serialize, Deserialize)]
82/// # Input for Question Answering
83/// Includes a context (containing the answer) and question strings
84pub struct QaInput {
85    /// Question string
86    pub question: String,
87    /// Context or query
88    pub context: String,
89}
90
91#[derive(Debug)]
92struct QaFeature {
93    pub input_ids: Vec<i64>,
94    pub offsets: Vec<Option<Offset>>,
95    pub token_type_ids: Vec<i8>,
96    pub p_mask: Vec<i8>,
97    pub example_index: i64,
98}
99
100#[derive(Debug, Clone, Serialize, Deserialize)]
101/// # Output for Question Answering
102pub struct Answer {
103    /// Confidence score
104    pub score: f64,
105    /// Start position of answer span
106    pub start: usize,
107    /// End position of answer span
108    pub end: usize,
109    /// Answer span
110    pub answer: String,
111}
112
113impl PartialEq for Answer {
114    fn eq(&self, other: &Self) -> bool {
115        (self.start == other.start) && (self.end == other.end) && (self.answer == other.answer)
116    }
117}
118
119fn remove_duplicates<T: PartialEq + Clone>(vector: &mut Vec<T>) -> &mut Vec<T> {
120    let mut potential_duplicates = vec![];
121    vector.retain(|item| {
122        if potential_duplicates.contains(item) {
123            false
124        } else {
125            potential_duplicates.push(item.clone());
126            true
127        }
128    });
129    vector
130}
131
132/// # Configuration for question answering
133/// Contains information regarding the model to load and device to place the model on.
134pub struct QuestionAnsweringConfig {
135    /// Model weights resource (default: pretrained DistilBERT model on SQuAD)
136    pub model_resource: ModelResource,
137    /// Config resource (default: pretrained DistilBERT model on SQuAD)
138    pub config_resource: Box<dyn ResourceProvider + Send>,
139    /// Vocab resource (default: pretrained DistilBERT model on SQuAD)
140    pub vocab_resource: Box<dyn ResourceProvider + Send>,
141    /// Merges resource (default: None)
142    pub merges_resource: Option<Box<dyn ResourceProvider + Send>>,
143    /// Device to place the model on (default: CUDA/GPU when available)
144    pub device: Device,
145    /// Model type
146    pub model_type: ModelType,
147    /// Flag indicating if the model expects a lower casing of the input
148    pub lower_case: bool,
149    /// Flag indicating if the tokenizer should strip accents (normalization). Only used for BERT / ALBERT models
150    pub strip_accents: Option<bool>,
151    /// Flag indicating if the tokenizer should add a white space before each tokenized input (needed for some Roberta models)
152    pub add_prefix_space: Option<bool>,
153    /// Maximum sequence length for the combined query and context
154    pub max_seq_length: usize,
155    /// Stride to apply if the context needs to be broken down due to a large length. Represents the number of overlapping tokens between sliding windows.
156    pub doc_stride: usize,
157    /// Maximum length for the query
158    pub max_query_length: usize,
159    /// Maximum length for the answer
160    pub max_answer_length: usize,
161    /// Model weights precision. If not provided, will default to full precision on CPU, or the loaded weights precision otherwise
162    pub kind: Option<Kind>,
163}
164
165impl QuestionAnsweringConfig {
166    /// Instantiate a new question answering configuration of the supplied type.
167    ///
168    /// # Arguments
169    ///
170    /// * `model_type` - `ModelType` indicating the model type to load (must match with the actual data to be loaded!)
171    /// * model_resource - The `ResourceProvider` pointing to the model to load (e.g.  model.ot)
172    /// * config_resource - The `ResourceProvider` pointing to the model configuration to load (e.g. config.json)
173    /// * vocab_resource - The `ResourceProvider` pointing to the tokenizer's vocabulary to load (e.g.  vocab.txt/vocab.json)
174    /// * merges_resource - An optional `ResourceProvider` pointing to the tokenizer's merge file to load (e.g.  merges.txt), needed only for Roberta.
175    /// * lower_case - A `bool` indicating whether the tokenizer should lower case all input (in case of a lower-cased model)
176    pub fn new<RC, RV>(
177        model_type: ModelType,
178        model_resource: ModelResource,
179        config_resource: RC,
180        vocab_resource: RV,
181        merges_resource: Option<RV>,
182        lower_case: bool,
183        strip_accents: impl Into<Option<bool>>,
184        add_prefix_space: impl Into<Option<bool>>,
185    ) -> QuestionAnsweringConfig
186    where
187        RC: ResourceProvider + Send + 'static,
188        RV: ResourceProvider + Send + 'static,
189    {
190        QuestionAnsweringConfig {
191            model_type,
192            model_resource,
193            config_resource: Box::new(config_resource),
194            vocab_resource: Box::new(vocab_resource),
195            merges_resource: merges_resource.map(|r| Box::new(r) as Box<_>),
196            lower_case,
197            strip_accents: strip_accents.into(),
198            add_prefix_space: add_prefix_space.into(),
199            device: Device::cuda_if_available(),
200            max_seq_length: 384,
201            doc_stride: 128,
202            max_query_length: 64,
203            max_answer_length: 15,
204            kind: None,
205        }
206    }
207
208    /// Instantiate a new question answering configuration of the supplied type.
209    ///
210    /// # Arguments
211    ///
212    /// * `model_type` - `ModelType` indicating the model type to load (must match with the actual data to be loaded!)
213    /// * model_resource - The `ResourceProvider` pointing to the model to load (e.g.  model.ot)
214    /// * config_resource - The `ResourceProvider` pointing to the model configuration to load (e.g. config.json)
215    /// * vocab_resource - The `ResourceProvider` pointing to the tokenizer's vocabulary to load (e.g.  vocab.txt/vocab.json)
216    /// * merges_resource - An optional `ResourceProvider` pointing to the tokenizer's merge file to load (e.g.  merges.txt), needed only for Roberta.
217    /// * lower_case - A `bool` indicating whether the tokenizer should lower case all input (in case of a lower-cased model)
218    /// * max_seq_length - Optional maximum sequence token length to limit memory footprint. If the context is too long, it will be processed with sliding windows. Defaults to 384.
219    /// * max_query_length - Optional maximum question token length. Defaults to 64.
220    /// * doc_stride - Optional stride to apply if a sliding window is required to process the input context. Represents the number of overlapping tokens between sliding windows. This should be lower than the max_seq_length minus max_query_length (otherwise there is a risk for the sliding window not to progress). Defaults to 128.
221    /// * max_answer_length - Optional maximum token length for the extracted answer. Defaults to 15.
222    pub fn custom_new<RC, RV>(
223        model_type: ModelType,
224        model_resource: ModelResource,
225        config_resource: RC,
226        vocab_resource: RV,
227        merges_resource: Option<RV>,
228        lower_case: bool,
229        strip_accents: impl Into<Option<bool>>,
230        add_prefix_space: impl Into<Option<bool>>,
231        max_seq_length: impl Into<Option<usize>>,
232        doc_stride: impl Into<Option<usize>>,
233        max_query_length: impl Into<Option<usize>>,
234        max_answer_length: impl Into<Option<usize>>,
235    ) -> QuestionAnsweringConfig
236    where
237        RC: ResourceProvider + Send + 'static,
238        RV: ResourceProvider + Send + 'static,
239    {
240        QuestionAnsweringConfig {
241            model_type,
242            model_resource,
243            config_resource: Box::new(config_resource),
244            vocab_resource: Box::new(vocab_resource),
245            merges_resource: merges_resource.map(|r| Box::new(r) as Box<_>),
246            lower_case,
247            strip_accents: strip_accents.into(),
248            add_prefix_space: add_prefix_space.into(),
249            device: Device::cuda_if_available(),
250            max_seq_length: max_seq_length.into().unwrap_or(384),
251            doc_stride: doc_stride.into().unwrap_or(128),
252            max_query_length: max_query_length.into().unwrap_or(64),
253            max_answer_length: max_answer_length.into().unwrap_or(15),
254            kind: None,
255        }
256    }
257}
258
259#[cfg(feature = "remote")]
260impl Default for QuestionAnsweringConfig {
261    fn default() -> QuestionAnsweringConfig {
262        QuestionAnsweringConfig {
263            model_resource: ModelResource::Torch(Box::new(RemoteResource::from_pretrained(
264                DistilBertModelResources::DISTIL_BERT_SQUAD,
265            ))),
266            config_resource: Box::new(RemoteResource::from_pretrained(
267                DistilBertConfigResources::DISTIL_BERT_SQUAD,
268            )),
269            vocab_resource: Box::new(RemoteResource::from_pretrained(
270                DistilBertVocabResources::DISTIL_BERT_SQUAD,
271            )),
272            merges_resource: None,
273            device: Device::cuda_if_available(),
274            kind: None,
275            model_type: ModelType::DistilBert,
276            lower_case: false,
277            add_prefix_space: None,
278            strip_accents: None,
279            max_seq_length: 384,
280            doc_stride: 128,
281            max_query_length: 64,
282            max_answer_length: 15,
283        }
284    }
285}
286
287#[allow(clippy::large_enum_variant)]
288/// # Abstraction that holds one particular question answering model, for any of the supported models
289pub enum QuestionAnsweringOption {
290    /// Bert for Question Answering
291    Bert(BertForQuestionAnswering),
292    /// DeBERTa for Question Answering
293    Deberta(DebertaForQuestionAnswering),
294    /// DeBERTa V2 for Question Answering
295    DebertaV2(DebertaV2ForQuestionAnswering),
296    /// DistilBert for Question Answering
297    DistilBert(DistilBertForQuestionAnswering),
298    /// MobileBert for Question Answering
299    MobileBert(MobileBertForQuestionAnswering),
300    /// Roberta for Question Answering
301    Roberta(RobertaForQuestionAnswering),
302    /// XLMRoberta for Question Answering
303    XLMRoberta(RobertaForQuestionAnswering),
304    /// Albert for Question Answering
305    Albert(AlbertForQuestionAnswering),
306    /// XLNet for Question Answering
307    XLNet(XLNetForQuestionAnswering),
308    /// Reformer for Question Answering
309    Reformer(ReformerForQuestionAnswering),
310    /// Longformer for Question Answering
311    Longformer(LongformerForQuestionAnswering),
312    /// FNet for Question Answering
313    FNet(FNetForQuestionAnswering),
314    /// ONNX model for Question Answering
315    #[cfg(feature = "onnx")]
316    ONNX(ONNXEncoder),
317}
318
319impl QuestionAnsweringOption {
320    /// Instantiate a new question answering model of the supplied type.
321    ///
322    /// # Arguments
323    ///
324    /// * `QuestionAnsweringConfig` - Question answering pipeline configuration. The type of model created will be inferred from the
325    ///     `ModelResources` (Torch or ONNX) and `ModelType` (Architecture for Torch models) variants provided and
326    pub fn new(config: &QuestionAnsweringConfig) -> Result<Self, RustBertError> {
327        match config.model_resource {
328            ModelResource::Torch(_) => Self::new_torch(config),
329            #[cfg(feature = "onnx")]
330            ModelResource::ONNX(_) => Self::new_onnx(config),
331        }
332    }
333
334    fn new_torch(config: &QuestionAnsweringConfig) -> Result<Self, RustBertError> {
335        let device = config.device;
336        let weights_path = config.model_resource.get_torch_local_path()?;
337        let mut var_store = VarStore::new(device);
338        let model_config = &mut ConfigOption::from_file(
339            config.model_type,
340            config.config_resource.get_local_path()?,
341        );
342        let model_type = config.model_type;
343        let model = match model_type {
344            ModelType::Bert => {
345                if let ConfigOption::Bert(config) = model_config {
346                    Ok(QuestionAnsweringOption::Bert(
347                        BertForQuestionAnswering::new(var_store.root(), config),
348                    ))
349                } else {
350                    Err(RustBertError::InvalidConfigurationError(
351                        "You can only supply a BertConfig for Bert!".to_string(),
352                    ))
353                }
354            }
355            ModelType::Deberta => {
356                if let ConfigOption::Deberta(config) = model_config {
357                    Ok(QuestionAnsweringOption::Deberta(
358                        DebertaForQuestionAnswering::new(var_store.root(), config),
359                    ))
360                } else {
361                    Err(RustBertError::InvalidConfigurationError(
362                        "You can only supply a DebertaConfig for DeBERTa!".to_string(),
363                    ))
364                }
365            }
366            ModelType::DebertaV2 => {
367                if let ConfigOption::DebertaV2(config) = model_config {
368                    Ok(QuestionAnsweringOption::DebertaV2(
369                        DebertaV2ForQuestionAnswering::new(var_store.root(), config),
370                    ))
371                } else {
372                    Err(RustBertError::InvalidConfigurationError(
373                        "You can only supply a DebertaV2Config for DeBERTa V2!".to_string(),
374                    ))
375                }
376            }
377            ModelType::DistilBert => {
378                if let ConfigOption::DistilBert(ref mut config) = model_config {
379                    config.sinusoidal_pos_embds = false;
380                    Ok(QuestionAnsweringOption::DistilBert(
381                        DistilBertForQuestionAnswering::new(var_store.root(), config),
382                    ))
383                } else {
384                    Err(RustBertError::InvalidConfigurationError(
385                        "You can only supply a DistilBertConfig for DistilBert!".to_string(),
386                    ))
387                }
388            }
389            ModelType::MobileBert => {
390                if let ConfigOption::MobileBert(config) = model_config {
391                    Ok(QuestionAnsweringOption::MobileBert(
392                        MobileBertForQuestionAnswering::new(var_store.root(), config),
393                    ))
394                } else {
395                    Err(RustBertError::InvalidConfigurationError(
396                        "You can only supply a MobileBertConfig for MobileBert!".to_string(),
397                    ))
398                }
399            }
400            ModelType::Roberta => {
401                if let ConfigOption::Roberta(config) = model_config {
402                    Ok(QuestionAnsweringOption::Roberta(
403                        RobertaForQuestionAnswering::new(var_store.root(), config),
404                    ))
405                } else {
406                    Err(RustBertError::InvalidConfigurationError(
407                        "You can only supply a RobertaConfig for Roberta!".to_string(),
408                    ))
409                }
410            }
411            ModelType::XLMRoberta => {
412                if let ConfigOption::Bert(config) = model_config {
413                    Ok(QuestionAnsweringOption::XLMRoberta(
414                        RobertaForQuestionAnswering::new(var_store.root(), config),
415                    ))
416                } else {
417                    Err(RustBertError::InvalidConfigurationError(
418                        "You can only supply a BertConfig for Roberta!".to_string(),
419                    ))
420                }
421            }
422            ModelType::Albert => {
423                if let ConfigOption::Albert(config) = model_config {
424                    Ok(QuestionAnsweringOption::Albert(
425                        AlbertForQuestionAnswering::new(var_store.root(), config),
426                    ))
427                } else {
428                    Err(RustBertError::InvalidConfigurationError(
429                        "You can only supply an AlbertConfig for Albert!".to_string(),
430                    ))
431                }
432            }
433            ModelType::XLNet => {
434                if let ConfigOption::XLNet(config) = model_config {
435                    Ok(QuestionAnsweringOption::XLNet(
436                        XLNetForQuestionAnswering::new(var_store.root(), config)?,
437                    ))
438                } else {
439                    Err(RustBertError::InvalidConfigurationError(
440                        "You can only supply a XLNetConfig for XLNet!".to_string(),
441                    ))
442                }
443            }
444            ModelType::Reformer => {
445                if let ConfigOption::Reformer(config) = model_config {
446                    Ok(QuestionAnsweringOption::Reformer(
447                        ReformerForQuestionAnswering::new(var_store.root(), config)?,
448                    ))
449                } else {
450                    Err(RustBertError::InvalidConfigurationError(
451                        "You can only supply a ReformerConfig for Reformer!".to_string(),
452                    ))
453                }
454            }
455            ModelType::Longformer => {
456                if let ConfigOption::Longformer(config) = model_config {
457                    Ok(QuestionAnsweringOption::Longformer(
458                        LongformerForQuestionAnswering::new(var_store.root(), config),
459                    ))
460                } else {
461                    Err(RustBertError::InvalidConfigurationError(
462                        "You can only supply a LongformerConfig for Longformer!".to_string(),
463                    ))
464                }
465            }
466            ModelType::FNet => {
467                if let ConfigOption::FNet(config) = model_config {
468                    Ok(QuestionAnsweringOption::FNet(
469                        FNetForQuestionAnswering::new(var_store.root(), config),
470                    ))
471                } else {
472                    Err(RustBertError::InvalidConfigurationError(
473                        "You can only supply a FNetConfig for FNet!".to_string(),
474                    ))
475                }
476            }
477            _ => Err(RustBertError::InvalidConfigurationError(format!(
478                "QuestionAnswering not implemented for {model_type:?}!",
479            ))),
480        }?;
481        var_store.load(weights_path)?;
482        cast_var_store(&mut var_store, config.kind, device);
483        Ok(model)
484    }
485
486    #[cfg(feature = "onnx")]
487    pub fn new_onnx(config: &QuestionAnsweringConfig) -> Result<Self, RustBertError> {
488        let onnx_config = ONNXEnvironmentConfig::from_device(config.device);
489        let environment = onnx_config.get_environment()?;
490        let encoder_file = config
491            .model_resource
492            .get_onnx_local_paths()?
493            .encoder_path
494            .ok_or(RustBertError::InvalidConfigurationError(
495                "An encoder file must be provided for question answering ONNX models.".to_string(),
496            ))?;
497
498        Ok(Self::ONNX(ONNXEncoder::new(
499            encoder_file,
500            &environment,
501            &onnx_config,
502        )?))
503    }
504
505    /// Returns the `ModelType` for this SequenceClassificationOption
506    pub fn model_type(&self) -> ModelType {
507        match *self {
508            Self::Bert(_) => ModelType::Bert,
509            Self::Deberta(_) => ModelType::Deberta,
510            Self::DebertaV2(_) => ModelType::DebertaV2,
511            Self::Roberta(_) => ModelType::Roberta,
512            Self::XLMRoberta(_) => ModelType::XLMRoberta,
513            Self::DistilBert(_) => ModelType::DistilBert,
514            Self::MobileBert(_) => ModelType::MobileBert,
515            Self::Albert(_) => ModelType::Albert,
516            Self::XLNet(_) => ModelType::XLNet,
517            Self::Reformer(_) => ModelType::Reformer,
518            Self::Longformer(_) => ModelType::Longformer,
519            Self::FNet(_) => ModelType::FNet,
520            #[cfg(feature = "onnx")]
521            Self::ONNX(_) => ModelType::ONNX,
522        }
523    }
524
525    /// Interface method to forward_t() of the particular models.
526    pub fn forward_t(
527        &self,
528        input_ids: Option<&Tensor>,
529        mask: Option<&Tensor>,
530        input_embeds: Option<&Tensor>,
531        _token_type_ids: Option<&Tensor>,
532        train: bool,
533    ) -> (Tensor, Tensor) {
534        match *self {
535            Self::Bert(ref model) => {
536                let outputs = model.forward_t(input_ids, mask, None, None, input_embeds, train);
537                (outputs.start_logits, outputs.end_logits)
538            }
539            Self::Deberta(ref model) => {
540                let outputs = model
541                    .forward_t(input_ids, mask, None, None, input_embeds, train)
542                    .expect("Error in Deberta forward_t");
543                (outputs.start_logits, outputs.end_logits)
544            }
545            Self::DebertaV2(ref model) => {
546                let outputs = model
547                    .forward_t(input_ids, mask, None, None, input_embeds, train)
548                    .expect("Error in Deberta V2 forward_t");
549                (outputs.start_logits, outputs.end_logits)
550            }
551            Self::DistilBert(ref model) => {
552                let outputs = model
553                    .forward_t(input_ids, mask, input_embeds, train)
554                    .expect("Error in distilbert forward_t");
555                (outputs.start_logits, outputs.end_logits)
556            }
557            Self::MobileBert(ref model) => {
558                let outputs = model
559                    .forward_t(input_ids, None, None, input_embeds, mask, train)
560                    .expect("Error in mobilebert forward_t");
561                (outputs.start_logits, outputs.end_logits)
562            }
563            Self::Roberta(ref model) | Self::XLMRoberta(ref model) => {
564                let outputs = model.forward_t(input_ids, mask, None, None, input_embeds, train);
565                (outputs.start_logits, outputs.end_logits)
566            }
567            Self::Albert(ref model) => {
568                let outputs = model.forward_t(input_ids, mask, None, None, input_embeds, train);
569                (outputs.start_logits, outputs.end_logits)
570            }
571            Self::XLNet(ref model) => {
572                let outputs =
573                    model.forward_t(input_ids, mask, None, None, None, None, input_embeds, train);
574                (outputs.start_logits, outputs.end_logits)
575            }
576            Self::Reformer(ref model) => {
577                let outputs = model
578                    .forward_t(input_ids, None, None, mask, None, train)
579                    .expect("Error in reformer forward pass");
580                (outputs.start_logits, outputs.end_logits)
581            }
582            Self::Longformer(ref model) => {
583                let outputs = model
584                    .forward_t(input_ids, mask, None, None, None, None, train)
585                    .expect("Error in reformer forward pass");
586                (outputs.start_logits, outputs.end_logits)
587            }
588            Self::FNet(ref model) => {
589                let outputs = model
590                    .forward_t(input_ids, None, None, None, train)
591                    .expect("Error in fnet forward pass");
592                (outputs.start_logits, outputs.end_logits)
593            }
594            #[cfg(feature = "onnx")]
595            Self::ONNX(ref model) => {
596                let outputs = model
597                    .forward(
598                        input_ids,
599                        mask.map(|tensor| tensor.to_kind(Kind::Int64)).as_ref(),
600                        _token_type_ids,
601                        None,
602                        input_embeds,
603                    )
604                    .expect("Error in ONNX forward pass.");
605                (outputs.start_logits.unwrap(), outputs.end_logits.unwrap())
606            }
607        }
608    }
609}
610
611/// # QuestionAnsweringModel to perform extractive question answering
612pub struct QuestionAnsweringModel {
613    tokenizer: TokenizerOption,
614    pad_idx: i64,
615    sep_idx: i64,
616    max_seq_len: usize,
617    doc_stride: usize,
618    max_query_length: usize,
619    max_answer_len: usize,
620    qa_model: QuestionAnsweringOption,
621    device: Device,
622}
623
624impl QuestionAnsweringModel {
625    /// Build a new `QuestionAnsweringModel`
626    ///
627    /// # Arguments
628    ///
629    /// * `question_answering_config` - `QuestionAnsweringConfig` object containing the resource references (model, vocabulary, configuration) and device placement (CPU/GPU)
630    ///
631    /// # Example
632    ///
633    /// ```no_run
634    /// # fn main() -> anyhow::Result<()> {
635    /// use rust_bert::pipelines::question_answering::QuestionAnsweringModel;
636    ///
637    /// let qa_model = QuestionAnsweringModel::new(Default::default())?;
638    /// # Ok(())
639    /// # }
640    /// ```
641    pub fn new(
642        question_answering_config: QuestionAnsweringConfig,
643    ) -> Result<QuestionAnsweringModel, RustBertError> {
644        let vocab_path = question_answering_config.vocab_resource.get_local_path()?;
645        let merges_path = question_answering_config
646            .merges_resource
647            .as_ref()
648            .map(|resource| resource.get_local_path())
649            .transpose()?;
650
651        let tokenizer = TokenizerOption::from_file(
652            question_answering_config.model_type,
653            vocab_path.to_str().unwrap(),
654            merges_path.as_deref().map(|path| path.to_str().unwrap()),
655            question_answering_config.lower_case,
656            question_answering_config.strip_accents,
657            question_answering_config.add_prefix_space,
658        )?;
659        Self::new_with_tokenizer(question_answering_config, tokenizer)
660    }
661
662    /// Build a new `QuestionAnsweringModel` with a provided tokenizer.
663    ///
664    /// # Arguments
665    ///
666    /// * `question_answering_config` - `QuestionAnsweringConfig` object containing the resource references (model, vocabulary, configuration) and device placement (CPU/GPU)
667    /// * `tokenizer` - `TokenizerOption` tokenizer to use for question answering.
668    ///
669    /// # Example
670    ///
671    /// ```no_run
672    /// # fn main() -> anyhow::Result<()> {
673    /// use rust_bert::pipelines::common::{ModelType, TokenizerOption};
674    /// use rust_bert::pipelines::question_answering::QuestionAnsweringModel;
675    /// let tokenizer = TokenizerOption::from_file(
676    ///     ModelType::Bert,
677    ///     "path/to/vocab.txt",
678    ///     None,
679    ///     false,
680    ///     None,
681    ///     None,
682    /// )?;
683    /// let qa_model = QuestionAnsweringModel::new_with_tokenizer(Default::default(), tokenizer)?;
684    /// # Ok(())
685    /// # }
686    /// ```
687    pub fn new_with_tokenizer(
688        question_answering_config: QuestionAnsweringConfig,
689        tokenizer: TokenizerOption,
690    ) -> Result<QuestionAnsweringModel, RustBertError> {
691        let qa_model = QuestionAnsweringOption::new(&question_answering_config)?;
692
693        let pad_idx = tokenizer
694            .get_pad_id()
695            .expect("The Tokenizer used for Question Answering should contain a PAD id");
696        let sep_idx = tokenizer
697            .get_sep_id()
698            .expect("The Tokenizer used for Question Answering should contain a SEP id");
699
700        if question_answering_config.max_seq_length
701            < (question_answering_config.max_query_length
702                + question_answering_config.doc_stride
703                + 24)
704        {
705            return Err(RustBertError::InvalidConfigurationError(format!(
706                "This configuration could cause an excessive number of sliding windows generated.\
707                Please ensure max_seq_length > max_query_length + doc_stride + 24.\
708                Got max_seq_length: {}, max_query_length: {}, doc_stride: {}",
709                question_answering_config.max_seq_length,
710                question_answering_config.max_query_length,
711                question_answering_config.doc_stride
712            )));
713        }
714        let device = get_device(
715            question_answering_config.model_resource,
716            question_answering_config.device,
717        );
718        Ok(QuestionAnsweringModel {
719            tokenizer,
720            pad_idx,
721            sep_idx,
722            max_seq_len: question_answering_config.max_seq_length,
723            doc_stride: question_answering_config.doc_stride,
724            max_query_length: question_answering_config.max_query_length,
725            max_answer_len: question_answering_config.max_answer_length,
726            qa_model,
727            device,
728        })
729    }
730
731    /// Get a reference to the model tokenizer.
732    pub fn get_tokenizer(&self) -> &TokenizerOption {
733        &self.tokenizer
734    }
735
736    /// Get a mutable reference to the model tokenizer.
737    pub fn get_tokenizer_mut(&mut self) -> &mut TokenizerOption {
738        &mut self.tokenizer
739    }
740
741    /// Perform extractive question answering given a list of `QaInputs`
742    ///
743    /// # Arguments
744    ///
745    /// * `qa_inputs` - `&[QaInput]` Array of Question Answering inputs (context and question pairs)
746    /// * `top_k` - return the top-k answers for each QaInput. Set to 1 to return only the best answer.
747    /// * `batch_size` - maximum batch size for the model forward pass.
748    ///
749    /// # Returns
750    /// * `Vec<Vec<Answer>>` Vector (same length as `qa_inputs`) of vectors (each of length `top_k`) containing the extracted answers.
751    ///
752    /// # Example
753    ///
754    /// ```no_run
755    /// # fn main() -> anyhow::Result<()> {
756    /// use rust_bert::pipelines::question_answering::{QaInput, QuestionAnsweringModel};
757    ///
758    /// let qa_model = QuestionAnsweringModel::new(Default::default())?;
759    ///
760    /// let question_1 = String::from("Where does Amy live ?");
761    /// let context_1 = String::from("Amy lives in Amsterdam");
762    /// let question_2 = String::from("Where does Eric live");
763    /// let context_2 = String::from("While Amy lives in Amsterdam, Eric is in The Hague.");
764    ///
765    /// let qa_input_1 = QaInput {
766    ///     question: question_1,
767    ///     context: context_1,
768    /// };
769    /// let qa_input_2 = QaInput {
770    ///     question: question_2,
771    ///     context: context_2,
772    /// };
773    /// let answers = qa_model.predict(&[qa_input_1, qa_input_2], 1, 32);
774    ///
775    /// # Ok(())
776    /// # }
777    /// ```
778    pub fn predict(
779        &self,
780        qa_inputs: &[QaInput],
781        top_k: i64,
782        batch_size: usize,
783    ) -> Vec<Vec<Answer>> {
784        let mut features: Vec<QaFeature> = qa_inputs
785            .iter()
786            .enumerate()
787            .flat_map(|(example_index, qa_example)| {
788                self.generate_features(
789                    qa_example,
790                    self.max_seq_len,
791                    self.doc_stride,
792                    self.max_query_length,
793                    example_index as i64,
794                )
795            })
796            .collect();
797
798        let mut example_top_k_answers_map: HashMap<usize, Vec<Answer>> = HashMap::new();
799        let mut start = 0usize;
800        let len_features = features.len();
801
802        while start < len_features {
803            let end = start + min(len_features - start, batch_size);
804            let batch_features = &mut features[start..end];
805            no_grad(|| {
806                let (input_ids, attention_masks, token_type_ids) =
807                    self.pad_features(batch_features);
808
809                let (start_logits, end_logits) = self.qa_model.forward_t(
810                    Some(&input_ids),
811                    Some(&attention_masks),
812                    None,
813                    Some(&token_type_ids),
814                    false,
815                );
816
817                let start_logits = start_logits.detach();
818                let end_logits = end_logits.detach();
819                let example_index_to_feature_end_position: Vec<(usize, i64)> = batch_features
820                    .iter()
821                    .enumerate()
822                    .map(|(feature_index, feature)| {
823                        (feature.example_index as usize, feature_index as i64 + 1)
824                    })
825                    .collect();
826
827                let mut feature_id_start = 0;
828
829                for (example_id, max_feature_id) in example_index_to_feature_end_position {
830                    let mut answers: Vec<Answer> = vec![];
831                    let example = &qa_inputs[example_id];
832                    for feature_idx in feature_id_start..max_feature_id {
833                        let feature = &batch_features[feature_idx as usize];
834                        let p_mask = (Tensor::from_slice(&feature.p_mask) - 1)
835                            .abs()
836                            .to_device(start_logits.device())
837                            .eq(0);
838
839                        let start = start_logits
840                            .get(feature_idx)
841                            .masked_fill(&p_mask, get_min(start_logits.kind()).unwrap());
842                        let end = end_logits
843                            .get(feature_idx)
844                            .masked_fill(&p_mask, get_min(start_logits.kind()).unwrap());
845
846                        let start = start.softmax(0, start.kind());
847                        let end = end.softmax(0, end.kind());
848
849                        let (starts, ends, scores) = self.decode(&start, &end, top_k);
850
851                        for idx in 0..starts.len() {
852                            let start_pos = feature.offsets[starts[idx] as usize]
853                                .unwrap_or(Offset { begin: 0, end: 0 })
854                                .begin as usize;
855                            let end_pos = feature.offsets[ends[idx] as usize]
856                                .unwrap_or(Offset { begin: 0, end: 0 })
857                                .end as usize;
858                            let answer = example
859                                .context
860                                .chars()
861                                .take(end_pos)
862                                .skip(start_pos)
863                                .collect::<String>();
864
865                            answers.push(Answer {
866                                score: scores[idx],
867                                start: start_pos,
868                                end: end_pos,
869                                answer,
870                            });
871                        }
872                    }
873                    feature_id_start = max_feature_id;
874                    let example_answers = example_top_k_answers_map.entry(example_id).or_default();
875                    example_answers.extend(answers);
876                }
877            });
878            start = end;
879        }
880        let mut all_answers = vec![];
881        for example_id in 0..qa_inputs.len() {
882            if let Some(answers) = example_top_k_answers_map.get_mut(&example_id) {
883                remove_duplicates(answers).sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
884                all_answers.push(answers[..min(answers.len(), top_k as usize)].to_vec());
885            } else {
886                all_answers.push(vec![]);
887            }
888        }
889        all_answers
890    }
891
892    fn decode(&self, start: &Tensor, end: &Tensor, top_k: i64) -> (Vec<i64>, Vec<i64>, Vec<f64>) {
893        let outer = start.unsqueeze(-1).matmul(&end.unsqueeze(0));
894        let start_dim = start.size()[0];
895        let end_dim = end.size()[0];
896        let candidates = outer
897            .triu(0)
898            .tril(self.max_answer_len as i64 - 1)
899            .flatten(0, -1);
900        let idx_sort = if top_k == 1 {
901            candidates.argmax(0, true)
902        } else if candidates.size()[0] < top_k {
903            candidates.argsort(0, true)
904        } else {
905            candidates.argsort(0, true).slice(0, 0, top_k, 1)
906        };
907        let mut start: Vec<i64> = vec![];
908        let mut end: Vec<i64> = vec![];
909        let mut scores: Vec<f64> = vec![];
910        for flat_index_position in 0..idx_sort.size()[0] {
911            let flat_index = idx_sort.int64_value(&[flat_index_position]);
912            scores.push(candidates.double_value(&[flat_index]));
913            start.push(flat_index / start_dim);
914            end.push(flat_index % end_dim);
915        }
916        (start, end, scores)
917    }
918
919    fn generate_features(
920        &self,
921        qa_example: &QaInput,
922        max_seq_length: usize,
923        doc_stride: usize,
924        max_query_length: usize,
925        example_index: i64,
926    ) -> Vec<QaFeature> {
927        let mut encoded_query = self.tokenizer.tokenize_with_offsets(&qa_example.question);
928        encoded_query.tokens.truncate(max_query_length);
929        encoded_query.offsets.truncate(max_query_length);
930        encoded_query.reference_offsets.truncate(max_query_length);
931        encoded_query.masks.truncate(max_query_length);
932        let encoded_query = TokenIdsWithOffsets {
933            ids: self.tokenizer.convert_tokens_to_ids(&encoded_query.tokens),
934            offsets: encoded_query.offsets,
935            reference_offsets: encoded_query.reference_offsets,
936            masks: encoded_query.masks,
937        };
938
939        let sequence_added_tokens = self
940            .tokenizer
941            .build_input_with_special_tokens(
942                TokenIdsWithOffsets {
943                    ids: vec![],
944                    offsets: vec![],
945                    reference_offsets: vec![],
946                    masks: vec![],
947                },
948                None,
949            )
950            .token_ids
951            .len();
952
953        let sequence_pair_added_tokens = self
954            .tokenizer
955            .build_input_with_special_tokens(
956                TokenIdsWithOffsets {
957                    ids: vec![],
958                    offsets: vec![],
959                    reference_offsets: vec![],
960                    masks: vec![],
961                },
962                Some(TokenIdsWithOffsets {
963                    ids: vec![],
964                    offsets: vec![],
965                    reference_offsets: vec![],
966                    masks: vec![],
967                }),
968            )
969            .token_ids
970            .len();
971
972        let mut spans: Vec<QaFeature> = vec![];
973
974        let tokenized_context = self.tokenizer.tokenize_with_offsets(&qa_example.context);
975        let encoded_context = TokenIdsWithOffsets {
976            ids: self
977                .tokenizer
978                .convert_tokens_to_ids(&tokenized_context.tokens),
979            offsets: tokenized_context.offsets,
980            reference_offsets: tokenized_context.reference_offsets,
981            masks: tokenized_context.masks,
982        };
983        let max_context_length =
984            max_seq_length - sequence_pair_added_tokens - encoded_query.ids.len();
985
986        let mut start_token = 0_usize;
987        while (spans.len() * doc_stride) < encoded_context.ids.len() {
988            let end_token = min(start_token + max_context_length, encoded_context.ids.len());
989            let sub_encoded_context = TokenIdsWithOffsets {
990                ids: encoded_context.ids[start_token..end_token].to_vec(),
991                offsets: encoded_context.offsets[start_token..end_token].to_vec(),
992                reference_offsets: encoded_context.reference_offsets[start_token..end_token]
993                    .to_vec(),
994                masks: encoded_context.masks[start_token..end_token].to_vec(),
995            };
996
997            let encoded_span = self
998                .tokenizer
999                .build_input_with_special_tokens(encoded_query.clone(), Some(sub_encoded_context));
1000            let p_mask = self.get_mask(
1001                &encoded_span,
1002                encoded_query.ids.len() + sequence_added_tokens,
1003            );
1004            let qa_feature = QaFeature {
1005                input_ids: encoded_span.token_ids,
1006                offsets: encoded_span.token_offsets,
1007                token_type_ids: encoded_span.segment_ids,
1008                p_mask,
1009                example_index,
1010            };
1011            spans.push(qa_feature);
1012            if end_token == encoded_context.ids.len() {
1013                break;
1014            }
1015            start_token = end_token - doc_stride;
1016        }
1017        spans
1018    }
1019
1020    fn pad_features(&self, features: &mut [QaFeature]) -> (Tensor, Tensor, Tensor) {
1021        let max_len = features
1022            .iter()
1023            .map(|feature| feature.input_ids.len())
1024            .max()
1025            .unwrap();
1026
1027        let attention_masks = features
1028            .iter()
1029            .map(|feature| &feature.input_ids)
1030            .map(|input| {
1031                let mut attention_mask = Vec::with_capacity(max_len);
1032                attention_mask.resize(input.len(), 1);
1033                attention_mask.resize(max_len, 0);
1034                attention_mask
1035            })
1036            .map(|input| Tensor::from_slice(&(input)))
1037            .collect::<Vec<_>>();
1038
1039        for feature in features.iter_mut() {
1040            feature.offsets.resize(max_len, None);
1041            feature.p_mask.resize(max_len, 1);
1042            feature.input_ids.resize(max_len, self.pad_idx);
1043            feature
1044                .token_type_ids
1045                .resize(max_len, *feature.token_type_ids.last().unwrap_or(&0));
1046        }
1047
1048        let padded_input_ids = features
1049            .iter_mut()
1050            .map(|input| Tensor::from_slice(input.input_ids.as_slice()))
1051            .collect::<Vec<_>>();
1052
1053        let padded_token_type_ids = features
1054            .iter_mut()
1055            .map(|input| Tensor::from_slice(input.token_type_ids.as_slice()))
1056            .collect::<Vec<_>>();
1057
1058        let input_ids = Tensor::stack(&padded_input_ids, 0).to(self.device);
1059        let attention_masks = Tensor::stack(&attention_masks, 0).to(self.device);
1060        let token_type_ids = Tensor::stack(&padded_token_type_ids, 0)
1061            .to(self.device)
1062            .to_kind(Kind::Int64);
1063        (input_ids, attention_masks, token_type_ids)
1064    }
1065
1066    fn get_mask(&self, encoded_span: &TokenizedInput, question_length: usize) -> Vec<i8> {
1067        let sep_indices: Vec<usize> = encoded_span
1068            .token_ids
1069            .iter()
1070            .enumerate()
1071            .filter(|(_, &value)| value == self.sep_idx)
1072            .map(|(position, _)| position)
1073            .collect();
1074
1075        let mut p_mask: Vec<i8> = Vec::with_capacity(encoded_span.token_ids.len());
1076        p_mask.extend(vec![1; question_length]);
1077        p_mask.extend(vec![0; encoded_span.token_ids.len() - question_length]);
1078        for sep_position in sep_indices {
1079            p_mask[sep_position] = 1;
1080        }
1081        p_mask
1082    }
1083}
1084
1085pub fn squad_processor(file_path: PathBuf) -> Vec<QaInput> {
1086    let file = fs::File::open(file_path).expect("unable to open file");
1087    let json: serde_json::Value =
1088        serde_json::from_reader(file).expect("JSON not properly formatted");
1089    let data = json
1090        .get("data")
1091        .expect("SQuAD file does not contain data field")
1092        .as_array()
1093        .expect("Data array not properly formatted");
1094
1095    let mut qa_inputs: Vec<QaInput> = Vec::with_capacity(data.len());
1096    for qa_input in data.iter() {
1097        let qa_input = qa_input.as_object().unwrap();
1098        let paragraphs = qa_input.get("paragraphs").unwrap().as_array().unwrap();
1099        for paragraph in paragraphs.iter() {
1100            let paragraph = paragraph.as_object().unwrap();
1101            let context = paragraph.get("context").unwrap().as_str().unwrap();
1102            let qas = paragraph.get("qas").unwrap().as_array().unwrap();
1103            for qa in qas.iter() {
1104                let question = qa
1105                    .as_object()
1106                    .unwrap()
1107                    .get("question")
1108                    .unwrap()
1109                    .as_str()
1110                    .unwrap();
1111                qa_inputs.push(QaInput {
1112                    question: question.to_owned(),
1113                    context: context.to_owned(),
1114                });
1115            }
1116        }
1117    }
1118    qa_inputs
1119}
1120
1121#[cfg(test)]
1122mod test {
1123    use super::*;
1124
1125    #[test]
1126    #[ignore] // no need to run, compilation is enough to verify it is Send
1127    fn test() {
1128        let config = QuestionAnsweringConfig::default();
1129        let _: Box<dyn Send> = Box::new(QuestionAnsweringModel::new(config));
1130    }
1131}
rust_bert/pipelines/question_answering.rs

rust_bert/pipelines/
question_answering.rs
