class_docstr: Here, we use three networks to predict the start scores, end scores and answerable scores.
  
  We formulate p(start, end, answerable | contextual_embedding) as the product of the
  following three terms:
  
  - p(start | contextual_embedding)
  - p(end | start, contextual_embedding)
  - p(answerable | contextual_embedding)
  
  In the inference phase, we are able to use beam search to do the inference.
  
  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: ["Here",",","we","use","three","networks","to","predict","the","start","scores",",","end","scores","and","answerable","scores",".","We","formulate","p","(","start",",","end",",","answerable","|","contextual_embedding",")","as","the","product","of","the","following","three","terms",":","-","p","(","start","|","contextual_embedding",")","-","p","(","end","|","start",",","contextual_embedding",")","-","p","(","answerable","|","contextual_embedding",")","In","the","inference","phase",",","we","are","able","to","use","beam","search","to","do","the","inference",".","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, start_top_n, end_top_n)
  end_top_index
      Index of the top end logits
      Shape (batch_size, start_top_n, end_top_n)
  answerable_logits
      The answerable logits. Here 0 --> answerable and 1 --> not answerable.
      Shape (batch_size, sequence_length, 2)
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, start_top_n, end_top_n)
      end_top_index
          Index of the top end logits
          Shape (batch_size, start_top_n, end_top_n)
      answerable_logits
          The answerable logits. Here 0 --> answerable and 1 --> not answerable.
          Shape (batch_size, sequence_length, 2)
      """
      # 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)
      start_logits = self.get_start_logits(
          contextual_embeddings, p_mask
      )
      # The shape of start_top_index will be (..., start_top_n)
      start_top_logits, start_top_index = npx.topk(
          start_logits, k=start_top_n, axis=-1, ret_typ="both"
      )
      end_logits = self.get_end_logits(
          contextual_embeddings, start_top_index, p_mask
      )
      # 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 = npx.topk(
          end_logits, k=end_top_n, axis=-1, ret_typ="both"
      )
      answerable_logits = self.get_answerable_logits(
          contextual_embeddings, p_mask
      )
      return (
          start_top_logits,
          start_top_index,
          end_top_logits,
          end_top_index,
          answerable_logits,
      )
  
imports: import numpy
  import gluonnlp.data.batchify
  import collections
  import dataclasses
  import mxnet
  from mxnet import np, npx
  from mxnet.gluon import nn
  from mxnet.gluon.nn import HybridBlock
  from dataclasses import dataclass
  from tqdm import tqdm
  from typing import Optional, List, Tuple
  from gluonnlp.data.tokenizers import BaseTokenizerWithVocab
  from gluonnlp.utils.preprocessing import match_tokens_with_char_spans
  from gluonnlp.layers import get_activation
  from gluonnlp.op import select_vectors_by_position
  from gluonnlp.attention_cell import masked_logsoftmax, masked_softmax
  import string
  import re
  import json
label_desc: Use image node_id 6 for calling the ModelForQAConditionalV1 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,  answerable_logits