code: def predict_extended(
      original_feature,
      chunked_features,
      results,
      n_best_size,
      max_answer_length=64,
      start_top_n=5,
      end_top_n=5,
  ):
      """Get prediction results for SQuAD.
  
      Start Logits: (B, N_start)
      End Logits: (B, N_start, N_end)
  
      Parameters
      ----------
      original_feature:
          The original SquadFeature before chunked
      chunked_features
          List of ChunkFeatures
      results
          List of model predictions for span start and span end.
      n_best_size
          Best N results written to file
      max_answer_length
          Maximum length of the answer tokens.
      start_top_n
          Number of start-position candidates
      end_top_n
          Number of end-position candidates
      Returns
      -------
      not_answerable_score
          Model's estimate that the question is not answerable.
      prediction
          The final prediction.
      nbest_json
          n-best predictions with their probabilities.
      """
      not_answerable_score = 1000000  # Score for not-answerable. We set it to be a large and positive
      # If one chunk votes for answerable, we will treat the context as answerable,
      # Thus, the overall not_answerable_score = min(chunk_not_answerable_score)
      all_start_idx = []
      all_end_idx = []
      all_pred_score = []
      context_length = len(original_feature.context_token_ids)
      token_max_context_score = _np.full(
          (len(chunked_features), context_length),
          -_np.inf,
          dtype=_np.float32,
      )
      for i, chunked_feature in enumerate(chunked_features):
          chunk_start = chunked_feature.chunk_start
          chunk_length = chunked_feature.chunk_length
          for j in range(chunk_start, chunk_start + chunk_length):
              # This is a heuristic score
              # TODO investigate the impact
              token_max_context_score[i, j] = (
                  min(
                      j - chunk_start,
                      chunk_start + chunk_length - 1 - j,
                  )
                  + 0.01 * chunk_length
              )
      token_max_chunk_id = token_max_context_score.argmax(axis=0)
  
      for chunk_id, (result, chunk_feature) in enumerate(
          zip(results, chunked_features)
      ):
          # We use the log-likelihood as the not answerable score.
          # Thus, a high score indicates that the answer is not answerable
          cur_not_answerable_score = float(result.answerable_logits[1])
          not_answerable_score = min(
              not_answerable_score, cur_not_answerable_score
          )
          # Calculate the start_logits + end_logits as the overall score
          context_offset = chunk_feature.context_offset
          chunk_start = chunk_feature.chunk_start
          chunk_length = chunk_feature.chunk_length
          for i in range(start_top_n):
              for j in range(end_top_n):
                  pred_score = (
                      result.start_top_logits[i]
                      + result.end_top_logits[i, j]
                  )
                  start_index = result.start_top_index[i]
                  end_index = result.end_top_index[i, j]
                  # We could hypothetically create invalid predictions, e.g., predict
                  # that the start of the answer span is in the query tokens or out of
                  # the chunk. We throw out all invalid predictions.
                  if (
                      not (
                          context_offset
                          <= start_index
                          < context_offset + chunk_length
                      )
                      or not (
                          context_offset
                          <= end_index
                          < context_offset + chunk_length
                      )
                      or end_index < start_index
                  ):
                      continue
                  pred_answer_length = end_index - start_index + 1
                  if pred_answer_length > max_answer_length:
                      continue
                  start_idx = int(
                      start_index - context_offset + chunk_start
                  )
                  end_idx = int(
                      end_index - context_offset + chunk_start
                  )
                  if token_max_chunk_id[start_idx] != chunk_id:
                      continue
                  all_start_idx.append(start_idx)
                  all_end_idx.append(end_idx)
                  all_pred_score.append(pred_score)
      sorted_start_end_score = sorted(
          zip(all_start_idx, all_end_idx, all_pred_score),
          key=lambda args: args[-1],
          reverse=True,
      )
      nbest = []
      context_text = original_feature.context_text
      context_token_offsets = original_feature.context_token_offsets
      seen_predictions = set()
      for start_idx, end_idx, pred_score in sorted_start_end_score:
          if len(seen_predictions) >= n_best_size:
              break
          pred_answer = context_text[
              context_token_offsets[start_idx][
                  0
              ] : context_token_offsets[end_idx][1]
          ]
          seen_predictions.add(pred_answer)
          nbest.append((pred_answer, pred_score))
  
      # In very rare edge cases we could have no valid predictions. So we
      # just create a nonce prediction in this case to avoid failure.
      if len(nbest) == 0:
          nbest.append(("", float("-inf")))
      all_scores = _np.array(
          [ele[1] for ele in nbest], dtype=_np.float32
      )
      probs = _np.exp(all_scores) / _np.sum(_np.exp(all_scores))
      nbest_json = []
      for i, (entry, prob) in enumerate(zip(nbest, probs)):
          output = collections.OrderedDict()
          output["text"] = entry[0]
          output["probability"] = float(prob)
          nbest_json.append(output)
  
      assert len(nbest_json) >= 1
      return not_answerable_score, nbest[0][0], nbest_json
  
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 1 for calling a global function with example usage: predict_extended(original_feature, chunked_features, results, n_best_size, max_answer_length, start_top_n, end_top_n) and returns: not_answerable_score,  nbest,  nbest_json