chore: import upstream snapshot with attribution

test/python/testpipeline/testimage/testcaption.py

@@ -0,0 +1,36 @@
+"""
+Caption module tests
+"""
+
+import unittest
+
+from PIL import Image
+
+from transformers import AutoModelForImageTextToText, AutoImageProcessor, AutoTokenizer
+from txtai.pipeline import Caption
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestCaption(unittest.TestCase):
+    """
+    Caption tests.
+    """
+
+    def testCaption(self):
+        """
+        Test captions
+        """
+
+        caption = Caption()
+        self.assertEqual(caption(Image.open(Utils.PATH + "/books.jpg")), "a book shelf filled with books and a stack of books")
+
+        # Load passing models directly
+        path = "ydshieh/vit-gpt2-coco-en"
+        model = AutoModelForImageTextToText.from_pretrained(path)
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        processor = AutoImageProcessor.from_pretrained(path)
+
+        caption = Caption((model, tokenizer, processor))
+        self.assertEqual(caption(Image.open(Utils.PATH + "/books.jpg")), "a book shelf filled with books and a stack of books")