Ошибка запроса () модуля pgmpy: IndexError: индекс 2 выходит за пределы оси 0 с размером 2

Question

Я пытаюсь использовать функцию запроса в модуле pgmpy для определения условной вероятности байесовской сети. Мой код выглядит следующим образом:

data_test = pd.DataFrame(np.random.randint(1, 3, size = (240, 5)),columns=['A','B','C','D','E'])

best_model.fit(data_test)
bayes_i = VariableElimination(best_model)
# find out the probability of B given other parameters' values
q = bayes_i.query(variables=['B'], evidence={'A': 1,'C': 2,'D': 2,'E': 1})
print (q['B'])
q = bayes_i.query(variables=['B'], evidence={'A': 2,'C': 1,'D': 1,'E': 1})
print (q['B'])

сообщение об ошибке: IndexError: индекс 2 выходит за пределы оси 0 с размером 2

редактировать: модель построена следующим образом:

def Bayes(dataset):
    start = time.time()
    # exhuastive search and find the set of the best DAG
    searcher = ExhaustiveSearch(dataset, scoring_method=K2Score(dataset))
    model_list = []
    for score, model in searcher.all_scores():
        model_list.append(model.edges())
    model = BayesianModel(model_list[0])
    end = time.time()
    print('Time used to construct the model is: ' + str(end - start))
    return model

best_model = Bayes(data_random)