def _get_sub_paras(para, tokenizer, max_seq_length, doc_stride, total):
  """Split paragraph object into sub-paragraphs with maximum length."""
  max_tokens_for_doc = max_seq_length - 2  # -2 for [CLS] and [SEP]
  para_tokens, para_char_to_token = bert_utils.preprocess_text(
      para["context"], tokenizer)

  # Get mention token start and ends.
  mentions = []
  for im, ment in enumerate(para["mentions"]):
    st_tok = para_char_to_token[ment["start"]][0]
    en_tok = para_char_to_token[ment["start"] + len(ment["text"]) - 1][1]
    mentions.append({
        "wikidata_id": ment["wikidata_id"],
        "name": ment["name"],
        "text": ment["text"],
        "start_token": st_tok,
        "end_token": en_tok,
        "orig_index": im,
    })

  # Get sub para spans.
  sub_paras = []
  start_offset = 0
  while start_offset < len(para_tokens):
    length = len(para_tokens) - start_offset
    if length > max_tokens_for_doc:
      length = max_tokens_for_doc
    sub_paras.append((start_offset, length))
    if start_offset + length == len(para_tokens):
      break
    start_offset += min(length, doc_stride)

  # Assign each mention to a sub_para.
  sub_para_to_mentions = {i: [] for i in range(len(sub_paras))}
  for ment in mentions:
    best_score, best_index = None, None
    for ii, subp in enumerate(sub_paras):
      subp_end = subp[0] + subp[1] - 1
      if ment["start_token"] < subp[0] or ment["end_token"] > subp_end:
        continue
      score = min(ment["start_token"] - subp[0], subp_end - ment["end_token"])
      if best_score is None or score > best_score:
        best_score = score
        best_index = ii
    ment["start_token"] -= sub_paras[best_index][0]
    ment["end_token"] -= sub_paras[best_index][0]
    sub_para_to_mentions[best_index].append(ment)

  # Create a list of sub_para objects.
  sub_para_objects = []
  for ii, subp in enumerate(sub_paras):
    sub_para_objects.append({
        "id": total[0],
        "wikidata_id": para["wikidata_id"],
        "mentions": sub_para_to_mentions[ii],
        "tokens": para_tokens[subp[0]:subp[0] + subp[1]],
    })
    total[0] += 1

  return sub_para_objects
Split paragraph object into sub-paragraphs with maximum length.
language/labs/drkit/wikidata/index.py