Keras vertical ensemble model with condition in between

Question

I have trained two separate models

• ModelA: Checks if the input text is related to my work (Binary Classifier [related/not-related])
• ModelB: Classifier of related texts (Classifier [good/normal/bad]). Only the related texts are relayed to this model from ModelA

I want

• ModelC: Ensemble classifier that outputs [good/normal/bad/not-related]
• I'll be training in batches. And there can be mix of not-related and good/normal/bad in one batch. I need them separated.

Some pseudo code of what i need

# Output of modelA will be a vector I presume `(1, None)` where `None` is batch
def ModelC.predict(input):
    outputA = ModelA(input)
    if outputA == 'not-related':
        return outputA
    return ModelB(outputA)

I don't know how to include the if logic inside the models inference. How can I achieve this?

Answer 1

Just define your own model. I'm surprised your other models are outputting strings instead of numbers, but without more info this is about all I can give you, so I will assume the output of model A is a string.

import tensorflow as tf

class ModelC(tf.keras.Model):

  def __init__(self, A, B):
    super(ModelC, self).__init__()
    self.A = A
    self.B = B

  def call(self, inputs, training=False):
    x = self.A(inputs, training)
    if x == 'not-related':
        return x
    return self.B(inputs, training)