-
Notifications
You must be signed in to change notification settings - Fork 4.2k
/
tensorrt_text_classification.py
126 lines (105 loc) · 4.16 KB
/
tensorrt_text_classification.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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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
#
# http://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.
#
"""A pipeline to demonstrate usage of TensorRT with RunInference
for a text classification model. This pipeline reads data from a text
file, preprocesses the data, and then uses RunInference to generate
predictions from the text classification TensorRT engine. Next,
it postprocesses the RunInference outputs to print the input and
the predicted class label.
It also prints metrics provided by RunInference.
"""
import argparse
import logging
import numpy as np
import apache_beam as beam
from apache_beam.ml.inference.base import RunInference
from apache_beam.ml.inference.tensorrt_inference import TensorRTEngineHandlerNumPy
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions
from transformers import AutoTokenizer
class Preprocess(beam.DoFn):
"""Processes the input sentences to tokenize them.
The input sentences are tokenized because the
model is expecting tokens.
"""
def __init__(self, tokenizer: AutoTokenizer):
self._tokenizer = tokenizer
def process(self, element):
inputs = self._tokenizer(
element, return_tensors="np", padding="max_length", max_length=128)
return inputs.input_ids
class Postprocess(beam.DoFn):
"""Processes the PredictionResult to get the predicted class.
The logits are the output of the TensorRT engine.
We can get the class label by getting the index of
maximum logit using argmax.
"""
def __init__(self, tokenizer: AutoTokenizer):
self._tokenizer = tokenizer
def process(self, element):
decoded_input = self._tokenizer.decode(
element.example, skip_special_tokens=True)
logits = element.inference[0]
argmax = np.argmax(logits)
output = "Positive" if argmax == 1 else "Negative"
yield decoded_input, output
def parse_known_args(argv):
"""Parses args for the workflow."""
parser = argparse.ArgumentParser()
parser.add_argument(
'--input',
dest='input',
required=True,
help='Path to the text file containing sentences.')
parser.add_argument(
'--trt_model_path',
dest='trt_model_path',
required=True,
help='Path to the pre-built textattack/bert-base-uncased-SST-2'
'TensorRT engine.')
parser.add_argument(
'--model_id',
dest='model_id',
default="textattack/bert-base-uncased-SST-2",
help="name of model.")
return parser.parse_known_args(argv)
def run(
argv=None,
save_main_session=True,
):
known_args, pipeline_args = parse_known_args(argv)
pipeline_options = PipelineOptions(pipeline_args)
pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
model_handler = TensorRTEngineHandlerNumPy(
min_batch_size=1,
max_batch_size=1,
engine_path=known_args.trt_model_path,
)
tokenizer = AutoTokenizer.from_pretrained(known_args.model_id)
with beam.Pipeline(options=pipeline_options) as pipeline:
_ = (
pipeline
| "ReadSentences" >> beam.io.ReadFromText(known_args.input)
| "Preprocess" >> beam.ParDo(Preprocess(tokenizer=tokenizer))
| "RunInference" >> RunInference(model_handler=model_handler)
| "PostProcess" >> beam.ParDo(Postprocess(tokenizer=tokenizer))
| "LogResult" >> beam.Map(logging.info))
metrics = pipeline.result.metrics().query(beam.metrics.MetricsFilter())
logging.info(metrics)
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
run()