Использование генераторов с функцией tf.py_function

Question

Как передать функцию генератора в tf.py_function?

def generate_context_target_pairs(words, window_size, vocab_size):
  context_length = window_size*2
  sentence_length = len(words)
  for index, word in enumerate(words):
      context_words = []
      label_word   = []            
      start = index - window_size
      end = index + window_size + 1

      context_words.append([words[i] 
                            for i in range(start, end) 
                            if 0 <= i < sentence_length 
                            and i != index])
      label_word.append(word)

      x = sequence.pad_sequences(context_words, maxlen=context_length,
                                  padding = 'post', truncating = 'post',value = encoder.encode("<UNK>"))
      y = to_categorical(label_word, vocab_size) #creates a one hot vector for each sequence in the list
      return (tf.convert_to_tensor(x, dtype=tf.float32), tf.convert_to_tensor(y, dtype=tf.float32))

def generate_context_target_pairs_map_fxn(word_list, window_size, vocab_size):
  # py_func doesn't set the shape of the returned tensors.
  encoded_text, label = tf.py_function(generate_context_target_pairs, 
                                       inp=[word_list, window_size, vocab_size], 
                                       Tout=(tf.float32, tf.float32))
  return encoded_text, label

Если в приведенной выше функции я хочу передать генератор вместо функции python, возможно ли это? поэтому, пожалуйста, поделитесь примером. Было бы очень полезно, я просмотрел официальную документацию, но не смог найти ничего конкретного.