print_example_preprocessor_config.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. #!/usr/bin/env python2
  2. # Copyright 2018 The Chromium Authors. All rights reserved.
  3. # Use of this source code is governed by a BSD-style license that can be
  4. # found in the LICENSE file.
  5. """Dumps info from a ExamplePreprocessorConfig protobuf file.
  6. Prints feature names, types, and bucket values.
  7. """
  8. import os
  9. import sys
  10. import textwrap
  11. from enum import Enum
  12. from google.protobuf import text_format
  13. class FeatureType(Enum):
  14. CATEGORICAL = 'categorical'
  15. BUCKETED = 'bucketed'
  16. SCALAR = 'scalar'
  17. def ReadConfig(pb_file):
  18. """Parses the protobuf containing the example preprocessor config."""
  19. import example_preprocessor_pb2
  20. config = example_preprocessor_pb2.ExamplePreprocessorConfig()
  21. with open(pb_file) as pb:
  22. config.ParseFromString(pb.read())
  23. return config
  24. def PrintExamplePreprocessorConfig(pb_file):
  25. """Prints the features listed the example preprocessor config."""
  26. config = ReadConfig(pb_file)
  27. features = set()
  28. for feature_index in sorted(config.feature_indices):
  29. # For string or string list feature types, remove the "_value" suffix to get
  30. # the base name.
  31. name_parts = feature_index.split('_')
  32. base_name = name_parts[0]
  33. # Skip additional values of the same base name.
  34. if base_name in features:
  35. continue
  36. features.add(base_name)
  37. if len(name_parts) == 1:
  38. feature_type = FeatureType.SCALAR
  39. elif base_name in config.bucketizers:
  40. feature_type = FeatureType.BUCKETED
  41. else:
  42. feature_type = FeatureType.CATEGORICAL
  43. description = '* %s (%s)' % (base_name, feature_type.value)
  44. if feature_type == FeatureType.BUCKETED:
  45. description += ':\n\t'
  46. boundaries = config.bucketizers[base_name].boundaries
  47. bucket_str = ', '.join(['%.1f' % bucket for bucket in boundaries])
  48. # Indent description by a tab and wrap text.
  49. max_len = 80 - 8 # Leave at least 8 columns for tab width.
  50. description += ('\n\t').join(textwrap.wrap(bucket_str, max_len))
  51. print description
  52. return 0
  53. def Main(args):
  54. if len(args) != 2:
  55. print 'Usage: %s <out_dir> <path/to/example_preprocessor_config.pb>' % (
  56. __file__)
  57. return 1
  58. out_dir = args[0]
  59. if not os.path.isdir(out_dir):
  60. print 'Could not find out directory: %s' % out_dir
  61. return 1
  62. pb_file = args[1]
  63. if not os.path.isfile(pb_file):
  64. print 'Protobuf file not found: %s' % pb_file
  65. return 1
  66. proto_dir = os.path.join(out_dir, 'pyproto/components/assist_ranker/proto')
  67. if not os.path.isdir(proto_dir):
  68. print 'Proto directory not found: %s' % proto_dir
  69. print 'Build the "components/assist_ranker/proto" target'
  70. print ' (usually built with chrome)'
  71. return 1
  72. # Allow importing the ExamplePreprocessorConfig proto definition.
  73. sys.path.insert(0, proto_dir)
  74. PrintExamplePreprocessorConfig(pb_file)
  75. if __name__ == '__main__':
  76. sys.exit(Main(sys.argv[1:]))