def _collect_sparse_feature_sizes(
featurized_example: Dict[Text, List["Features"]],
training_example: Message,
featurizers: Optional[List[Text]] = None,
) -> Dict[Text, Dict[Text, List[int]]]:
"""Collects sparse feature sizes for all attributes that have sparse features.
Returns sparse feature sizes for each attribute. It could look like this:
{TEXT: {FEATURE_TYPE_SEQUENCE: [16, 32], FEATURE_TYPE_SENTENCE: [16, 32]}}.
Args:
featurized_example: a featurized example
training_example: a training example
featurizers: the featurizers to consider
Returns:
A dictionary of attribute to feature sizes.
"""
sparse_feature_sizes = {}
sparse_attributes = []
for attribute, features in featurized_example.items():
if features and features[0].is_sparse():
sparse_attributes.append(attribute)
for attribute in sparse_attributes:
sparse_feature_sizes[
attribute] = training_example.get_sparse_feature_sizes(
attribute=attribute, featurizers=featurizers)
return sparse_feature_sizes
def test_get_sparse_feature_sizes(
features: Optional[List[Features]],
attribute: Text,
featurizers: List[Text],
expected_sequence_sizes: List[int],
expected_sentence_sizes: List[int],
):
message = Message(data={TEXT: "This is a test sentence."}, features=features)
feature_sizes = message.get_sparse_feature_sizes(attribute, featurizers)
assert feature_sizes[FEATURE_TYPE_SEQUENCE] == expected_sequence_sizes
assert feature_sizes[FEATURE_TYPE_SENTENCE] == expected_sentence_sizes
