class LearnBinarySearchTest(parameterized.TestCase):
@parameterized.named_parameters(
{
'testcase_name':
'ReachesVocabSize',
'word_counts': [('apple', 2), ('peach', 1), ('pear', 1)],
'lower':
1,
'upper':
10,
'delta':
0,
'expected_vocab':
['a', 'c', 'e', 'h', 'l', 'p', 'r', 'apple', 'peach', 'pear'],
'params':
learner.Params(upper_thresh=4,
lower_thresh=1,
num_iterations=4,
max_input_tokens=1000,
max_token_length=50,
max_unique_chars=50,
vocab_size=10,
slack_ratio=0,
include_joiner_token=False,
joiner='##',
reserved_tokens=[])
}, {
'testcase_name':
'VocabSizeWithinSlack',
'word_counts': [('apple', 2), ('peach', 1), ('pear', 1),
('app', 2)],
'lower':
1,
'upper':
10,
'delta':
6,
'expected_vocab': ['a', 'c', 'e', 'h', 'l', 'p', 'r'],
'params':
learner.Params(upper_thresh=4,
lower_thresh=1,
num_iterations=4,
max_input_tokens=1000,
max_token_length=50,
max_unique_chars=50,
vocab_size=12,
slack_ratio=0.5,
include_joiner_token=False,
joiner='##',
reserved_tokens=[])
})
def testBinarySearch(self, word_counts, lower, upper, delta,
expected_vocab, params):
vocab = learner.learn_binary_search(word_counts, lower, upper, params)
self.assertAlmostEqual(len(vocab), params.vocab_size, delta=delta)
self.assertLessEqual(len(vocab), params.vocab_size)
self.assertEqual(vocab, expected_vocab)
def main(_):
# Define schema.
raw_metadata = dataset_metadata.DatasetMetadata(
dataset_schema.from_feature_spec({
'text':
tf.FixedLenFeature([], tf.string),
'language_code':
tf.FixedLenFeature([], tf.string),
}))
# Add in padding tokens.
reserved_tokens = FLAGS.reserved_tokens
if FLAGS.num_pad_tokens:
padded_tokens = ['<pad>']
padded_tokens += [
'<pad%d>' % i for i in range(1, FLAGS.num_pad_tokens)
]
reserved_tokens = padded_tokens + reserved_tokens
params = learner.Params(FLAGS.upper_thresh, FLAGS.lower_thresh,
FLAGS.num_iterations, FLAGS.max_input_tokens,
FLAGS.max_token_length, FLAGS.max_unique_chars,
FLAGS.vocab_size, FLAGS.slack_ratio,
FLAGS.include_joiner_token, FLAGS.joiner,
reserved_tokens)
generate_vocab(FLAGS.data_file, FLAGS.vocab_file, FLAGS.metrics_file,
raw_metadata, params)
def main(_):
# Read in wordcount file.
with open(FLAGS.input_path) as wordcount_file:
word_counts = [(line.split()[0], int(line.split()[1]))
for line in wordcount_file]
# Add in padding tokens.
reserved_tokens = FLAGS.reserved_tokens
if FLAGS.num_pad_tokens:
padded_tokens = ['<pad>']
padded_tokens += ['<pad%d>' % i for i in range(1, FLAGS.num_pad_tokens)]
reserved_tokens = padded_tokens + reserved_tokens
params = learner.Params(FLAGS.upper_thresh, FLAGS.lower_thresh,
FLAGS.num_iterations, FLAGS.max_input_tokens,
FLAGS.max_token_length, FLAGS.max_unique_chars,
FLAGS.vocab_size, FLAGS.slack_ratio,
FLAGS.include_joiner_token, FLAGS.joiner,
reserved_tokens)
vocab = learner.learn(word_counts, params)
vocab = ''.join([line + '\n' for line in vocab])
# Write vocab to file.
with open(FLAGS.output_path, 'w') as vocab_file:
vocab_file.write(vocab)
class LearnWithThreshTest(parameterized.TestCase):
@parameterized.named_parameters(
{
'testcase_name':
'LearnWithOneIteration',
'word_counts': [('apple', 1), ('app', 1)],
'thresh':
1,
'expected_vocab':
['a', 'e', 'l', 'p', 'app', 'apple', 'le', 'ple', 'pp', 'pple'],
'params':
learner.Params(upper_thresh=4,
lower_thresh=1,
num_iterations=1,
max_input_tokens=1000,
max_token_length=50,
max_unique_chars=5,
vocab_size=10,
slack_ratio=0,
include_joiner_token=False,
joiner='##',
reserved_tokens=[])
}, {
'testcase_name':
'LearnWithTwoIterations',
'word_counts': [('apple', 1), ('app', 1)],
'thresh':
1,
'expected_vocab': ['a', 'e', 'l', 'p', 'app', 'apple'],
'params':
learner.Params(upper_thresh=4,
lower_thresh=1,
num_iterations=2,
max_input_tokens=1000,
max_token_length=50,
max_unique_chars=5,
vocab_size=10,
slack_ratio=0,
include_joiner_token=False,
joiner='##',
reserved_tokens=[])
}, {
'testcase_name':
'LearnWithHigherThresh',
'word_counts': [('apple', 1), ('app', 2)],
'thresh':
2,
'expected_vocab': ['a', 'e', 'l', 'p', 'app', 'pp'],
'params':
learner.Params(upper_thresh=4,
lower_thresh=1,
num_iterations=1,
max_input_tokens=1000,
max_token_length=50,
max_unique_chars=5,
vocab_size=10,
slack_ratio=0,
include_joiner_token=False,
joiner='##',
reserved_tokens=[])
})
def testLearnWithThresh(self, word_counts, thresh, expected_vocab, params):
vocab = learner.learn_with_thresh(word_counts, thresh, params)
self.assertEqual(vocab, expected_vocab)
