code: def predict(
      trainer,
      data_generator,
      batch_size,
      max_seqlen,
      use_regex_tokens,
      num_workers,
  ):
      """
      The guts of the prediction method
  
      Calls trainer.predict() over and over until we have predictions for all of the text
      """
      all_preds = []
      all_raw = []
  
      sorted_data = SortedDataset(data_generator)
      dataloader = TorchDataLoader(
          sorted_data,
          batch_size=batch_size,
          collate_fn=sorted_data.collate,
          num_workers=num_workers,
      )
      for batch_idx, batch in enumerate(dataloader):
          num_sentences = len(batch[3])
          N = len(batch[3][0])
          for paragraph in batch[3]:
              all_raw.append(list(paragraph))
  
          if N <= max_seqlen:
              pred = np.argmax(trainer.predict(batch), axis=2)
          else:
              # TODO: we could shortcircuit some processing of
              # long strings of PAD by tracking which rows are finished
              idx = [0] * num_sentences
              adv = [0] * num_sentences
              para_lengths = [x.index("<PAD>") for x in batch[3]]
              pred = [[] for _ in range(num_sentences)]
              while True:
                  ens = [
                      min(N - idx1, max_seqlen)
                      for idx1, N in zip(idx, para_lengths)
                  ]
                  en = max(ens)
                  batch1 = (
                      batch[0][:, :en],
                      batch[1][:, :en],
                      batch[2][:, :en],
                      [x[:en] for x in batch[3]],
                  )
                  pred1 = np.argmax(trainer.predict(batch1), axis=2)
  
                  for j in range(num_sentences):
                      sentbreaks = np.where(
                          (pred1[j] == 2) + (pred1[j] == 4)
                      )[0]
                      if (
                          len(sentbreaks) <= 0
                          or idx[j] >= para_lengths[j] - max_seqlen
                      ):
                          advance = ens[j]
                      else:
                          advance = np.max(sentbreaks) + 1
  
                      pred[j] += [pred1[j, :advance]]
                      idx[j] += advance
                      adv[j] = advance
  
                  if all(
                      [idx1 >= N for idx1, N in zip(idx, para_lengths)]
                  ):
                      break
                  # once we've made predictions on a certain number of characters for each paragraph (recorded in `adv`),
                  # we skip the first `adv` characters to make the updated batch
                  batch = data_generator.advance_old_batch(adv, batch)
  
              pred = [np.concatenate(p, 0) for p in pred]
  
          for par_idx in range(num_sentences):
              offset = batch_idx * batch_size + par_idx
  
              raw = all_raw[offset]
              par_len = raw.index("<PAD>")
              raw = raw[:par_len]
              all_raw[offset] = raw
              if pred[par_idx][par_len - 1] < 2:
                  pred[par_idx][par_len - 1] = 2
              elif pred[par_idx][par_len - 1] > 2:
                  pred[par_idx][par_len - 1] = 4
              if use_regex_tokens:
                  all_preds.append(
                      update_pred_regex(raw, pred[par_idx][:par_len])
                  )
              else:
                  all_preds.append(pred[par_idx][:par_len])
  
      all_preds = sorted_data.unsort(all_preds)
      all_raw = sorted_data.unsort(all_raw)
  
      return all_preds, all_raw
  
imports: from collections import Counter
  from copy import copy
  import json
  import numpy
  import re
  import logging
  import os
  from torch.utils.data import DataLoader
  import stanza.utils.default_paths
  from stanza.models.common.utils import ud_scores, harmonic_mean
  from stanza.models.common.doc import Document
  from stanza.utils.conll import CoNLL
  from stanza.models.common.doc import *
  from stanza.models.tokenization.data import SortedDataset
label_desc: Use image node_id 8 for calling a global function with example usage: predict(trainer, data_generator, batch_size, max_seqlen, use_regex_tokens, num_workers) and returns: all_preds,  all_raw