[Python] How to apply different model on different rows of a pandas dataframe?

I have a pandas dataframe that looks like this:

import pandas as pd
df = pd.DataFrame({'id': [1,2], 'var1': [5,6], 'var2': [20,60], 'var3': [8, -2], 'model_version': ['model_a', 'model_b']})


I have 2 models different models, saved in pkl files, which I load them like this:

import pickle
with open('model_a.pkl', 'rb') as file:
model_a = pickle.load(file)
with open('model_b.pkl', 'rb') as file:
model_b = pickle.load(file)


I would like to apply and predict model_a for id==1 and model_b for id==2, as indicated in the model_version column in the df.

How can I do that in "one-go" ?

What I mean in "one-go" is that: I am NOT looking for a solution that looks like this:

df_a = df.query('model_version=="model_a"')
predictions_a = model_a.predict(df_a)

df_b = df.query('model_version=="model_b"')
predictions_b = model_b.predict(df_b)

Continue reading...