| 68 | return bleu |
| 69 | |
| 70 | def getT5EncoderPlugin(arg): |
| 71 | nBatchSize = arg['batch_size'] |
| 72 | nMaxSeqLen = arg['max_seq_len'] |
| 73 | nBeamSize = arg['beam_width'], |
| 74 | nSM = globalNSM |
| 75 | useFP16 = int(arg['data_type'] == 'fp16') |
| 76 | ckpt_path = arg['ckpt_path'].encode() |
| 77 | for c in trt.get_plugin_registry().plugin_creator_list: |
| 78 | if c.name == 'T5EncoderPlugin': |
| 79 | pList = [ |
| 80 | trt.PluginField('max_batch_size', np.int32(nBatchSize), npToPFT[np.int32]), |
| 81 | trt.PluginField('max_seq_len', np.int32(nMaxSeqLen), npToPFT[np.int32]), |
| 82 | trt.PluginField('beam_width', np.int32(nBeamSize), npToPFT[np.int32]), |
| 83 | trt.PluginField('sm', np.int32(nSM), npToPFT[np.int32]), |
| 84 | trt.PluginField('useFP16', np.int32(useFP16), npToPFT[np.int32]), |
| 85 | trt.PluginField('ckpt_path', ckpt_path, trt.PluginFieldType.CHAR), |
| 86 | ] |
| 87 | return c.create_plugin(c.name, trt.PluginFieldCollection(pList)) |
| 88 | return None |
| 89 | |
| 90 | def getT5DecodingPlugin(arg): |
| 91 | nBatchSize = arg['batch_size'] |