Model predict get label name

[Zjq9409]

training data as follows:
class,gender,age,fare,label
1,female,50.0,28.7125,died
1,female,2.0,151.55,died
1,female,25.0,151.55,died
2,female,29.0,26.0,survived
2,female,33.0,26.0,survived
2,female,36.0,26.0,survived
2,female,42.0,26.0,survived

test_data as follows:
class,gender,age,fare
3,male,24,7.1417
1,male,47,227

I use tfdf.keras.RandomForestModel predict test data:

output = model.predict(test_ds)

[[0.6399997 ]
 [0.13999999]]

How can I get specific category names 'died' or 'survived' according to probability?

[Zjq9409]

when I use sklearn RandomForestClassifier, the probability of two classifications will be output

prediction = model.predict_proba(vec.transform(data))
[[0.27861378 0.72138625]]
[[0.99693835 0.00306163]]

[Zjq9409]

adv_args = tfdf.keras.AdvancedArguments(
    predict_single_probability_for_binary_classification=False
)

model = tfdf.keras.GradientBoostedTreesModel(verbose=0, advanced_arguments=adv_args)

I solved this problem by this code.