-
Notifications
You must be signed in to change notification settings - Fork 273
/
writer.py
58 lines (50 loc) · 2.26 KB
/
writer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Writer types."""
from typing import NamedTuple, Optional, Union
import apache_beam as beam
from tensorflow_model_analysis.evaluators import evaluator
from tensorflow_model_analysis.validators import validator
# A writer is a PTransform that takes Evaluation or Validation output as input
# and serializes the associated PCollections of data to a sink.
Writer = NamedTuple(
'Writer',
[
('stage_name', str),
# PTransform (Evaluation|Validation) -> Beam write result
('ptransform', beam.PTransform)
])
@beam.ptransform_fn
def Write(evaluation_or_validation: Union[evaluator.Evaluation,
validator.Validation], key: str,
ptransform: beam.PTransform) -> Optional[beam.PCollection]:
"""Writes given Evaluation or Validation data using given writer PTransform.
Args:
evaluation_or_validation: Evaluation or Validation data.
key: Key for Evaluation or Validation output to write. It is valid for the
key to not exist in the dict (in which case the write is a no-op).
ptransform: PTransform to use for writing.
Raises:
ValueError: If Evaluation or Validation is empty. The key does not need to
exist in the Evaluation or Validation, but the dict must not be empty.
Returns:
The result of the underlying beam write PTransform. This makes it possible
for interactive environments to execute your writer, as well as for
downstream Beam stages to make use of the files that are written.
"""
if not evaluation_or_validation:
raise ValueError('Evaluations and Validations cannot be empty')
if key in evaluation_or_validation:
return evaluation_or_validation[key] | ptransform
return None