1  """
  2  ONNX module tests
  3  """
  4  
  5  import os
  6  import tempfile
  7  import unittest
  8  
  9  from unittest.mock import patch
 10  
 11  from sklearn.feature_extraction.text import TfidfVectorizer
 12  from sklearn.linear_model import LogisticRegression
 13  from sklearn.pipeline import Pipeline
 14  
 15  from txtai.embeddings import Embeddings
 16  from txtai.models import OnnxModel
 17  from txtai.pipeline import HFOnnx, HFTrainer, Labels, MLOnnx, Questions
 18  
 19  
 20  class TestOnnx(unittest.TestCase):
 21      """
 22      ONNX tests.
 23      """
 24  
 25      @classmethod
 26      def setUpClass(cls):
 27          """
 28          Create default datasets.
 29          """
 30  
 31          cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100
 32  
 33      def testDefault(self):
 34          """
 35          Test exporting an ONNX model with default parameters
 36          """
 37  
 38          # Export model to ONNX, use default parameters
 39          onnx = HFOnnx()
 40          model = onnx("google/bert_uncased_L-2_H-128_A-2")
 41  
 42          # Validate model has data
 43          self.assertGreater(len(model), 0)
 44  
 45          # Validate model device properly works
 46          self.assertEqual(OnnxModel(model).device, -1)
 47  
 48      def testClassification(self):
 49          """
 50          Test exporting a classification model to ONNX and running inference
 51          """
 52  
 53          path = "google/bert_uncased_L-2_H-128_A-2"
 54  
 55          trainer = HFTrainer()
 56          model, tokenizer = trainer(path, self.data)
 57  
 58          # Output file path
 59          output = os.path.join(tempfile.gettempdir(), "onnx")
 60  
 61          # Export model to ONNX
 62          onnx = HFOnnx()
 63          model = onnx((model, tokenizer), "text-classification", output, True)
 64  
 65          # Test classification
 66          labels = Labels((model, path), dynamic=False)
 67          self.assertEqual(labels("cat")[0][0], 1)
 68  
 69      @patch("onnxruntime.get_available_providers")
 70      @patch("torch.cuda.is_available")
 71      def testPooling(self, cuda, providers):
 72          """
 73          Test exporting a pooling model to ONNX and running inference
 74          """
 75  
 76          path = "sentence-transformers/paraphrase-MiniLM-L3-v2"
 77  
 78          # Export model to ONNX
 79          onnx = HFOnnx()
 80          model = onnx(path, "pooling", quantize=True)
 81  
 82          # Test no CUDA and onnxruntime installed
 83          cuda.return_value = False
 84          providers.return_value = ["CPUExecutionProvider"]
 85  
 86          embeddings = Embeddings({"path": model, "tokenizer": path})
 87          self.assertEqual(embeddings.similarity("animal", ["dog", "book", "rug"])[0][0], 0)
 88  
 89          # Test no CUDA and onnxruntime-gpu installed
 90          cuda.return_value = False
 91          providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
 92  
 93          embeddings = Embeddings({"path": model, "tokenizer": path})
 94          self.assertIsNotNone(embeddings)
 95  
 96          # Test CUDA and only onnxruntime installed
 97          cuda.return_value = True
 98          providers.return_value = ["CPUExecutionProvider"]
 99  
100          embeddings = Embeddings({"path": model, "tokenizer": path})
101          self.assertIsNotNone(embeddings)
102  
103          # Test CUDA and onnxruntime-gpu installed
104          cuda.return_value = True
105          providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
106  
107          embeddings = Embeddings({"path": model, "tokenizer": path})
108          self.assertIsNotNone(embeddings)
109  
110      def testQA(self):
111          """
112          Test exporting a QA model to ONNX and running inference
113          """
114  
115          path = "distilbert-base-cased-distilled-squad"
116  
117          # Export model to ONNX
118          onnx = HFOnnx()
119          model = onnx(path, "question-answering")
120  
121          questions = Questions((model, path))
122          self.assertEqual(questions(["What is the price?"], ["The price is $30"])[0], "$30")
123  
124      def testScikit(self):
125          """
126          Test exporting a scikit-learn model to ONNX and running inference
127          """
128  
129          # pylint: disable=W0613
130          def tokenizer(inputs, **kwargs):
131              if isinstance(inputs, str):
132                  inputs = [inputs]
133  
134              return {"input_ids": [[x] for x in inputs]}
135  
136          # Train a scikit-learn model
137          model = Pipeline([("tfidf", TfidfVectorizer()), ("lr", LogisticRegression())])
138          model.fit([x["text"] for x in self.data], [x["label"] for x in self.data])
139  
140          # Export model to ONNX
141          onnx = MLOnnx()
142          model = onnx(model)
143  
144          # Test classification
145          labels = Labels((model, tokenizer), dynamic=False)
146          self.assertEqual(labels("cat")[0][0], 1)
147  
148      def testZeroShot(self):
149          """
150          Test exporting a zero shot classification model to ONNX and running inference
151          """
152  
153          path = "prajjwal1/bert-medium-mnli"
154  
155          # Export model to ONNX
156          onnx = HFOnnx()
157          model = onnx(path, "zero-shot-classification", quantize=True)
158  
159          # Test zero shot classification
160          labels = Labels((model, path))
161          self.assertEqual(labels("That is great news", ["negative", "positive"])[0][0], 1)