class_docstr: The basic pretrained model for QA. It is used in the original BERT paper for SQuAD 1.1.
  
  Here, we directly use the backbone network to extract the contextual embeddings and use
  another dense layer to map the contextual embeddings to the start scores and end scores.
  
  use_segmentation is used to mark whether we segment the input sentence. In RoBERTa and XLMR,
  this flag is set to True, then the QA model no longer accept `token_types` as valid input.
  
  - use_segmentation=True:
      tokens :      <CLS> Question <SEP> Context <SEP>
      token_typess:  0       0       0      1      1
  
  - use_segmentation=False:
      tokens :      <CLS> Question <SEP> Context <SEP>
      token_typess:  None
class_docstr_tok: ["The","basic","pretrained","model","for","QA",".","It","is","used","in","the","original","BERT","paper","for","SQuAD","1.1",".","Here",",","we","directly","use","the","backbone","network","to","extract","the","contextual","embeddings","and","use","another","dense","layer","to","map","the","contextual","embeddings","to","the","start","scores","and","end","scores",".","use_segmentation","is","used","to","mark","whether","we","segment","the","input","sentence",".","In","RoBERTa","and","XLMR",",","this","flag","is","set","to","True",",","then","the","QA","model","no","longer","accept","`","token_types","`","as","valid","input",".","-","use_segmentation=True",":","tokens",":","<","CLS",">","Question","<","SEP",">","Context","<","SEP",">","token_typess",":","0","0","0","1","1","-","use_segmentation=False",":","tokens",":","<","CLS",">","Question","<","SEP",">","Context","<","SEP",">","token_typess",":","None"]
docstr: Get the inference result with beam search
  
  Parameters
  ----------
  tokens
      The input tokens. Shape (batch_size, sequence_length)
  token_types
      The input token types. Shape (batch_size, sequence_length)
  valid_length
      The valid length of the tokens. Shape (batch_size,)
  p_mask
      The mask which indicates that some tokens won't be used in the calculation.
      Shape (batch_size, sequence_length)
  start_top_n
      The number of candidates to select for the start position.
  end_top_n
      The number of candidates to select for the end position.
  
  Returns
  -------
  start_top_logits
      The top start logits
      Shape (batch_size, start_top_n)
  start_top_index
      Index of the top start logits
      Shape (batch_size, start_top_n)
  end_top_logits
      The top end logits.
      Shape (batch_size, end_top_n)
  end_top_index
      Index of the top end logits
      Shape (batch_size, end_top_n)
docstr_tok: ["Get","the","inference","result","with","beam","search","Parameters","--","--","--","--","--","tokens","The","input","tokens",".","Shape","(","batch_size",",","sequence_length",")","token_types","The","input","token","types",".","Shape","(","batch_size",",","sequence_length",")","valid_length","The","valid","length","of","the","tokens",".","Shape","(","batch_size",",",")","p_mask","The","mask","which","indicates","that","some","tokens","wo","n't","be","used","in","the","calculation",".","Shape","(","batch_size",",","sequence_length",")","start_top_n","The","number","of","candidates","to","select","for","the","start","position",".","end_top_n","The","number","of","candidates","to","select","for","the","end","position",".","Returns","--","--","--","-","start_top_logits","The","top","start","logits","Shape","(","batch_size",",","start_top_n",")","start_top_index","Index","of","the","top","start","logits","Shape","(","batch_size",",","start_top_n",")","end_top_logits","The","top","end","logits",".","Shape","(","batch_size",",","end_top_n",")","end_top_index","Index","of","the","top","end","logits","Shape","(","batch_size",",","end_top_n",")"]
code: def inference(
      self,
      tokens,
      token_types,
      valid_length,
      p_mask,
      start_top_n: int = 5,
      end_top_n: int = 5,
  ):
      """Get the inference result with beam search
  
      Parameters
      ----------
      tokens
          The input tokens. Shape (batch_size, sequence_length)
      token_types
          The input token types. Shape (batch_size, sequence_length)
      valid_length
          The valid length of the tokens. Shape (batch_size,)
      p_mask
          The mask which indicates that some tokens won't be used in the calculation.
          Shape (batch_size, sequence_length)
      start_top_n
          The number of candidates to select for the start position.
      end_top_n
          The number of candidates to select for the end position.
  
      Returns
      -------
      start_top_logits
          The top start logits
          Shape (batch_size, start_top_n)
      start_top_index
          Index of the top start logits
          Shape (batch_size, start_top_n)
      end_top_logits
          The top end logits.
          Shape (batch_size, end_top_n)
      end_top_index
          Index of the top end logits
          Shape (batch_size, end_top_n)
      """
      # Shape (batch_size, sequence_length, C)
      if self.use_segmentation:
          contextual_embeddings = self.backbone(
              tokens, token_types, valid_length
          )
      else:
          contextual_embeddings = self.backbone(tokens, valid_length)
      scores = self.qa_outputs(contextual_embeddings)
      start_scores = scores[:, :, 0]
      end_scores = scores[:, :, 1]
      start_logits = masked_logsoftmax(
          start_scores, mask=p_mask, axis=-1
      )
      end_logits = masked_logsoftmax(end_scores, mask=p_mask, axis=-1)
      # The shape of start_top_index will be (..., start_top_n)
      start_top_logits, start_top_index = mx.npx.topk(
          start_logits, k=start_top_n, axis=-1, ret_typ="both"
      )
      # Note that end_top_index and end_top_log_probs have shape (bsz, start_n_top, end_n_top)
      # So that for each start position, there are end_n_top end positions on the third dim.
      end_top_logits, end_top_index = mx.npx.topk(
          end_logits, k=end_top_n, axis=-1, ret_typ="both"
      )
      return (
          start_top_logits,
          start_top_index,
          end_top_logits,
          end_top_index,
      )
  
file_path: gluon-nlp/docs/model_zoo/question_answering/models.py
imports: import mxnet
  from mxnet import np, npx
  from mxnet.gluon import nn, HybridBlock
  from mxnet.util import use_np
  from gluonnlp.layers import get_activation
  from gluonnlp.op import select_vectors_by_position
  from gluonnlp.attention_cell import masked_logsoftmax, masked_softmax
label_desc: Use image node_id 3 for calling the ModelForQABasic obj's underlying member method code with example usage: obj.inference(tokens, token_types, valid_length, p_mask, start_top_n, end_top_n) and returns: start_top_logits,  start_top_index,  end_top_logits,  end_top_index