def test_partition_with_max_labels_discards_labels(self):
mock_label_reader = labelled_data.Reader(
io.StringIO("""
input,name,qty,range_end,unit,comment
1 cup foo,foo,1.0,0.0,cup,
2 drops foz,foz,2.0,0.0,drop,
3 ml faa,faa,3.0,0.0,ml,
4 cloves bar,bar,4.0,0.0,cloves,
5 oz baz,baz,5.0,0.0,oz,
""".strip()))
partitioner.split_labels(mock_label_reader,
self.mock_training_writer,
self.mock_testing_writer,
training_fraction=0.67,
max_labels=3)
self.assertMultiLineEqual(
"""
input,name,qty,range_end,unit,comment
1 cup foo,foo,1.0,0.0,cup,
2 drops foz,foz,2.0,0.0,drop,
""".strip(),
self.mock_training_file.getvalue().strip())
self.assertMultiLineEqual(
"""
input,name,qty,range_end,unit,comment
3 ml faa,faa,3.0,0.0,ml,
""".strip(),
self.mock_testing_file.getvalue().strip())
def main(args):
with open(args.label_path) as label_file, open(
args.training_path, 'wb') as training_file, open(
args.testing_path, 'wb') as testing_file:
label_reader = labelled_data.Reader(label_file)
training_writer = labelled_data.Writer(training_file)
testing_writer = labelled_data.Writer(testing_file)
partitioner.split_labels(label_reader, training_writer, testing_writer,
args.training_fraction, args.max_labels)
def test_partition_20_percent_training(self):
mock_label_reader = labelled_data.Reader(
io.StringIO(
"""\
input,name,qty,range_end,unit,comment
1 cup foo,foo,1.0,0.0,cup,
2 drops foz,foz,2.0,0.0,drop,
3 ml faa,faa,3.0,0.0,ml,
4 cloves bar,bar,4.0,0.0,cloves,
5 oz baz,baz,5.0,0.0,oz,\
"""
)
)
partitioner.split_labels(
mock_label_reader,
self.mock_training_writer,
self.mock_testing_writer,
training_fraction=0.2,
)
self.assertMultiLineEqual(
"""\
input,name,qty,range_end,unit,comment
1 cup foo,foo,1.0,0.0,cup,\
""",
self.mock_training_file.getvalue().strip(),
)
self.assertMultiLineEqual(
"""\
input,name,qty,range_end,unit,comment
2 drops foz,foz,2.0,0.0,drop,
3 ml faa,faa,3.0,0.0,ml,
4 cloves bar,bar,4.0,0.0,cloves,
5 oz baz,baz,5.0,0.0,oz,\
""",
self.mock_testing_file.getvalue().strip(),
)