I can't use someone else's model to generate inputs to mine, in a generator. How do I fix this?

I am trying to train a neural network on the SQuAD v1.1 dataset, using a pretrained BERT model. Someone suggested that I first grab the output of the BERT model and then feed those into my neural network as inputs. Due to the large amount of data, I feel that I need to create a generator, on which my neural network can then fit:

# @title Preparation
!pip install -q keras-bert
!wget -q https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip
!unzip -o uncased_L-12_H-768_A-12.zip

import os

pretrained_path = 'uncased_L-12_H-768_A-12'
config_path = os.path.join(pretrained_path, 'bert_config.json')
checkpoint_path = os.path.join(pretrained_path, 'bert_model.ckpt')
vocab_path = os.path.join(pretrained_path, 'vocab.txt')

# TF_KERAS must be added to environment variables in order to use TPU
os.environ['TF_KERAS'] = '1'
import codecs
from keras_bert import load_trained_model_from_checkpoint

token_dict = {}
with codecs.open(vocab_path, 'r', 'utf8') as reader:
    for line in reader:
        token = line.strip()
        token_dict[token] = len(token_dict)

model = load_trained_model_from_checkpoint(config_path, checkpoint_path)
import numpy as np
from keras_bert import Tokenizer
tokenizer = Tokenizer(token_dict)
def tokenize(text):
  tokens = tokenizer.tokenize(text)
  indices, segments = tokenizer.encode(first=text, max_len=512)
  return indices,segments
def feature_extraction(texts):
  return_values = []
  for text_ in texts:
    try:
      text_.split(" ")
    except AttributeError as e:
      raise TypeError("Expected array of strings.")
    try:
      indices,segments = tokenize(text_)
      predicts = model.predict([np.array([indices] * 8), np.array([segments] * 8)])[0]
      return_values.append(predicts)
    except ValueError as v:
      print(v)
  return_values = np.array(return_values)
  return return_values
print(feature_extraction(text_array).shape)
def batch_generator(dataframe,batch_size):
  while True:
    batch = dataframe.sample(n=batch_size)
    try:
      batch_features = feature_extraction(batch["question"].values)
    except ValueError as v:
      print("Oops, I'm getting a ValueError for batch_features.")
      print(v)
    try:
      batch_targets = batch["answer_start"]
    except ValueError as v:
      print("Oops, I'm getting a ValueError for batch_targets.")
      print(v)

    yield batch_features,batch_targets

This works when I feed it test data:

def batch_generator(dataframe,batch_size):
  while True:
    batch = dataframe.sample(n=batch_size)
    try:
      batch_features = feature_extraction(batch["question"].values)
    except ValueError as v:
      print("Oops, I'm getting a ValueError for batch_features.")
      print(v)
    try:
      batch_targets = batch["answer_start"]
    except ValueError as v:
      print("Oops, I'm getting a ValueError for batch_targets.")
      print(v)

    yield batch_features,batch_targets

This works when I use this test code:

testDataframe = pd.DataFrame({"question":["Does she sell seashells by the seashore?"],"answer":["She sells seashells by the seashore"],"answer_start":[0]})
for x,y in batch_generator(testDataframe,1):
  print (x)
  print (y)
  break

Output:

[[[-0.11251544 -0.09277309 0.04996187 ... -0.43535435 0.23852573 0.3206718 ] [ 0.35688528 0.43881682 -0.1390086 ... -0.32458037 0.64422214 -0.11743623] [ 0.6213926 -0.9945548 0.07564903 ... -0.87357795 0.2069801 -0.25303575] ... [-0.06796454 -0.24819699 -0.25508618 ... 0.20477912 0.36703664 0.04691853] [ 0.15030818 -0.05989693 0.17198643 ... 0.19960165 0.0324061 -0.31075317] [ 0.05091426 -0.14167279 0.18194658 ... 0.12112649 0.05029908 -0.15253511]]] 0 0 Name: answer_start, dtype: int64

I create and compile my neural network and inputs like so:

import json
import re
#regex = re.compile(r'\W+')
def readFile(filename):
  with open(filename) as file:
    fields = []
    JSON = json.loads(file.read())
    articles = []
    for article in JSON["data"]:
      articleTitle = article["title"]
      article_body = []
      for paragraph in article["paragraphs"]:
        paragraphContext = paragraph["context"]
        article_body.append(paragraphContext)
        for qas in paragraph["qas"]:
          question = qas["question"]
          answer = qas["answers"][0]

      article_body = "\\n".join(article_body)
      article = {"title":articleTitle,"body":article_body}
      articles.append(article)
  fields = pd.DataFrame(fields)
  #fields["question"] = fields["question"].str.replace(regex," ")
  assert not (fields["question"].str.contains("catalanswhat").any())
  #fields["paragraph_context"] = fields["paragraph_context"].str.replace(regex," ")
  #fields["answer_text"] = fields["answer_text"].str.replace(regex," ")
  assert not (fields["paragraph_context"].str.contains("catalanswhat").any())
  fields["article_title"] = fields["article_title"].str.replace("_"," ")
  assert not (fields["article_title"].str.contains("catalanswhat").any())
  return fields

trainingData = readFile("train-v1.1.json")

answers_network = Sequential()
answers_network.add(Dense(32,input_shape=(512,768)))
answers_network.summary()
answers_network.compile("rmsprop","categorical_crossentropy")
answers_network_checkpoint = ModelCheckpoint('answers_network-rnn-best.h5', verbose=1, monitor='val_loss',save_best_only=True, mode='auto')
answers_network.fit_generator(batch_generator(trainingData[["question","paragraph_context","answer_start"]],100),steps_per_epoch=8)

This fails with an error:

Tensor Input-Token:0, specified in either feed_devices or fetch_devices was not found in the Graph

Now, Input-Token is the name of one of the input layers in the BERT model.

I think TensorFlow is suggesting that the BERT model is using a different graph than my model.

Apparently, the BERT model uses custom layers and activation functions, so making a deep copy of the model may not be the best course of action.

What should I do?

from I can't use someone else's model to generate inputs to mine, in a generator. How do I fix this?