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 os
  import json
  import time
  import logging
  import argparse
  import ast
  import functools
  import collections
  import dataclasses
  from dataclasses import dataclass
  from multiprocessing import Pool, cpu_count
  from neural_compressor.experimental import Quantization
  import mxnet
  import numpy
  from mxnet.lr_scheduler import PolyScheduler
  import gluonnlp.data.batchify
  from models import ModelForQABasic, ModelForQAConditionalV1
  from eval_utils import squad_eval
  from squad_utils import SquadFeature, get_squad_examples, convert_squad_example_to_feature, get_squad_examples_from_json
  from gluonnlp.models import get_backbone
  from gluonnlp.utils.misc import repeat, grouper, set_seed, init_comm, logging_config, parse_ctx
  from gluonnlp.initializer import TruncNorm
  from gluonnlp.data.sampler import SplitSampler
  from gluonnlp.utils.parameter import grad_global_norm, clip_grad_global_norm, count_parameters, deduplicate_param_dict
  import custom_strategy
label_desc: Use image node_id 5 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