def word_count(input_path, output_path, raw_metadata, min_token_frequency=2):
"""Returns a pipeline counting words and writing the output.
Args:
input_path: recordio file to read
output_path: path in which to write the output
raw_metadata: metadata of input tf.Examples
min_token_frequency: the min frequency for a token to be included
"""
lang_set = set(FLAGS.lang_set.split(','))
# Create pipeline.
pipeline = beam.Pipeline()
with tft_beam.Context(temp_dir=tempfile.mkdtemp()):
converter = tft.coders.ExampleProtoCoder(raw_metadata.schema,
serialized=False)
# Read raw data and convert to TF Transform encoded dict.
raw_data = (
pipeline
| 'ReadInputData' >> beam.io.tfrecordio.ReadFromTFRecord(
input_path, coder=beam.coders.ProtoCoder(tf.train.Example))
| 'DecodeInputData' >> beam.Map(converter.decode))
# Apply TF Transform.
(transformed_data, _), _ = (
(raw_data, raw_metadata)
|
'FilterLangAndExtractToken' >> tft_beam.AnalyzeAndTransformDataset(
utils.count_preprocessing_fn(FLAGS.text_key,
FLAGS.language_code_key)))
# Filter by languages.
tokens = (
transformed_data
| 'FilterByLang' >> beam.ParDo(utils.FilterTokensByLang(lang_set)))
# Calculate smoothing coefficients.
coeffs = (tokens
| 'CalculateSmoothingCoefficients' >> beam.CombineGlobally(
utils.CalculateCoefficients(FLAGS.smoothing_exponent)))
# Apply smoothing, aggregate counts, and sort words by count.
_ = (tokens
| 'ApplyExponentialSmoothing' >> beam.ParDo(
utils.ExponentialSmoothing(), beam.pvalue.AsSingleton(coeffs))
| 'SumCounts' >> beam.CombinePerKey(sum)
| 'FilterLowCounts' >> beam.ParDo(
utils.FilterByCount(FLAGS.max_word_length,
min_token_frequency))
|
'MergeAndSortCounts' >> beam.CombineGlobally(utils.SortByCount())
| 'Flatten' >> beam.FlatMap(lambda x: x)
| 'FormatCounts' >> beam.Map(lambda tc: '%s\t%s' % (tc[0], tc[1]))
| 'WriteSortedCount' >> beam.io.WriteToText(
output_path, shard_name_template=''))
return pipeline
def run_vocab():
"""Creates a pipeline to generate wordpiece vocab over a corpus."""
vocab_pipeline = beam.Pipeline()
with tft_beam.Context(temp_dir=tempfile.mkdtemp()):
# Read raw data and convert to TF Transform encoded dict.
raw_data = (
vocab_pipeline
| 'ReadInputData' >> beam.io.tfrecordio.ReadFromTFRecord(
data_file, coder=beam.coders.ProtoCoder(tf.train.Example))
| 'DecodeInputData' >> beam.Map(example_converter.decode))
# Apply TF Transform.
(transformed_data,
_), _ = ((raw_data, raw_metadata)
| 'FilterLangAndExtractToken' >>
tft_beam.AnalyzeAndTransformDataset(
utils.count_preprocessing_fn(
FLAGS.text_key, FLAGS.language_code_key)))
# Filter by languages.
tokens = (transformed_data
| 'FilterByLang' >> beam.ParDo(
utils.FilterTokensByLang(lang_set)))
# Calculate smoothing coefficients.
coeffs = (
tokens
| 'CalculateSmoothingCoefficients' >> beam.CombineGlobally(
utils.CalculateCoefficients(FLAGS.smoothing_exponent)))
# Apply smoothing, aggregate counts, and sort words by count.
_ = (tokens
| 'ApplyExponentialSmoothing' >> beam.ParDo(
utils.ExponentialSmoothing(),
beam.pvalue.AsSingleton(coeffs))
| 'SumCounts' >> beam.CombinePerKey(sum)
| 'FilterLowCounts' >> beam.ParDo(
utils.FilterByCount(FLAGS.max_word_length,
min_token_frequency))
| 'MergeAndSortCounts' >> beam.CombineGlobally(
utils.SortByCount())
| 'LearnVocab' >> beam.ParDo(utils.LearnVocab(params))
| 'Flatten' >> beam.FlatMap(lambda x: x + '\n')
| 'WriteVocab' >> beam.io.WriteToText(
vocab_file,
shard_name_template='',
append_trailing_newlines=False))
return vocab_pipeline
def testUseGivenLang(self):
preprocessing_fn = utils.count_preprocessing_fn(
'text', 'language_code')
with tf.Session() as sess:
expected_tokens = [
'Let', '\'', 's', 'make', 'this', 'Chinese', 'even', 'though',
'it', '\'', 's', 'English', '.'
]
outputs = preprocessing_fn(self.raw_data)
outputs = sess.run(outputs)
self.assertEqual(outputs['lang'], 'zh')
self.assertSequenceAlmostEqual(outputs['tokens'].values,
expected_tokens)