docstr: Predict lemmas of words for a given (list of) sentence(s).
  
  Args:
      sentences: sentences to predict
      label_name: label name used for predicted lemmas
      mini_batch_size: number of tokens that are send through the RNN simultaneously, assuming batching_in_rnn is set to True
      embedding_storage_mode: default is 'none' which is always best. Only set to 'cpu' or 'gpu' if you wish to not only predict, but also keep the generated embeddings in CPU or GPU memory respectively.
      return_loss: whether to compute and return loss. Setting it to True only makes sense if labels are provided
      verbose: If True, lemmatized sentences will be printed in the console.
      return_probabilities_for_all_classes: unused parameter.
docstr_tok: ["Predict","lemmas","of","words","for","a","given","(","list","of",")","sentence","(","s",")",".","Args",":","sentences",":","sentences","to","predict","label_name",":","label","name","used","for","predicted","lemmas","mini_batch_size",":","number","of","tokens","that","are","send","through","the","RNN","simultaneously",",","assuming","batching_in_rnn","is","set","to","True","embedding_storage_mode",":","default","is","'none","'","which","is","always","best",".","Only","set","to","'cpu","'","or","'gpu","'","if","you","wish","to","not","only","predict",",","but","also","keep","the","generated","embeddings","in","CPU","or","GPU","memory","respectively",".","return_loss",":","whether","to","compute","and","return","loss",".","Setting","it","to","True","only","makes","sense","if","labels","are","provided","verbose",":","If","True",",","lemmatized","sentences","will","be","printed","in","the","console",".","return_probabilities_for_all_classes",":","unused","parameter","."]
code: def predict(
      self,
      sentences: Union[List[Sentence], Sentence],
      mini_batch_size: int = 16,
      return_probabilities_for_all_classes: bool = False,
      verbose: bool = False,
      label_name="predicted",
      return_loss=False,
      embedding_storage_mode="none",
  ):
      """Predict lemmas of words for a given (list of) sentence(s).
  
      Args:
          sentences: sentences to predict
          label_name: label name used for predicted lemmas
          mini_batch_size: number of tokens that are send through the RNN simultaneously, assuming batching_in_rnn is set to True
          embedding_storage_mode: default is 'none' which is always best. Only set to 'cpu' or 'gpu' if you wish to not only predict, but also keep the generated embeddings in CPU or GPU memory respectively.
          return_loss: whether to compute and return loss. Setting it to True only makes sense if labels are provided
          verbose: If True, lemmatized sentences will be printed in the console.
          return_probabilities_for_all_classes: unused parameter.
      """
      if isinstance(sentences, Sentence):
          sentences = [sentences]
  
      Sentence.set_context_for_sentences(sentences)
  
      # filter empty sentences
      sentences = [
          sentence for sentence in sentences if len(sentence) > 0
      ]
      if len(sentences) == 0:
          return sentences
  
      # max length of the predicted sequences
      if not self.dependent_on_input:
          max_length = self.max_sequence_length
      else:
          max_length = max(
              [
                  len(token.text) + 1
                  for sentence in sentences
                  for token in sentence
              ]
          )
  
      # for printing
      line_to_print = ""
  
      overall_loss = 0.0
      number_tokens_in_total = 0
  
      with torch.no_grad():
          dataloader = DataLoader(
              dataset=FlairDatapointDataset(sentences),
              batch_size=mini_batch_size,
          )
  
          for batch in dataloader:
              # stop if all sentences are empty
              if not batch:
                  continue
  
              # remove previously predicted labels of this type
              for sentence in batch:
                  for token in sentence:
                      token.remove_labels(label_name)
  
              # create list of tokens in batch
              tokens_in_batch = [
                  token for sentence in batch for token in sentence
              ]
              number_tokens = len(tokens_in_batch)
              number_tokens_in_total += number_tokens
  
              # encode inputs
              hidden, all_encoder_outputs = self.encode(batch)
  
              # create input for first pass (batch_size, 1, input_size), first letter is special character <S>
              # sequence length is always set to one in prediction
              input_indices = self.start_index * torch.ones(
                  number_tokens, dtype=torch.long, device=flair.device
              ).unsqueeze(1)
  
              # option 1: greedy decoding
              if self.beam_size == 1:
                  # predictions
                  predicted: List[List[int]] = [
                      [] for _ in range(number_tokens)
                  ]
  
                  for _decode_step in range(max_length):
                      # decode next character
                      output_vectors, hidden = self.decode(
                          input_indices, hidden, all_encoder_outputs
                      )
  
                      log_softmax_probs = (
                          torch.nn.functional.log_softmax(
                              output_vectors, dim=2
                          )
                      )
                      # pick top beam size many outputs with highest probabilities
                      input_indices = log_softmax_probs.argmax(dim=2)
  
                      for i in range(number_tokens):
                          if (
                              len(predicted[i]) > 0
                              and predicted[i][-1] == self.end_index
                          ):
                              continue
                          predicted[i].append(input_indices[i].item())
  
                  for t_id, token in enumerate(tokens_in_batch):
                      predicted_lemma = "".join(
                          self.char_dictionary.get_item_for_index(idx)
                          if idx != self.end_index
                          else ""
                          for idx in predicted[t_id]
                      )
                      token.set_label(
                          typename=label_name, value=predicted_lemma
                      )
  
              # option 2: beam search
              else:
                  output_vectors, hidden = self.decode(
                      input_indices, hidden, all_encoder_outputs
                  )
  
                  # out_probs = self.softmax(output_vectors).squeeze(1)
                  log_softmax_probs = torch.nn.functional.log_softmax(
                      output_vectors, dim=2
                  ).squeeze(1)
                  # make sure no dummy symbol <> or start symbol <S> is predicted
                  log_softmax_probs[:, self.dummy_index] = -inf
                  log_softmax_probs[:, self.start_index] = -inf
  
                  # pick top beam size many outputs with highest probabilities
                  # probabilities, leading_indices = out_probs.topk(self.beam_size, 1)  # max prob along dimension 1
                  (
                      log_probabilities,
                      leading_indices,
                  ) = log_softmax_probs.topk(self.beam_size, 1)
                  # leading_indices and probabilities have size (batch_size, beam_size)
  
                  # keep scores of beam_size many hypothesis for each token in the batch
                  scores = log_probabilities.view(-1, 1)
                  # stack all leading indices of all hypothesis and corresponding hidden states in two tensors
                  leading_indices = leading_indices.view(
                      -1, 1
                  )  # this vector goes through RNN in each iteration
  
                  hidden_states_beam = torch.stack(
                      self.beam_size * [hidden], dim=2
                  ).view(self.rnn_layers, -1, self.rnn_hidden_size)
  
                  # save sequences so far
                  sequences = torch.tensor(
                      [[i.item()] for i in leading_indices],
                      device=flair.device,
                  )
  
                  # keep track of how many hypothesis were completed for each token
                  n_completed = [0 for _ in range(number_tokens)]  # cpu
                  final_candidates: List[
                      List[Tuple[torch.Tensor, float]]
                  ] = [
                      [] for _ in range(number_tokens)
                  ]  # cpu
  
                  # if all_encoder_outputs returned, expand them to beam size (otherwise keep this as None)
                  batched_encoding_output = (
                      torch.stack(
                          self.beam_size * [all_encoder_outputs], dim=1
                      ).view(
                          self.beam_size * number_tokens,
                          -1,
                          self.rnn_hidden_size,
                      )
                      if self.use_attention
                      else None
                  )
  
                  for _j in range(1, max_length):
                      output_vectors, hidden_states_beam = self.decode(
                          leading_indices,
                          hidden_states_beam,
                          batched_encoding_output,
                      )
  
                      # decode with log softmax
                      out_log_probs = torch.nn.functional.log_softmax(
                          output_vectors, dim=2
                      )
                      # make sure no dummy symbol <> or start symbol <S> is predicted
                      out_log_probs[:, 0, self.dummy_index] = -inf
                      out_log_probs[:, 0, self.start_index] = -inf
                      (
                          log_probabilities,
                          index_candidates,
                      ) = out_log_probs.topk(self.beam_size, 2)
                      log_probabilities.squeeze_(1)
                      index_candidates.squeeze_(1)
  
                      # check if an end symbol <E> has been predicted and, in that case, set hypothesis aside
                      end_symbols = (
                          index_candidates == self.end_index
                      ).nonzero(as_tuple=False)
                      for tuple in end_symbols:
                          # if the sequence is already ended, do not record as candidate
                          if (
                              sequences[tuple[0], -1].item()
                              == self.end_index
                          ):
                              continue
  
                          # index of token in in list tokens_in_batch
                          token_number = torch.div(
                              tuple[0],
                              self.beam_size,
                              rounding_mode="trunc",
                          )
                          # print(token_number)
                          seq = sequences[
                              tuple[0], :
                          ]  # hypothesis sequence
                          # hypothesis score
                          score = (
                              scores[tuple[0]]
                              + log_probabilities[tuple[0], tuple[1]]
                          ) / (len(seq) + 1)
  
                          final_candidates[token_number].append(
                              (seq, score.item())
                          )
                          # TODO: remove token if number of completed hypothesis exceeds given value
                          n_completed[token_number] += 1
  
                          # set score of corresponding entry to -inf so it will not be expanded
                          log_probabilities[tuple[0], tuple[1]] = -inf
  
                      # get leading_indices for next expansion
                      # find highest scoring hypothesis among beam_size*beam_size possible ones for each token
  
                      # take beam_size many copies of scores vector and add scores of possible new extensions
                      # size (beam_size*batch_size, beam_size)
                      hypothesis_scores = (
                          torch.cat(self.beam_size * [scores], dim=1)
                          + log_probabilities
                      )
                      # print(hypothesis_scores)
  
                      # reshape to vector of size (batch_size, beam_size*beam_size),
                      # each row contains beam_size*beam_size scores of the new possible hypothesis
                      hypothesis_scores_per_token = (
                          hypothesis_scores.view(
                              number_tokens, self.beam_size**2
                          )
                      )
                      # print(hypothesis_scores_per_token)
  
                      # choose beam_size best for each token - size (batch_size, beam_size)
                      (
                          best_scores,
                          indices_per_token,
                      ) = hypothesis_scores_per_token.topk(
                          self.beam_size, 1
                      )
  
                      # out of indices_per_token we now need to recompute the original indices of the hypothesis in
                      # a list of length beam_size*batch_size
                      # where the first three inidices belong to the first token, the next three to the second token,
                      # and so on
                      beam_numbers: List[int] = []
                      seq_numbers: List[int] = []
  
                      for i, row in enumerate(indices_per_token):
                          beam_numbers.extend(
                              i * self.beam_size
                              + index.item() // self.beam_size
                              for index in row
                          )
  
                          seq_numbers.extend(
                              index.item() % self.beam_size
                              for index in row
                          )
  
                      # with these indices we can compute the tensors for the next iteration
                      # expand sequences with corresponding index
                      sequences = torch.cat(
                          (
                              sequences[beam_numbers],
                              index_candidates[
                                  beam_numbers, seq_numbers
                              ].unsqueeze(1),
                          ),
                          dim=1,
                      )
  
                      # add log-probabilities to the scores
                      scores = scores[beam_numbers] + log_probabilities[
                          beam_numbers, seq_numbers
                      ].unsqueeze(1)
  
                      # save new leading indices
                      leading_indices = index_candidates[
                          beam_numbers, seq_numbers
                      ].unsqueeze(1)
  
                      # save corresponding hidden states
                      hidden_states_beam = hidden_states_beam[
                          :, beam_numbers, :
                      ]
  
                  # it may happen that no end symbol <E> is predicted for a token in all of the max_length iterations
                  # in that case we append one of the final seuqences without end symbol to the final_candidates
                  best_scores, indices = scores.view(
                      number_tokens, -1
                  ).topk(1, 1)
  
                  for j, (score, index) in enumerate(
                      zip(best_scores.squeeze(1), indices.squeeze(1))
                  ):
                      if len(final_candidates[j]) == 0:
                          beam = j * self.beam_size + index.item()
                          final_candidates[j].append(
                              (
                                  sequences[beam, :],
                                  score.item() / max_length,
                              )
                          )
  
                  # get best final hypothesis for each token
                  output_sequences = []
                  for candidate in final_candidates:
                      l_ordered = sorted(
                          candidate,
                          key=lambda tup: tup[1],
                          reverse=True,
                      )
                      output_sequences.append(l_ordered[0])
  
                  # get characters from index sequences and add predicted label to token
                  for i, out_seq in enumerate(output_sequences):
                      predicted_lemma = ""
                      for idx in out_seq[0]:
                          predicted_lemma += (
                              self.char_dictionary.get_item_for_index(
                                  idx
                              )
                          )
                      line_to_print += predicted_lemma
                      line_to_print += " "
                      tokens_in_batch[i].add_tag(
                          tag_type=label_name, tag_value=predicted_lemma
                      )
  
              if return_loss:
                  overall_loss += self.forward_loss(batch)[0].item()
  
              store_embeddings(
                  batch, storage_mode=embedding_storage_mode
              )
  
          if verbose:
              log.info(line_to_print)
  
          if return_loss:
              return overall_loss, number_tokens_in_total
          return None
  
imports: import logging
  from math import inf
  from typing import List, Optional, Tuple, Union
  import torch
  from torch import nn
  import flair.embeddings
  import flair.nn
  from flair.data import Dictionary, Sentence, Token
  from flair.datasets import DataLoader, FlairDatapointDataset
  from flair.training_utils import Result, store_embeddings
label_desc: Use image node_id 12 for calling the Lemmatizer obj's underlying member method code with example usage: obj.predict(sentences, mini_batch_size, return_probabilities_for_all_classes, verbose, label_name, return_loss, embedding_storage_mode) and returns: sentences, None, overall_loss,  number_tokens_in_total