chore: import upstream snapshot with attribution

test/python/testpipeline/testtrain/testonnx.py

@@ -0,0 +1,161 @@
+"""
+ONNX module tests
+"""
+
+import os
+import tempfile
+import unittest
+
+from unittest.mock import patch
+
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.pipeline import Pipeline
+
+from txtai.embeddings import Embeddings
+from txtai.models import OnnxModel
+from txtai.pipeline import HFOnnx, HFTrainer, Labels, MLOnnx, Questions
+
+
+class TestOnnx(unittest.TestCase):
+    """
+    ONNX tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create default datasets.
+        """
+
+        cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100
+
+    def testDefault(self):
+        """
+        Test exporting an ONNX model with default parameters
+        """
+
+        # Export model to ONNX, use default parameters
+        onnx = HFOnnx()
+        model = onnx("google/bert_uncased_L-2_H-128_A-2")
+
+        # Validate model has data
+        self.assertGreater(len(model), 0)
+
+        # Validate model device properly works
+        self.assertEqual(OnnxModel(model).device, -1)
+
+    def testClassification(self):
+        """
+        Test exporting a classification model to ONNX and running inference
+        """
+
+        path = "google/bert_uncased_L-2_H-128_A-2"
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer(path, self.data)
+
+        # Output file path
+        output = os.path.join(tempfile.gettempdir(), "onnx")
+
+        # Export model to ONNX
+        onnx = HFOnnx()
+        model = onnx((model, tokenizer), "text-classification", output, True)
+
+        # Test classification
+        labels = Labels((model, path), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    @patch("onnxruntime.get_available_providers")
+    @patch("torch.cuda.is_available")
+    def testPooling(self, cuda, providers):
+        """
+        Test exporting a pooling model to ONNX and running inference
+        """
+
+        path = "sentence-transformers/paraphrase-MiniLM-L3-v2"
+
+        # Export model to ONNX
+        onnx = HFOnnx()
+        model = onnx(path, "pooling", quantize=True)
+
+        # Test no CUDA and onnxruntime installed
+        cuda.return_value = False
+        providers.return_value = ["CPUExecutionProvider"]
+
+        embeddings = Embeddings({"path": model, "tokenizer": path})
+        self.assertEqual(embeddings.similarity("animal", ["dog", "book", "rug"])[0][0], 0)
+
+        # Test no CUDA and onnxruntime-gpu installed
+        cuda.return_value = False
+        providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+        embeddings = Embeddings({"path": model, "tokenizer": path})
+        self.assertIsNotNone(embeddings)
+
+        # Test CUDA and only onnxruntime installed
+        cuda.return_value = True
+        providers.return_value = ["CPUExecutionProvider"]
+
+        embeddings = Embeddings({"path": model, "tokenizer": path})
+        self.assertIsNotNone(embeddings)
+
+        # Test CUDA and onnxruntime-gpu installed
+        cuda.return_value = True
+        providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+        embeddings = Embeddings({"path": model, "tokenizer": path})
+        self.assertIsNotNone(embeddings)
+
+    def testQA(self):
+        """
+        Test exporting a QA model to ONNX and running inference
+        """
+
+        path = "distilbert-base-cased-distilled-squad"
+
+        # Export model to ONNX
+        onnx = HFOnnx()
+        model = onnx(path, "question-answering")
+
+        questions = Questions((model, path))
+        self.assertEqual(questions(["What is the price?"], ["The price is $30"])[0], "$30")
+
+    def testScikit(self):
+        """
+        Test exporting a scikit-learn model to ONNX and running inference
+        """
+
+        # pylint: disable=W0613
+        def tokenizer(inputs, **kwargs):
+            if isinstance(inputs, str):
+                inputs = [inputs]
+
+            return {"input_ids": [[x] for x in inputs]}
+
+        # Train a scikit-learn model
+        model = Pipeline([("tfidf", TfidfVectorizer()), ("lr", LogisticRegression())])
+        model.fit([x["text"] for x in self.data], [x["label"] for x in self.data])
+
+        # Export model to ONNX
+        onnx = MLOnnx()
+        model = onnx(model)
+
+        # Test classification
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testZeroShot(self):
+        """
+        Test exporting a zero shot classification model to ONNX and running inference
+        """
+
+        path = "prajjwal1/bert-medium-mnli"
+
+        # Export model to ONNX
+        onnx = HFOnnx()
+        model = onnx(path, "zero-shot-classification", quantize=True)
+
+        # Test zero shot classification
+        labels = Labels((model, path))
+        self.assertEqual(labels("That is great news", ["negative", "positive"])[0][0], 1)

test/python/testpipeline/testtrain/testquantization.py

@@ -0,0 +1,35 @@
+"""
+Quantization module tests
+"""
+
+import platform
+import unittest
+
+from transformers import AutoModel
+
+from txtai.pipeline import HFModel, HFPipeline
+
+
+class TestQuantization(unittest.TestCase):
+    """
+    Quantization tests.
+    """
+
+    @unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
+    def testModel(self):
+        """
+        Test quantizing a model through HFModel.
+        """
+
+        model = HFModel(quantize=True, gpu=False)
+        model = model.prepare(AutoModel.from_pretrained("google/bert_uncased_L-2_H-128_A-2"))
+        self.assertIsNotNone(model)
+
+    @unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
+    def testPipeline(self):
+        """
+        Test quantizing a model through HFPipeline.
+        """
+
+        pipeline = HFPipeline("text-classification", "google/bert_uncased_L-2_H-128_A-2", True, False)
+        self.assertIsNotNone(pipeline)

test/python/testpipeline/testtrain/testtrainer.py

@@ -0,0 +1,335 @@
+"""
+Trainer module tests
+"""
+
+import os
+import unittest
+import tempfile
+
+import torch
+
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+from txtai.data import Data
+from txtai.pipeline import HFTrainer, Labels, Questions, Sequences
+
+
+class TestTrainer(unittest.TestCase):
+    """
+    Trainer tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create default datasets.
+        """
+
+        cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100
+
+    def testBasic(self):
+        """
+        Test training a model with basic parameters
+        """
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testCLM(self):
+        """
+        Test training a model with causal language modeling
+        """
+
+        trainer = HFTrainer()
+
+        # Test default parameters
+        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation")
+        self.assertIsNotNone(model)
+
+        # Test pack merging
+        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge="pack")
+        self.assertIsNotNone(model)
+
+        # Test no merging
+        model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge=None)
+        self.assertIsNotNone(model)
+
+    def testCustom(self):
+        """
+        Test training a model with custom parameters
+        """
+
+        # pylint: disable=E1120
+        model = AutoModelForSequenceClassification.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
+        tokenizer = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer(
+            (model, tokenizer),
+            self.data,
+            self.data,
+            columns=("text", "label"),
+            do_eval=True,
+            output_dir=os.path.join(tempfile.gettempdir(), "trainer"),
+        )
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testDataFrame(self):
+        """
+        Test training a model with a mock pandas DataFrame
+        """
+
+        class TestDataFrame:
+            """
+            Test DataFrame
+            """
+
+            def __init__(self, data):
+                # Get list of columns
+                self.columns = list(data[0].keys())
+
+                # Build columnar data view
+                self.data = {}
+                for column in self.columns:
+                    self.data[column] = Values([row[column] for row in data])
+
+            def __getitem__(self, column):
+                return self.data[column]
+
+        class Values:
+            """
+            Test values list
+            """
+
+            def __init__(self, values):
+                self.values = list(values)
+
+            def __getitem__(self, index):
+                return self.values[index]
+
+            def unique(self):
+                """
+                Returns a list of unique values.
+
+                Returns:
+                    unique list of values
+                """
+
+                return set(self.values)
+
+        # Mock DataFrame
+        df = TestDataFrame(self.data)
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", df)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testDataset(self):
+        """
+        Test training a model with a mock Hugging Face Dataset
+        """
+
+        class TestDataset(torch.utils.data.Dataset):
+            """
+            Test Dataset
+            """
+
+            def __init__(self, data):
+                self.data = data
+                self.unique = lambda _: [0, 1]
+
+            def __len__(self):
+                return len(self.data)
+
+            def __getitem__(self, index):
+                return self.data[index]
+
+            def column_names(self):
+                """
+                Returns column names for this dataset
+
+                Returns:
+                    list of columns
+                """
+
+                return ["text", "label"]
+
+            # pylint: disable=W0613
+            def map(self, fn, batched, num_proc, remove_columns):
+                """
+                Map each dataset row using fn.
+
+                Args:
+                    fn: function
+                    batched: batch records
+
+                Returns:
+                    updated Dataset
+                """
+
+                self.data = [fn(x) for x in self.data]
+                return self
+
+        ds = TestDataset(self.data)
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", ds)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testEmpty(self):
+        """
+        Test an empty training data object
+        """
+
+        self.assertIsNone(Data(None, None, None).process(None))
+
+    def testKD(self):
+        """
+        Test knowledge distillation
+        """
+
+        # Base model
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)
+
+        # Train with knowledge distillation
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data, teacher=(model, tokenizer))
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testMLM(self):
+        """
+        Test training a model with masked language modeling.
+        """
+
+        trainer = HFTrainer()
+        model, _ = trainer("hf-internal-testing/tiny-random-bert", self.data, task="language-modeling")
+
+        # Test model completed successfully
+        self.assertIsNotNone(model)
+
+    def testMultiLabel(self):
+        """
+        Test training model with labels provided as a list
+        """
+
+        data = []
+        for x in self.data:
+            data.append({"text": x["text"], "label": [0.0, 1.0] if x["label"] else [1.0, 0.0]})
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+        self.assertEqual(labels("cat")[0][0], 1)
+
+    def testPEFT(self):
+        """
+        Test training a model with causal language modeling and PEFT
+        """
+
+        trainer = HFTrainer()
+        model, _ = trainer(
+            "hf-internal-testing/tiny-random-gpt2",
+            self.data,
+            maxlength=16,
+            task="language-generation",
+            quantize=True,
+            lora=True,
+        )
+
+        # Test model completed successfully
+        self.assertIsNotNone(model)
+
+    def testQA(self):
+        """
+        Test training a QA model
+        """
+
+        # Training data
+        data = [
+            {"question": "What ingredient?", "context": "1 can whole tomatoes", "answers": "tomatoes"},
+            {"question": "What ingredient?", "context": "Crush 1 tomato", "answers": "tomato"},
+            {"question": "What ingredient?", "context": "1 yellow onion", "answers": "onion"},
+            {"question": "What ingredient?", "context": "Unwrap 2 red onions", "answers": "onions"},
+            {"question": "What ingredient?", "context": "1 red pepper", "answers": "pepper"},
+            {"question": "What ingredient?", "context": "Clean 3 red peppers", "answers": "peppers"},
+            {"question": "What ingredient?", "context": "1 clove garlic", "answers": "garlic"},
+            {"question": "What ingredient?", "context": "Unwrap 3 cloves of garlic", "answers": "garlic"},
+            {"question": "What ingredient?", "context": "3 pieces of ginger", "answers": "ginger"},
+            {"question": "What ingredient?", "context": "Peel 1 orange", "answers": "orange"},
+            {"question": "What ingredient?", "context": "1/2 lb beef", "answers": "beef"},
+            {"question": "What ingredient?", "context": "Roast 3 lbs of beef", "answers": "beef"},
+            {"question": "What ingredient?", "context": "1 pack of chicken", "answers": "chicken"},
+            {"question": "What ingredient?", "context": "Forest through the trees", "answers": None},
+        ]
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data, data, task="question-answering", num_train_epochs=40)
+
+        questions = Questions((model, tokenizer), gpu=True)
+        self.assertTrue("onion" in questions(["What ingredient?"], ["Peel 1 onion"])[0])
+
+    def testRegression(self):
+        """
+        Test training a model with a regression (continuous) output
+        """
+
+        data = []
+        for x in self.data:
+            data.append({"text": x["text"], "label": x["label"] + 0.1})
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)
+
+        labels = Labels((model, tokenizer), dynamic=False)
+
+        # Regression tasks return a single entry with the regression output
+        self.assertGreater(labels("cat")[0][1], 0.5)
+
+    def testRTD(self):
+        """
+        Test training a language model with replaced token detection
+        """
+
+        # Save directory
+        output = os.path.join(tempfile.gettempdir(), "trainer.rtd")
+
+        trainer = HFTrainer()
+        model, _ = trainer("hf-internal-testing/tiny-random-electra", self.data, task="token-detection", output_dir=output)
+
+        # Test model completed successfully
+        self.assertIsNotNone(model)
+
+        # Test output directories exist
+        self.assertTrue(os.path.exists(os.path.join(output, "generator")))
+        self.assertTrue(os.path.exists(os.path.join(output, "discriminator")))
+
+    def testSeqSeq(self):
+        """
+        Test training a sequence-sequence model
+        """
+
+        data = [
+            {"source": "Running again", "target": "Sleeping again"},
+            {"source": "Run", "target": "Sleep"},
+            {"source": "running", "target": "sleeping"},
+        ]
+
+        trainer = HFTrainer()
+        model, tokenizer = trainer("t5-small", data, task="sequence-sequence", prefix="translate Run to Sleep: ", learning_rate=1e-3)
+
+        # Run run-sleep translation
+        sequences = Sequences((model, tokenizer))
+        result = sequences("translate Run to Sleep: run")
+        self.assertEqual(result.lower(), "sleep")