def test_create_lit_model_from_endpoint_name_returns_model(
self, feature_types, label_types, model_id):
lit_model = create_lit_model_from_endpoint(_TEST_ENDPOINT_NAME,
feature_types, label_types,
model_id)
test_inputs = [
{
"feature_1": 1.0,
"feature_2": 2.0
},
]
outputs = lit_model.predict_minibatch(test_inputs)
assert lit_model.input_spec() == dict(feature_types)
assert lit_model.output_spec() == dict(label_types)
assert len(outputs) == 1
for item in outputs:
assert item.keys() == {"label"}
assert len(item.values()) == 1
def test_create_lit_model_from_endpoint_name_with_xai_returns_model(
self, feature_types, label_types, model_id):
lit_model = create_lit_model_from_endpoint(_TEST_ENDPOINT_NAME,
feature_types, label_types,
model_id)
test_inputs = [
{
"feature_1": 1.0,
"feature_2": 2.0
},
]
outputs = lit_model.predict_minibatch(test_inputs)
assert lit_model.input_spec() == dict(feature_types)
assert lit_model.output_spec() == dict({
**label_types,
"feature_attribution":
lit_types.FeatureSalience(signed=True),
})
assert len(outputs) == 1
for item in outputs:
assert item.keys() == {"label", "feature_attribution"}
assert len(item.values()) == 2
