clean

Author: nikg4

src/oumi/datasets/vision_language/coco_captions.py

@@ -16,13 +16,13 @@ def transform_conversation(self, example: dict) -> Conversation:
             if required_key not in example:
                 raise ValueError(
                     f"Example doesn't contain '{required_key}' key. "
-                    f"Available keys: {example.keys()}"
+                    f"Available keys: {example.keys()}."
                 )
 
         if "raw" not in example["sentences"]:
             raise ValueError(
-                f"Example doesn't contain 'sentences.raw' key. "
-                f"Available keys under 'sentences.': {example['sentences'].keys()}"
+                f"Training example doesn't contain 'sentences.raw' key. "
+                f"Available keys under 'sentences.': {example['sentences'].keys()}."
             )
         output_text = example["sentences"]["raw"]
 
@@ -46,8 +46,9 @@ def transform_conversation(self, example: dict) -> Conversation:
             )
         else:
             raise ValueError(
-                f"Example contains none of required keys: 'image.bytes', 'image.path'. "
-                f"Available keys under 'image.': {example['image'].keys()}"
+                "Training example contains none of required keys: "
+                "'image.bytes', 'image.path'. "
+                f"Available keys under 'image.': {example['image'].keys()}."
             )
 
         messages.append(Message(role=Role.ASSISTANT, content=output_text))

tests/core/datasets/test_vision_language_dataset.py

@@ -1,3 +1,6 @@
+import functools
+import io
+from typing import Optional, Tuple
 from unittest.mock import Mock, patch
 
 import pytest
@@ -8,6 +11,17 @@
 from oumi.core.types.turn import Conversation, Message, Role, Type
 
 
+class EqBytesIO:
+    def __init__(self, bytes_io: io.BytesIO):
+        self._byte_io = bytes_io
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, io.BytesIO)
+            and other.getvalue() == self._byte_io.getvalue()
+        )
+
+
 @pytest.fixture
 def mock_processor():
     processor = Mock()
@@ -22,8 +36,18 @@ def mock_processor():
     return processor
 
 
+@functools.lru_cache(maxsize=None)  # same as @cache added in Python 3.9
+def _create_test_image(image_size: Optional[Tuple[int, int]] = None) -> bytes:
+    if image_size is None:
+        image_size = (80, 40)
+    image = Image.new(mode="RGBA", size=image_size)
+    bytes_io = io.BytesIO()
+    image.save(bytes_io, format="PNG")
+    return bytes_io.getvalue()
+
+
 @pytest.fixture
