def test_container_feature_lookup(): arbitrary_attribute = "other" messages = [ Message(data={TEXT: "A"}, features=[_dummy_features(1, TEXT)]), Message( data={ INTENT: "B", arbitrary_attribute: "C" }, features=[_dummy_features(2, arbitrary_attribute)], ), Message(data={TEXT: "A2"}, features=[_dummy_features(3, TEXT)]), Message( data={ INTENT: "B2", arbitrary_attribute: "C2" }, features=[_dummy_features(4, arbitrary_attribute)], ), ] table = MessageContainerForCoreFeaturization() table.add_all(messages) # If we don't specify a list of attributes, the resulting features dictionary will # only contain those attributes for which there are features. sub_state = {TEXT: "A", INTENT: "B", arbitrary_attribute: "C"} features = table.collect_features(sub_state=sub_state) for attribute, feature_value in [ (TEXT, 1), (INTENT, None), (arbitrary_attribute, 2), ]: if feature_value is not None: assert attribute in features assert len(features[attribute]) == 1 assert feature_value == features[attribute][0].features[0] else: assert attribute not in features # If we query features for `INTENT`, then a key will be there, even if there are # no features features = table.collect_features(sub_state=sub_state, attributes=list(sub_state.keys())) assert INTENT in features assert len(features[INTENT]) == 0 # We only get the list of features we want... features = table.collect_features(sub_state, attributes=[arbitrary_attribute]) assert TEXT not in features assert INTENT not in features assert len(features[arbitrary_attribute]) == 1 # ... even if there are no features: YET_ANOTHER = "another" features = table.collect_features(sub_state, attributes=[YET_ANOTHER]) assert len(features[YET_ANOTHER]) == 0
def test_container_feature_lookup_fails_if_different_features_for_same_attribute(): broken_table = MessageContainerForCoreFeaturization() broken_table._table = { TEXT: {"A": Message(data={}, features=[_dummy_features(2, TEXT)])}, INTENT: {"B": Message(data={}, features=[_dummy_features(1, TEXT)])}, } with pytest.raises( RuntimeError, match=f"Feature for attribute {TEXT} has already been" ): broken_table.collect_features({TEXT: "A", INTENT: "B"})
def test_container_feature_lookup_works_if_messages_are_broken_but_consistent(): not_broken_but_strange_table = MessageContainerForCoreFeaturization() not_broken_but_strange_table._table = { TEXT: {"A": Message(data=dict())}, INTENT: {"B": Message(data=dict(), features=[_dummy_features(1, TEXT)])}, } features = not_broken_but_strange_table.collect_features({TEXT: "A", INTENT: "B"}) assert TEXT in features and len(features[TEXT]) == 1
def test_container_feature_lookup_fails_without_key_attribute(): table = MessageContainerForCoreFeaturization() with pytest.raises(ValueError, match="Unknown key"): table.collect_features({TEXT: "A-unknown"})