()
| 108 | |
| 109 | |
| 110 | def main(): |
| 111 | # Load the shared object file containing the Hardmax plugin implementation. |
| 112 | # By doing this, you will also register the Hardmax plugin with the TensorRT |
| 113 | # PluginRegistry through use of the macro REGISTER_TENSORRT_PLUGIN present |
| 114 | # in the plugin implementation. Refer to plugin/customHardmaxPlugin.cpp for more details. |
| 115 | load_plugin_lib() |
| 116 | |
| 117 | # Load pretrained model |
| 118 | if not os.path.isfile(TRT_MODEL_PATH): |
| 119 | raise IOError( |
| 120 | "\n{}\n{}\n{}\n".format( |
| 121 | "Failed to load model file ({}).".format(TRT_MODEL_PATH), |
| 122 | "Please use 'python3 model.py' to generate the ONNX model.", |
| 123 | "For more information, see README.md", |
| 124 | ) |
| 125 | ) |
| 126 | |
| 127 | if os.path.exists(ENGINE_FILE_PATH): |
| 128 | print(f"Loading saved TRT engine from {ENGINE_FILE_PATH}") |
| 129 | with open(ENGINE_FILE_PATH, "rb") as f: |
| 130 | runtime = trt.Runtime(TRT_LOGGER) |
| 131 | runtime.max_threads = 10 |
| 132 | engine = runtime.deserialize_cuda_engine(f.read()) |
| 133 | else: |
| 134 | print("Engine plan not saved. Building new engine...") |
| 135 | engine = build_engine(TRT_MODEL_PATH) |
| 136 | |
| 137 | inputs, outputs, bindings, stream = common.allocate_buffers(engine, profile_idx=0) |
| 138 | |
| 139 | testcases = [ |
| 140 | ('Garry the lion is 5 years old. He lives in the savanna.', 'Where does the lion live?'), |
| 141 | ('A quick brown fox jumps over the lazy dog.', 'What color is the fox?') |
| 142 | ] |
| 143 | |
| 144 | print("\n=== Testing ===") |
| 145 | |
| 146 | interactive = '--interactive' in sys.argv |
| 147 | if interactive: |
| 148 | context_text = input("Enter context: ") |
| 149 | query_text = input("Enter query: ") |
| 150 | testcases = [(context_text, query_text)] |
| 151 | |
| 152 | trt_context = engine.create_execution_context() |
| 153 | for (context_text, query_text) in testcases: |
| 154 | |
| 155 | context_words, _ = preprocess(context_text) |
| 156 | |
| 157 | load_test_case(inputs, context_text, query_text, trt_context) |
| 158 | if not interactive: |
| 159 | print(f"Input context: {context_text}") |
| 160 | print(f"Input query: {query_text}") |
| 161 | trt_outputs = common.do_inference_v2( |
| 162 | trt_context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream |
| 163 | ) |
| 164 | start = trt_outputs[1].item() |
| 165 | end = trt_outputs[0].item() |
| 166 | answer = context_words[start : end + 1].flatten() |
| 167 | print(f"Model prediction: ", " ".join(answer)) |
no test coverage detected