damccorm commented on code in PR #28203:
URL: https://github.com/apache/beam/pull/28203#discussion_r1309335291
##########
sdks/python/apache_beam/ml/transforms/tft_test.py:
##########
@@ -730,6 +730,41 @@ def test_bag_of_words_on_by_splitting_input_text(self):
], [b'yum', b'yum yum', b'yum yum pie', b'yum pie', b'pie']]
assert_that(result, equal_to(expected_data, equals_fn=np.array_equal))
+ def test_count_per_key_on_list(self):
+ def map_element_to_count(elements, counts):
+ d = {elements[i]: counts[i] for i in range(len(elements))}
+ return d
+
+ data = [{
+ 'x': ['I', 'like', 'pie', 'pie', 'pie'],
+ }, {
+ 'x': ['yum', 'yum', 'pie']
+ }, {
+ 'x': ['Banana', 'Banana', 'Apple', 'Apple', 'Apple', 'Apple']
+ }]
+ with beam.Pipeline() as p:
+ result = (
+ p
+ | "Create" >> beam.Create(data)
+ | "MLTransform" >> base.MLTransform(
+ write_artifact_location=self.artifact_location,
+ transforms=[
+ tft.BagOfWords(columns=['x'], compute_word_count=True)
+ ]))
+
+ # the unique elements and counts are artifacts and will be
+ # stored in the result and same for all the elements in the
+ # PCollection.
+ result = result | beam.Map(
+ lambda x: map_element_to_count(x.x_unique_elements, x.x_counts))
+
+ expected_data = [{
+ b'Apple': 4, b'Banana': 2, b'I': 1, b'like': 1, b'pie': 4, b'yum': 2
+ }] * 3 # since there are 3 elements in input.
+ assert_that(result, equal_to(expected_data))
+
+ # result | beam.Map(print)
Review Comment:
```suggestion
assert_that(result, equal_to(expected_data))
```
##########
sdks/python/apache_beam/ml/transforms/tft.py:
##########
@@ -598,15 +605,34 @@ def __init__(
self.ngrams_separator = ngrams_separator
self.name = name
self.split_string_by_delimiter = split_string_by_delimiter
+ self.compute_word_count = compute_word_count
+ if self.compute_word_count:
Review Comment:
```suggestion
if compute_word_count:
```
I don't think we actually use this again
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