-def sample_conversation():
+def sample_conversation_using_image_path():
     return Conversation(
         messages=[
             Message(role=Role.USER, content="Describe this image:", type=Type.TEXT),
@@ -38,36 +62,101 @@ def sample_conversation():
 
 
 @pytest.fixture
-def test_dataset(mock_processor, sample_conversation):
-    class TestDataset(VisionLanguageSftDataset):
+def sample_conversation_using_image_binary():
+    return Conversation(
+        messages=[
+            Message(role=Role.USER, content="Describe this image:", type=Type.TEXT),
+            Message(
+                role=Role.USER, binary=_create_test_image(), type=Type.IMAGE_BINARY
+            ),
+            Message(
+                role=Role.ASSISTANT,
+                content="A beautiful sunset over the ocean.",
+                type=Type.TEXT,
+            ),
+        ]
+    )
+
+
+@pytest.fixture
+def test_dataset_using_image_path(
+    mock_processor: Mock, sample_conversation_using_image_path: Conversation
+):
+    class TestDatasetImagePath(VisionLanguageSftDataset):
+        default_dataset = "custom"
+
+        def transform_conversation(self, example):
+            return sample_conversation_using_image_path
+
+        def _load_data(self):
+            pass
+
+    return TestDatasetImagePath(processor=mock_processor)
+
+
+@pytest.fixture
+def test_dataset_using_image_binary(
+    mock_processor: Mock, sample_conversation_using_image_binary: Conversation
+):
+    class TestDatasetImageBinary(VisionLanguageSftDataset):
         default_dataset = "custom"
 
         def transform_conversation(self, example):
-            return sample_conversation
+            return sample_conversation_using_image_binary
 
         def _load_data(self):
             pass
 
-    return TestDataset(processor=mock_processor)
+    return TestDatasetImageBinary(processor=mock_processor)
 
 
-def test_transform_image(test_dataset):
+def test_transform_image_using_image_path(test_dataset_using_image_path):
     with patch("PIL.Image.open") as mock_open:
         mock_image = Mock(spec=Image.Image)
         mock_open.return_value.convert.return_value = mock_image
 
-        test_dataset.transform_image("path/to/image.jpg")
+        test_dataset_using_image_path.transform_image("path/to/image.jpg")
 
         mock_open.assert_called_once_with("path/to/image.jpg")
-        test_dataset._image_processor.assert_called_once()
+        test_dataset_using_image_path._image_processor.assert_called_once()
 
 
-def test_transform_simple_model(test_dataset):
-    with patch.object(test_dataset, "_load_image") as mock_load_image:
+def test_transform_image_using_image_binary(test_dataset_using_image_binary):
+    with patch("PIL.Image.open") as mock_open:
+        mock_image = Mock(spec=Image.Image)
+        mock_open.return_value.convert.return_value = mock_image
+
+        test_image_bytes = _create_test_image()
+        test_dataset_using_image_binary.transform_image(
+            Message(type=Type.IMAGE_BINARY, binary=test_image_bytes, role=Role.USER)
+        )
+
+        mock_open.assert_called_once_with(EqBytesIO(io.BytesIO(test_image_bytes)))
+        test_dataset_using_image_binary._image_processor.assert_called_once()
+
+
+def test_transform_simple_model_using_image_path(test_dataset_using_image_path):
+    with patch.object(test_dataset_using_image_path, "_load_image") as mock_load_image:
+        mock_image = Mock(spec=Image.Image)
+        mock_load_image.return_value = mock_image
+
+        result = test_dataset_using_image_path.transform({"example": "data"})
+
+    assert isinstance(result, dict)
+    assert "input_ids" in result
+    assert "attention_mask" in result
+    assert "labels" in result
+    assert "pixel_values" in result
+
+
+def test_transform_simple_model_using_image_binary(test_dataset_using_image_binary):
+    with patch.object(
+        test_dataset_using_image_binary, "_load_image"
+    ) as mock_load_image:
         mock_image = Mock(spec=Image.Image)
         mock_load_image.return_value = mock_image
 
-        result = test_dataset.transform({"example": "data"})
+        result = test_dataset_using_image_binary.transform({"example": "data"})
 
     assert isinstance(result, dict)
     assert "input_ids" in result
@@ -76,15 +165,39 @@ def test_transform_simple_model(test_dataset):
     assert "pixel_values" in result
 
 
-def test_transform_instruct_model(test_dataset, mock_processor):
+def test_transform_instruct_model_using_image_path(
+    test_dataset_using_image_path, mock_processor: Mock
+):
+    mock_processor.chat_template = "Template"
+    mock_processor.apply_chat_template = Mock(return_value="Processed template")
+
+    with patch.object(test_dataset_using_image_path, "_load_image") as mock_load_image:
+        mock_image = Mock(spec=Image.Image)
+        mock_load_image.return_value = mock_image
+
+        result = test_dataset_using_image_path.transform({"example": "data"})
+
+    assert isinstance(result, dict)
+    assert "input_ids" in result
+    assert "attention_mask" in result
+    assert "labels" in result
+    assert "pixel_values" in result
+    mock_processor.apply_chat_template.assert_called_once()
+
+
+def test_transform_instruct_model_using_image_binary(
+    test_dataset_using_image_binary, mock_processor: Mock
+):
     mock_processor.chat_template = "Template"
     mock_processor.apply_chat_template = Mock(return_value="Processed template")
 
-    with patch.object(test_dataset, "_load_image") as mock_load_image:
+    with patch.object(
+        test_dataset_using_image_binary, "_load_image"
+    ) as mock_load_image:
         mock_image = Mock(spec=Image.Image)
         mock_load_image.return_value = mock_image
 
-        result = test_dataset.transform({"example": "data"})
+        result = test_dataset_using_image_binary.transform({"example": "data"})
 
     assert isinstance(result, dict)
     assert "input_ids" in result