HFOnnx

pipeline pipeline

Exports a Hugging Face Transformer model to ONNX. Currently, this works best with classification/pooling/qa models. Work is ongoing for sequence to sequence models (summarization, transcription, translation).

Example

The following shows a simple example using this pipeline.

  1. from txtai.pipeline import HFOnnx, Labels
  3. # Model path
  4. path = "distilbert-base-uncased-finetuned-sst-2-english"
  6. # Export model to ONNX
  7. onnx = HFOnnx()
  8. model = onnx(path, "text-classification", "model.onnx", True)
  10. # Run inference and validate
  11. labels = Labels((model, path), dynamic=False)
  12. labels("I am happy")

See the link below for a more detailed example.

NotebookDescription
Export and run models with ONNXExport models with ONNX, run natively in JavaScript, Java and RustOpen In Colab

Methods

Python documentation for the pipeline.

Exports a Hugging Face Transformer model to ONNX.

Parameters:

NameTypeDescriptionDefault
path

path to model, accepts Hugging Face model hub id, local path or (model, tokenizer) tuple

required
task

optional model task or category, determines the model type and outputs, defaults to export hidden state

‘default’
output

optional output model path, defaults to return byte array if None

None
quantize

if model should be quantized (requires onnx to be installed), defaults to False

False
opset

onnx opset, defaults to 12

12

Returns:

TypeDescription

path to model output or model as bytes depending on output parameter

Source code in txtai/pipeline/train/hfonnx.py

  1. 32
  2. 33
  3. 34
  4. 35
  5. 36
  6. 37
  7. 38
  8. 39
  9. 40
  10. 41
  11. 42
  12. 43
  13. 44
  14. 45
  15. 46
  16. 47
  17. 48
  18. 49
  19. 50
  20. 51
  21. 52
  22. 53
  23. 54
  24. 55
  25. 56
  26. 57
  27. 58
  28. 59
  29. 60
  30. 61
  31. 62
  32. 63
  33. 64
  34. 65
  35. 66
  36. 67
  37. 68
  38. 69
  39. 70
  40. 71
  41. 72
  42. 73
  43. 74
  44. 75
  45. 76
  46. 77
  47. 78
  48. 79
  49. 80
  50. 81
  51. 82
  52. 83
  53. 84
  54. 85
  55. 86
  1. def call(self, path, task=”default”, output=None, quantize=False, opset=12):
  2.     “””
  3.     Exports a Hugging Face Transformer model to ONNX.
  5.     Args:
  6.         path: path to model, accepts Hugging Face model hub id, local path or (model, tokenizer) tuple
  7.         task: optional model task or category, determines the model type and outputs, defaults to export hidden state
  8.         output: optional output model path, defaults to return byte array if None
  9.         quantize: if model should be quantized (requires onnx to be installed), defaults to False
  10.         opset: onnx opset, defaults to 12
  12.     Returns:
  13.         path to model output or model as bytes depending on output parameter
  14.     “””
  16.     inputs, outputs, model = self.parameters(task)
  18.     if isinstance(path, (list, tuple)):
  19.         model, tokenizer = path
  20.         model = model.cpu()
  21.     else:
  22.         model = model(path)
  23.         tokenizer = AutoTokenizer.from_pretrained(path)
  25.     # Generate dummy inputs
  26.     dummy = dict(tokenizer([“test inputs”], return_tensors=”pt”))
  28.     # Default to BytesIO if no output file provided
  29.     output = output if output else BytesIO()
  31.     # Export model to ONNX
  32.     export(
  33.         model,
  34.         (dummy,),
  35.         output,
  36.         opset_version=opset,
  37.         do_constant_folding=True,
  38.         input_names=list(inputs.keys()),
  39.         output_names=list(outputs.keys()),
  40.         dynamic_axes=dict(chain(inputs.items(), outputs.items())),
  41.     )
  43.     # Quantize model
  44.     if quantize:
  45.         if not ONNX_RUNTIME:
  46.             raise ImportError(‘onnxruntime is not available - install “pipeline” extra to enable’)
  48.         output = self.quantization(output)
  50.     if isinstance(output, BytesIO):
  51.         # Reset stream and return bytes
  52.         output.seek(0)
  53.         output = output.read()
  55.     return output