Package: transformers-inference-toolkit
Transformers Inference Toolkit🤗Transformerslibrary provides great API for manipulating pre-trained NLP (as well as CV and Audio-related) models. However, preparing 🤗 Transformers models for use in production usually requires additional effort. The purpose oftransformers-inference-toolkitis to get rid of boilerplate code and to simplify automatic optimization and inference process of Huggingface Transformers models.InstallationUsingpip:pipinstalltransformers-inference-toolkitOptimizationThe original 🤗 Transformers library includestransformers.onnxpackage, which can be used to convert PyTorch or TensorFlow models intoONNXformat. This Toolkit extends this functionality by giving the user an opportunity to automaticallyoptimize ONNX model graph- this is similar to what 🤗Optimumlibrary does, but 🤗 Optimum currently has limited support for locally stored pre-trained models as well as for models of less popular architectures (for example,MPNet).Aside from ONNX conversion the Toolkit also supports resaving PyTorch models with half-precision and setting upDeepSpeed Inference.PrerequisiteThe Toolkit expects your pretrained model (in PyTorch format) and tokenizer to be saved (usingsave_pretrained()method) inside a common parent directory inmodelandtokenizerfolders respectively. This is how a file structure oftoxic-bertmodel should look like:toxic-bert
├──model
│├──config.json
│└──pytorch_model.bin
└──tokenizer├──merges.txt├──special_tokens_map.json├──tokenizer_config.json└──vocab.jsonHow to useMost of the popular Transformer model architectures (like BERT and its variations) can be converted with a single command:fromtransformers_inference_toolkitimport(Feature,OnnxModelType,OnnxOptimizationLevel,optimizer,)optimizer.pack_onnx(input_path="toxic-bert",output_path="toxic-bert-optimized",feature=Feature.SEQUENCE_CLASSIFICATION,for_gpu=True,fp16=True,optimization_level=OnnxOptimizationLevel.FULL,)If your model architecture is not supported out-of-the-box (describedhere) you can try writing a custom OnnxConfig class:fromcollectionsimportOrderedDictfromtransformers.onnximportOnnxConfigclassMPNetOnnxConfig(OnnxConfig):@propertydefdefault_onnx_opset(self):return14@propertydefinputs(self):dynamic_axis={0:"batch",1:"sequence"}returnOrderedDict([("input_ids",dynamic_axis),("attention_mask",dynamic_axis),])optimizer.pack_onnx(input_path="all-mpnet-base-v2",output_path="all-mpnet-base-v2-optimized",feature=Feature.DEFAULT,custom_onnx_config_cls=MPNetOnnxConfig,)ONNX is not the only option, it is also possible to resave the model for future inference simply using PyTorch (optimizer.pack_transformers()method,force_fp16argument to save in half-precision) orDeepSpeed Inference(optimizer.pack_deepspeed()method):optimizer.pack_deepspeed(input_path="gpt-neo",output_path="gpt-neo-optimized",feature=Feature.CAUSAL_LM,replace_with_kernel_inject=True,mp_size=1,)After callingoptimizermethods the model and tokenizer would be saved atoutput_path. The output directory will also containmetadata.jsonfile that is necessary for thePredictorobject (described below) to correctly load the model:toxic-bert-optimized
├──metadata.json
├──model
│├──config.json
│└──model.onnx
└──tokenizer├──special_tokens_map.json├──tokenizer.json└──tokenizer_config.jsonPredictionAfter model and tokenizer are packaged using one of theoptimizermethods, it is possible to initialize aPredictorobject:>>>fromtransformers_inference_toolkitimportPredictor>>>>>>predictor=Predictor("toxic-bert-optimized",cuda=True)>>>print(predictor("I hate this!")){'logits':array([[0.02940369,-7.0195312,-4.7890625,-6.0664062,-5.625,-6.09375]],dtype=float32)}ThePredictorobject can be simply called with tokenizer arguments (similar to 🤗 Transformerspipelines,return_tensorsargument can be omitted,paddingandtruncationareTrueby default). For text generation tasksPredictor.generate()method (withgeneration arguments) can be used:>>>predictor=Predictor("gpt-neo-optimized",cuda=True)>>>predictor.generate(..."Tommy: Hi Mark!",...do_sample=True,...top_p=0.9,...num_return_sequences=3,...max_new_tokens=5,...)['Tommy: Hi Mark!\nMadelyn: Hello','Tommy: Hi Mark! It’s so','Tommy: Hi Mark! How are you?\n']