From 077eae79f33c8d6367419077895e9cc5f8eda84f Mon Sep 17 00:00:00 2001
From: movchan74 <aleksandr@mobius.ml>
Date: Wed, 8 Nov 2023 10:53:17 +0000
Subject: [PATCH 1/4] Implemented suggestions from PR review

---
 aana/deployments/hf_blip2_deployment.py    |  2 +-
 aana/models/core/dtype.py                  | 19 ++---
 aana/models/core/image.py                  | 93 ++++++++++------------
 aana/models/pydantic/exception_response.py |  2 +-
 aana/models/pydantic/image_input.py        | 24 +++++-
 aana/tests/test_image_input.py             |  8 ++
 6 files changed, 79 insertions(+), 69 deletions(-)

diff --git a/aana/deployments/hf_blip2_deployment.py b/aana/deployments/hf_blip2_deployment.py
index 21596a55..249f8eb4 100644
--- a/aana/deployments/hf_blip2_deployment.py
+++ b/aana/deployments/hf_blip2_deployment.py
@@ -83,7 +83,7 @@ async def apply_config(self, config: Dict[str, Any]):
         # Load the model and processor for BLIP2 from HuggingFace
         self.model_id = config_obj.model
         self.dtype = config_obj.dtype
-        self.torch_dtype = Dtype.to_torch(self.dtype)
+        self.torch_dtype = self.dtype.to_torch()
         self.model = Blip2ForConditionalGeneration.from_pretrained(
             self.model_id, torch_dtype=self.torch_dtype
         )
diff --git a/aana/models/core/dtype.py b/aana/models/core/dtype.py
index 9e044279..65d152de 100644
--- a/aana/models/core/dtype.py
+++ b/aana/models/core/dtype.py
@@ -22,13 +22,9 @@ class Dtype(str, Enum):
     FLOAT16 = "float16"
     INT8 = "int8"
 
-    @classmethod
-    def to_torch(cls, dtype: "Dtype") -> Union[torch.dtype, str]:
+    def to_torch(self) -> Union[torch.dtype, str]:
         """
-        Convert a dtype to a torch dtype.
-
-        Args:
-            dtype (Dtype): the dtype
+        Convert the instance's dtype to a torch dtype.
 
         Returns:
             Union[torch.dtype, str]: the torch dtype or "auto"
@@ -36,14 +32,13 @@ def to_torch(cls, dtype: "Dtype") -> Union[torch.dtype, str]:
         Raises:
             ValueError: if the dtype is unknown
         """
-
-        if dtype == cls.AUTO:
+        if self.value == self.AUTO:
             return "auto"
-        elif dtype == cls.FLOAT32:
+        elif self.value == self.FLOAT32:
             return torch.float32
-        elif dtype == cls.FLOAT16:
+        elif self.value == self.FLOAT16:
             return torch.float16
-        elif dtype == cls.INT8:
+        elif self.value == self.INT8:
             return torch.int8
         else:
-            raise ValueError(f"Unknown dtype: {dtype}")
+            raise ValueError(f"Unknown dtype: {self}")
diff --git a/aana/models/core/image.py b/aana/models/core/image.py
index 4c8bf474..f353d273 100644
--- a/aana/models/core/image.py
+++ b/aana/models/core/image.py
@@ -147,12 +147,22 @@ class Image:
     image_lib: Type[
         AbstractImageLibrary
     ] = OpenCVWrapper  # The image library to use, TODO: add support for PIL and allow to choose the library
-    saved: bool = False  # Whether the image is saved on disc by the class or not (used for cleanup)
+    is_saved: bool = False  # Whether the image is saved on disc by the class or not (used for cleanup)
 
     def __post_init__(self):
         """
-        Save the image on disc after initialization if save_on_disc is True.
+        Post-initialization method.
+
+        Performs checks:
+        - Checks that path is a Path object.
+        - Checks that at least one of 'path', 'url', 'content' or 'numpy' is provided.
+        - Checks if path exists if provided.
+
+        Saves the image on disk if needed.
         """
+        # check that path is a Path object
+        if self.path:
+            assert isinstance(self.path, Path)
         # check that at least one of 'path', 'url', 'content' or 'numpy' is provided
         if not any(
             [
@@ -205,7 +215,7 @@ def save_image(self):
                 "At least one of 'path', 'url', 'content' or 'numpy' must be provided."
             )
         self.path = file_path
-        self.saved = True
+        self.is_saved = True
 
     def save_from_content(self, file_path: Path):
         """
@@ -235,7 +245,9 @@ def save_from_url(self, file_path: Path):
         Args:
             file_path (Path): The path of the file to write.
         """
-        self.image_lib.write_file(file_path, self.load_numpy_from_url())
+        assert self.url is not None
+        content = download_file(self.url)
+        file_path.write_bytes(content)
 
     def get_numpy(self) -> np.ndarray:
         """
@@ -251,76 +263,63 @@ def get_numpy(self) -> np.ndarray:
         if self.numpy is not None:
             return self.numpy
         elif self.path:
-            self.numpy = self.load_numpy_from_path()
+            self.load_numpy_from_path()
         elif self.url:
-            self.numpy = self.load_numpy_from_url()
+            self.load_numpy_from_url()
         elif self.content:
-            self.numpy = self.load_numpy_from_content()
+            self.load_numpy_from_content()
         else:
             raise ValueError(
                 "At least one of 'path', 'url', 'content' or 'numpy' must be provided."
             )
+        assert self.numpy is not None
         return self.numpy
 
-    def load_numpy_from_path(self) -> np.ndarray:
+    def load_numpy_from_path(self):
         """
         Load the image as a numpy array from a path.
 
-        Returns:
-            np.ndarray: The image as a numpy array.
-
         Raises:
             ImageReadingException: If there is an error reading the image.
         """
         assert self.path is not None
         try:
-            numpy = self.image_lib.read_file(self.path)
+            self.numpy = self.image_lib.read_file(self.path)
         except Exception as e:
             raise ImageReadingException(self) from e
-        return numpy
 
-    def load_numpy_from_image_bytes(self, img_bytes: bytes) -> np.ndarray:
+    def load_numpy_from_image_bytes(self, img_bytes: bytes):
         """
         Load the image as a numpy array from image bytes (downloaded from URL or read from file).
 
-        Returns:
-            np.ndarray: The image as a numpy array.
-
         Raises:
             ImageReadingException: If there is an error reading the image.
         """
         try:
-            numpy = self.image_lib.read_bytes(img_bytes)
+            self.numpy = self.image_lib.read_bytes(img_bytes)
         except Exception as e:
             raise ImageReadingException(self) from e
-        return numpy
 
-    def load_numpy_from_url(self) -> np.ndarray:
+    def load_numpy_from_url(self):
         """
         Load the image as a numpy array from a URL.
 
-        Returns:
-            np.ndarray: The image as a numpy array.
-
         Raises:
             ImageReadingException: If there is an error reading the image.
         """
         assert self.url is not None
         content: bytes = download_file(self.url)
-        return self.load_numpy_from_image_bytes(content)
+        self.load_numpy_from_image_bytes(content)
 
-    def load_numpy_from_content(self) -> np.ndarray:
+    def load_numpy_from_content(self):
         """
         Load the image as a numpy array from content.
 
-        Returns:
-            np.ndarray: The image as a numpy array.
-
         Raises:
             ImageReadingException: If there is an error reading the image.
         """
         assert self.content is not None
-        return self.load_numpy_from_image_bytes(self.content)
+        self.load_numpy_from_image_bytes(self.content)
 
     def get_content(self) -> bytes:
         """
@@ -335,57 +334,45 @@ def get_content(self) -> bytes:
         if self.content:
             return self.content
         elif self.path:
-            self.content = self.load_content_from_path()
+            self.load_content_from_path()
         elif self.url:
-            self.content = self.load_content_from_url()
+            self.load_content_from_url()
         elif self.numpy is not None:
-            self.content = self.load_content_from_numpy()
+            self.load_content_from_numpy()
         else:
             raise ValueError(
                 "At least one of 'path', 'url', 'content' or 'numpy' must be provided."
             )
+        assert self.content is not None
         return self.content
 
-    def load_content_from_numpy(self) -> bytes:
+    def load_content_from_numpy(self):
         """
         Load the content of the image from numpy.
-
-        Note: This method is not implemented yet.
-
-        Returns:
-            bytes: The content of the image as bytes.
         """
         assert self.numpy is not None
-        return self.image_lib.write_bytes(self.numpy)
+        self.content = self.image_lib.write_bytes(self.numpy)
 
-    def load_content_from_path(self) -> bytes:
+    def load_content_from_path(self):
         """
         Load the content of the image from the path.
 
-        Returns:
-            bytes: The content of the image as bytes.
-
         Raises:
             FileNotFoundError: If the image file does not exist.
         """
         assert self.path is not None
         with open(self.path, "rb") as f:
-            content = f.read()
-        return content
+            self.content = f.read()
 
-    def load_content_from_url(self) -> bytes:
+    def load_content_from_url(self):
         """
         Load the content of the image from the URL using requests.
 
-        Returns:
-            bytes: The content of the image as bytes.
-
         Raises:
             DownloadException: If there is an error downloading the image.
         """
         assert self.url is not None
-        content = download_file(self.url)
-        return content
+        self.content = download_file(self.url)
 
     def __repr__(self) -> str:
         """
@@ -429,5 +416,5 @@ def cleanup(self):
         Remove the image from disc if it was saved by the class.
         """
         # Remove the image from disc if it was saved by the class
-        if self.saved and self.path and os.path.exists(self.path):
-            os.remove(self.path)
+        if self.is_saved and self.path:
+            self.path.unlink(missing_ok=True)
diff --git a/aana/models/pydantic/exception_response.py b/aana/models/pydantic/exception_response.py
index 6bb9f635..0ced696b 100644
--- a/aana/models/pydantic/exception_response.py
+++ b/aana/models/pydantic/exception_response.py
@@ -9,7 +9,7 @@ class ExceptionResponseModel(BaseModel):
     Attributes:
         error (str): The error that occurred.
         message (str): The message of the error.
-        data (Optional[Dict]): The extra data that helps to debug the error.
+        data (Optional[Any]): The extra data that helps to debug the error.
         stacktrace (Optional[str]): The stacktrace of the error.
     """
 
diff --git a/aana/models/pydantic/image_input.py b/aana/models/pydantic/image_input.py
index 1af8474a..9019abd3 100644
--- a/aana/models/pydantic/image_input.py
+++ b/aana/models/pydantic/image_input.py
@@ -2,7 +2,7 @@
 from pathlib import Path
 import numpy as np
 from typing import Dict, List, Optional
-from pydantic import BaseModel, Field, ValidationError, root_validator
+from pydantic import BaseModel, Field, ValidationError, root_validator, validator
 from pydantic.error_wrappers import ErrorWrapper
 
 from aana.models.core.image import Image
@@ -27,7 +27,9 @@ class ImageInput(BaseModel):
     """
 
     path: Optional[str] = Field(None, description="The file path of the image.")
-    url: Optional[str] = Field(None, description="The URL of the image.")
+    url: Optional[str] = Field(
+        None, description="The URL of the image."
+    )  # TODO: validate url
     content: Optional[bytes] = Field(
         None,
         description=(
@@ -161,6 +163,24 @@ class ImageListInput(BaseListModel):
 
     __root__: List[ImageInput]
 
+    @validator("__root__", pre=True)
+    def check_non_empty(cls, v: List[ImageInput]) -> List[ImageInput]:
+        """
+        Check that the list of images isn't empty.
+
+        Args:
+            v (List[ImageInput]): the list of images
+
+        Returns:
+            List[ImageInput]: the list of images
+
+        Raises:
+            ValueError: if the list of images is empty
+        """
+        if len(v) == 0:
+            raise ValueError("The list of images must not be empty.")
+        return v
+
     def set_files(self, files: List[bytes]):
         """
         Set the files for the images.
diff --git a/aana/tests/test_image_input.py b/aana/tests/test_image_input.py
index 3a1cced8..1591c187 100644
--- a/aana/tests/test_image_input.py
+++ b/aana/tests/test_image_input.py
@@ -242,3 +242,11 @@ def test_imagelistinput_set_files():
 
     assert image_list_input[0].content == files[0]
     assert image_list_input[1].numpy == files[1]
+
+
+def test_imagelistinput_non_empty():
+    """
+    Test that ImageListInput must not be empty.
+    """
+    with pytest.raises(ValidationError):
+        ImageListInput(__root__=[])

From 26142e92ce092f3a47ffd4d5481a5f59ad281074 Mon Sep 17 00:00:00 2001
From: Aleksandr Movchan <movchan174@gmail.com>
Date: Thu, 9 Nov 2023 15:25:04 +0100
Subject: [PATCH 2/4] Update issue templates

---
 .github/ISSUE_TEMPLATE/bug_report.md      | 29 +++++++++++++++++++++++
 .github/ISSUE_TEMPLATE/documentation.md   | 21 ++++++++++++++++
 .github/ISSUE_TEMPLATE/enhancement.md     | 17 +++++++++++++
 .github/ISSUE_TEMPLATE/feature_request.md | 17 +++++++++++++
 .github/ISSUE_TEMPLATE/question.md        | 18 ++++++++++++++
 .github/ISSUE_TEMPLATE/testing.md         | 18 ++++++++++++++
 6 files changed, 120 insertions(+)
 create mode 100644 .github/ISSUE_TEMPLATE/bug_report.md
 create mode 100644 .github/ISSUE_TEMPLATE/documentation.md
 create mode 100644 .github/ISSUE_TEMPLATE/enhancement.md
 create mode 100644 .github/ISSUE_TEMPLATE/feature_request.md
 create mode 100644 .github/ISSUE_TEMPLATE/question.md
 create mode 100644 .github/ISSUE_TEMPLATE/testing.md

diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
new file mode 100644
index 00000000..153c5a93
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -0,0 +1,29 @@
+---
+name: Bug report
+about: Report a malfunction, glitch, or error in the application. This includes any
+  performance-related issues that may arise.
+title: "[BUG]"
+labels: bug
+assignees: ''
+
+---
+
+### Bug Description
+- Brief summary of the issue
+
+### Steps to Reproduce
+1.
+2.
+3.
+
+### Expected Behavior
+- What should have happened?
+
+### Actual Behavior
+- What actually happened?
+
+### Performance Details (if applicable)
+- Specifics of the performance issue encountered
+
+### Environment
+- Version (commit hash, tag, branch)
diff --git a/.github/ISSUE_TEMPLATE/documentation.md b/.github/ISSUE_TEMPLATE/documentation.md
new file mode 100644
index 00000000..43230c39
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/documentation.md
@@ -0,0 +1,21 @@
+---
+name: Documentation
+about: Propose updates or corrections to both internal development documentation and
+  client-facing documentation.
+title: "[DOCS]"
+labels: documentation
+assignees: ''
+
+---
+
+### Documentation Area (Development/Client)
+- Specify the area of documentation
+
+### Current Content
+- Quote the current content or describe the issue
+
+### Proposed Changes
+- Detail the proposed changes
+
+### Reasons for Changes
+- Why these changes will improve the documentation
diff --git a/.github/ISSUE_TEMPLATE/enhancement.md b/.github/ISSUE_TEMPLATE/enhancement.md
new file mode 100644
index 00000000..33bdbadf
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/enhancement.md
@@ -0,0 +1,17 @@
+---
+name: Enhancement
+about: Recommend improvement to existing features and suggest code quality improvement.
+title: "[ENHANCEMENT]"
+labels: enhancement
+assignees: ''
+
+---
+
+### Enhancement Description
+- Overview of the enhancement
+
+### Advantages
+- Benefits of implementing this enhancement
+
+### Possible Implementation
+- Suggested methods for implementing the enhancement
diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md
new file mode 100644
index 00000000..32f74f10
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/feature_request.md
@@ -0,0 +1,17 @@
+---
+name: Feature request
+about: Suggest new functionalities or modifications to enhance the application's capabilities.
+title: "[FEATURE REQUEST]"
+labels: feature request
+assignees: ''
+
+---
+
+### Feature Summary
+- Concise description of the feature
+
+### Justification/Rationale
+- Why is the feature beneficial?
+
+### Proposed Implementation (if any)
+- How do you envision this feature's implementation?
diff --git a/.github/ISSUE_TEMPLATE/question.md b/.github/ISSUE_TEMPLATE/question.md
new file mode 100644
index 00000000..ecb0dd64
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/question.md
@@ -0,0 +1,18 @@
+---
+name: Question
+about: Ask questions to clarify project-related queries, seek further information,
+  or understand functionalities better.
+title: "[QUESTION]"
+labels: question
+assignees: ''
+
+---
+
+### Context
+- Background or context of the question
+
+### Question
+- Specific question being asked
+
+### What You've Tried
+- List any solutions or research already conducted
diff --git a/.github/ISSUE_TEMPLATE/testing.md b/.github/ISSUE_TEMPLATE/testing.md
new file mode 100644
index 00000000..f09852b7
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/testing.md
@@ -0,0 +1,18 @@
+---
+name: Testing
+about: Address needs for creating new tests, enhancing existing tests, or reporting
+  test failures.
+title: "[TESTS]"
+labels: tests
+assignees: ''
+
+---
+
+### Testing Requirement
+- Describe the testing need or issue
+
+### Test Scenarios
+- Detail specific test scenarios to be addressed
+
+### Acceptance Criteria
+- What are the criteria for the test to be considered successful?

From bc1c60971da1674672e9f043d3b6d140e4f57cd8 Mon Sep 17 00:00:00 2001
From: Aleksandr Movchan <aleksandr@mobiuslabs.com>
Date: Thu, 9 Nov 2023 16:11:54 +0000
Subject: [PATCH 3/4] Applied requested changes from PR

---
 aana/configs/settings.py                      |  2 +-
 aana/deployments/hf_blip2_deployment.py       | 23 ++++++---
 aana/deployments/vllm_deployment.py           | 47 ++++++++++++++----
 aana/models/core/dtype.py                     | 21 ++++----
 aana/models/core/image.py                     | 48 ++++++++++---------
 .../deployments/test_hf_blip2_deployment.py   |  2 +-
 aana/tests/test_batch_processor.py            | 19 +++++---
 aana/tests/test_image.py                      | 26 +++++-----
 aana/tests/test_settings.py                   |  2 +-
 9 files changed, 118 insertions(+), 72 deletions(-)

diff --git a/aana/configs/settings.py b/aana/configs/settings.py
index 36716463..44a661dd 100644
--- a/aana/configs/settings.py
+++ b/aana/configs/settings.py
@@ -8,7 +8,7 @@ class Settings(BaseSettings):
 
     """
 
-    tmp_data_dir: Path = Path("/tmp/aana/data")
+    tmp_data_dir: Path = Path("/tmp/aana_data")
 
 
 settings = Settings()
diff --git a/aana/deployments/hf_blip2_deployment.py b/aana/deployments/hf_blip2_deployment.py
index 249f8eb4..800bfe14 100644
--- a/aana/deployments/hf_blip2_deployment.py
+++ b/aana/deployments/hf_blip2_deployment.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, TypedDict
 from pydantic import BaseModel, Field, validator
 from ray import serve
 import torch
@@ -48,6 +48,17 @@ def validate_dtype(cls, value: Dtype) -> Dtype:
         return value
 
 
+class CaptioningOutput(TypedDict):
+    """
+    The output of the captioning model.
+
+    Attributes:
+        captions (List[str]): the list of captions
+    """
+
+    captions: List[str]
+
+
 @serve.deployment
 class HFBlip2Deployment(BaseDeployment):
     """
@@ -92,7 +103,7 @@ async def apply_config(self, config: Dict[str, Any]):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model.to(self.device)
 
-    async def generate_captions(self, **kwargs) -> Dict[str, Any]:
+    async def generate_captions(self, **kwargs) -> CaptioningOutput:
         """
         Generate captions for the given images.
 
@@ -100,7 +111,7 @@ async def generate_captions(self, **kwargs) -> Dict[str, Any]:
             images (List[Image]): the images
 
         Returns:
-            Dict[str, Any]: the dictionary with one key "captions"
+            CaptioningOutput: the dictionary with one key "captions"
                             and the list of captions for the images as value
 
         Raises:
@@ -110,7 +121,7 @@ async def generate_captions(self, **kwargs) -> Dict[str, Any]:
         # The actual inference is done in _generate_captions()
         return await self.batch_processor.process(kwargs)
 
-    def _generate_captions(self, images: List[Image]) -> Dict[str, Any]:
+    def _generate_captions(self, images: List[Image]) -> CaptioningOutput:
         """
         Generate captions for the given images.
 
@@ -120,7 +131,7 @@ def _generate_captions(self, images: List[Image]) -> Dict[str, Any]:
             images (List[Image]): the images
 
         Returns:
-            Dict[str, Any]: the dictionary with one key "captions"
+            CaptioningOutput: the dictionary with one key "captions"
                             and the list of captions for the images as value
 
         Raises:
@@ -141,6 +152,6 @@ def _generate_captions(self, images: List[Image]) -> Dict[str, Any]:
             generated_texts = [
                 generated_text.strip() for generated_text in generated_texts
             ]
-            return {"captions": generated_texts}
+            return CaptioningOutput(captions=generated_texts)
         except Exception as e:
             raise InferenceException(self.model_id) from e
diff --git a/aana/deployments/vllm_deployment.py b/aana/deployments/vllm_deployment.py
index 1175c783..a58a9d69 100644
--- a/aana/deployments/vllm_deployment.py
+++ b/aana/deployments/vllm_deployment.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, AsyncGenerator, Dict, List, Optional, TypedDict
 from pydantic import BaseModel, Field
 from ray import serve
 from vllm.engine.arg_utils import AsyncEngineArgs
@@ -34,6 +34,28 @@ class VLLMConfig(BaseModel):
     max_model_len: Optional[int] = Field(default=None)
 
 
+class LLMOutput(TypedDict):
+    """
+    The output of the LLM model.
+
+    Attributes:
+        text (str): the generated text
+    """
+
+    text: str
+
+
+class LLMBatchOutput(TypedDict):
+    """
+    The output of the LLM model for a batch of inputs.
+
+    Attributes:
+        texts (List[str]): the list of generated texts
+    """
+
+    texts: List[str]
+
+
 @serve.deployment
 class VLLMDeployment(BaseDeployment):
     """
@@ -78,7 +100,9 @@ async def apply_config(self, config: Dict[str, Any]):
         # create the engine
         self.engine = AsyncLLMEngine.from_engine_args(args)
 
-    async def generate_stream(self, prompt: str, sampling_params: SamplingParams):
+    async def generate_stream(
+        self, prompt: str, sampling_params: SamplingParams
+    ) -> AsyncGenerator[LLMOutput, None]:
         """
         Generate completion for the given prompt and stream the results.
 
@@ -87,7 +111,7 @@ async def generate_stream(self, prompt: str, sampling_params: SamplingParams):
             sampling_params (SamplingParams): the sampling parameters
 
         Yields:
-            dict: the generated text
+            LLMOutput: the dictionary with the key "text" and the generated text as the value
         """
         prompt = str(prompt)
         sampling_params = merged_options(self.default_sampling_params, sampling_params)
@@ -108,7 +132,7 @@ async def generate_stream(self, prompt: str, sampling_params: SamplingParams):
             num_returned = 0
             async for request_output in results_generator:
                 text_output = request_output.outputs[0].text[num_returned:]
-                yield {"text": text_output}
+                yield LLMOutput(text=text_output)
                 num_returned += len(text_output)
         except GeneratorExit as e:
             # If the generator is cancelled, we need to cancel the request
@@ -118,7 +142,7 @@ async def generate_stream(self, prompt: str, sampling_params: SamplingParams):
         except Exception as e:
             raise InferenceException(model_name=self.model) from e
 
-    async def generate(self, prompt: str, sampling_params: SamplingParams):
+    async def generate(self, prompt: str, sampling_params: SamplingParams) -> LLMOutput:
         """
         Generate completion for the given prompt.
 
@@ -127,14 +151,16 @@ async def generate(self, prompt: str, sampling_params: SamplingParams):
             sampling_params (SamplingParams): the sampling parameters
 
         Returns:
-            dict: the generated text
+            LLMOutput: the dictionary with the key "text" and the generated text as the value
         """
         generated_text = ""
         async for chunk in self.generate_stream(prompt, sampling_params):
             generated_text += chunk["text"]
-        return {"text": generated_text}
+        return LLMOutput(text=generated_text)
 
-    async def generate_batch(self, prompts: List[str], sampling_params: SamplingParams):
+    async def generate_batch(
+        self, prompts: List[str], sampling_params: SamplingParams
+    ) -> LLMBatchOutput:
         """
         Generate completion for the batch of prompts.
 
@@ -143,11 +169,12 @@ async def generate_batch(self, prompts: List[str], sampling_params: SamplingPara
             sampling_params (SamplingParams): the sampling parameters
 
         Returns:
-            dict: the generated texts
+            LLMBatchOutput: the dictionary with the key "texts"
+                            and the list of generated texts as the value
         """
         texts = []
         for prompt in prompts:
             text = await self.generate(prompt, sampling_params)
             texts.append(text["text"])
 
-        return {"texts": texts}
+        return LLMBatchOutput(texts=texts)
diff --git a/aana/models/core/dtype.py b/aana/models/core/dtype.py
index 65d152de..c656c6c1 100644
--- a/aana/models/core/dtype.py
+++ b/aana/models/core/dtype.py
@@ -32,13 +32,14 @@ def to_torch(self) -> Union[torch.dtype, str]:
         Raises:
             ValueError: if the dtype is unknown
         """
-        if self.value == self.AUTO:
-            return "auto"
-        elif self.value == self.FLOAT32:
-            return torch.float32
-        elif self.value == self.FLOAT16:
-            return torch.float16
-        elif self.value == self.INT8:
-            return torch.int8
-        else:
-            raise ValueError(f"Unknown dtype: {self}")
+        match self.value:
+            case self.AUTO:
+                return "auto"
+            case self.FLOAT32:
+                return torch.float32
+            case self.FLOAT16:
+                return torch.float16
+            case self.INT8:
+                return torch.int8
+            case _:
+                raise ValueError(f"Unknown dtype: {self}")
diff --git a/aana/models/core/image.py b/aana/models/core/image.py
index f353d273..22ae8dbe 100644
--- a/aana/models/core/image.py
+++ b/aana/models/core/image.py
@@ -32,7 +32,7 @@ def read_file(cls, path: Path) -> np.ndarray:
         raise NotImplementedError
 
     @classmethod
-    def read_bytes(cls, content: bytes) -> np.ndarray:
+    def read_from_bytes(cls, content: bytes) -> np.ndarray:
         """
         Read bytes using the image library.
 
@@ -56,7 +56,7 @@ def write_file(cls, path: Path, img: np.ndarray):
         raise NotImplementedError
 
     @classmethod
-    def write_bytes(cls, img: np.ndarray) -> bytes:
+    def write_to_bytes(cls, img: np.ndarray) -> bytes:
         """
         Write bytes using the image library.
 
@@ -90,7 +90,7 @@ def read_file(cls, path: Path) -> np.ndarray:
         return img
 
     @classmethod
-    def read_bytes(cls, content: bytes) -> np.ndarray:
+    def read_from_bytes(cls, content: bytes) -> np.ndarray:
         """
         Read bytes using OpenCV.
 
@@ -117,7 +117,7 @@ def write_file(cls, path: Path, img: np.ndarray):
         cv2.imwrite(str(path), img)
 
     @classmethod
-    def write_bytes(cls, img: np.ndarray) -> bytes:
+    def write_to_bytes(cls, img: np.ndarray) -> bytes:
         """
         Write bytes using OpenCV.
 
@@ -143,11 +143,11 @@ class Image:
     image_id: Optional[str] = field(
         default_factory=lambda: str(uuid.uuid4())
     )  # The ID of the image, generated automatically
-    save_on_disc: bool = True  # Whether to save the image on disc or not
+    save_on_disk: bool = True  # Whether to save the image on disk or not
     image_lib: Type[
         AbstractImageLibrary
     ] = OpenCVWrapper  # The image library to use, TODO: add support for PIL and allow to choose the library
-    is_saved: bool = False  # Whether the image is saved on disc by the class or not (used for cleanup)
+    is_saved: bool = False  # Whether the image is saved on disk by the class or not (used for cleanup)
 
     def __post_init__(self):
         """
@@ -162,7 +162,9 @@ def __post_init__(self):
         """
         # check that path is a Path object
         if self.path:
-            assert isinstance(self.path, Path)
+            if not isinstance(self.path, Path):
+                raise ValueError("Path must be a Path object.")
+
         # check that at least one of 'path', 'url', 'content' or 'numpy' is provided
         if not any(
             [
@@ -180,18 +182,18 @@ def __post_init__(self):
         if self.path and not self.path.exists():
             raise FileNotFoundError(f"Image file not found: {self.path}")
 
-        if self.save_on_disc:
-            self.save_image()
+        if self.save_on_disk:
+            self.save()
 
-    def save_image(self):
+    def save(self):
         """
-        Save the image on disc.
-        If the image is already available on disc, do nothing.
-        If the image represented as a byte string, save it on disc.
-        If the image is represented as a URL, download it and save it on disc.
-        If the image is represented as a numpy array, convert it to BMP and save it on disc.
+        Save the image on disk.
+        If the image is already available on disk, do nothing.
+        If the image represented as a byte string, save it on disk.
+        If the image is represented as a URL, download it and save it on disk.
+        If the image is represented as a numpy array, convert it to BMP and save it on disk.
 
-        First check if the image is already available on disc, then content, then url, then numpy
+        First check if the image is already available on disk, then content, then url, then numpy
         to avoid unnecessary operations (e.g. downloading the image or converting it to BMP).
 
         Raises:
@@ -219,7 +221,7 @@ def save_image(self):
 
     def save_from_content(self, file_path: Path):
         """
-        Save the image from content on disc.
+        Save the image from content on disk.
 
         Args:
             file_path (Path): The path of the file to write.
@@ -230,7 +232,7 @@ def save_from_content(self, file_path: Path):
 
     def save_from_numpy(self, file_path: Path):
         """
-        Save the image from numpy on disc.
+        Save the image from numpy on disk.
 
         Args:
             file_path (Path): The path of the file to write.
@@ -240,7 +242,7 @@ def save_from_numpy(self, file_path: Path):
 
     def save_from_url(self, file_path: Path):
         """
-        Save the image from URL on disc.
+        Save the image from URL on disk.
 
         Args:
             file_path (Path): The path of the file to write.
@@ -296,7 +298,7 @@ def load_numpy_from_image_bytes(self, img_bytes: bytes):
             ImageReadingException: If there is an error reading the image.
         """
         try:
-            self.numpy = self.image_lib.read_bytes(img_bytes)
+            self.numpy = self.image_lib.read_from_bytes(img_bytes)
         except Exception as e:
             raise ImageReadingException(self) from e
 
@@ -351,7 +353,7 @@ def load_content_from_numpy(self):
         Load the content of the image from numpy.
         """
         assert self.numpy is not None
-        self.content = self.image_lib.write_bytes(self.numpy)
+        self.content = self.image_lib.write_to_bytes(self.numpy)
 
     def load_content_from_path(self):
         """
@@ -413,8 +415,8 @@ def cleanup(self):
         """
         Cleanup the image.
 
-        Remove the image from disc if it was saved by the class.
+        Remove the image from disk if it was saved by the class.
         """
-        # Remove the image from disc if it was saved by the class
+        # Remove the image from disk if it was saved by the class
         if self.is_saved and self.path:
             self.path.unlink(missing_ok=True)
diff --git a/aana/tests/deployments/test_hf_blip2_deployment.py b/aana/tests/deployments/test_hf_blip2_deployment.py
index 787c87ca..3c501967 100644
--- a/aana/tests/deployments/test_hf_blip2_deployment.py
+++ b/aana/tests/deployments/test_hf_blip2_deployment.py
@@ -30,7 +30,7 @@ async def test_hf_blip2_deployments():
         handle = ray_setup(deployment)
 
         path = resources.path("aana.tests.files.images", "Starry_Night.jpeg")
-        image = Image(path=path, save_on_disc=False)
+        image = Image(path=path, save_on_disk=False)
 
         images = [image] * 8
 
diff --git a/aana/tests/test_batch_processor.py b/aana/tests/test_batch_processor.py
index d6a888be..c2ff6923 100644
--- a/aana/tests/test_batch_processor.py
+++ b/aana/tests/test_batch_processor.py
@@ -3,20 +3,24 @@
 
 from aana.utils.batch_processor import BatchProcessor
 
+NUM_IMAGES = 5
+IMAGE_SIZE = 100
+FEATURE_SIZE = 10
+
 
 @pytest.fixture
 def images():
-    return np.array([np.random.rand(100, 100) for _ in range(5)])  # 5 images
+    return np.array([np.random.rand(IMAGE_SIZE, IMAGE_SIZE) for _ in range(NUM_IMAGES)])
 
 
 @pytest.fixture
 def texts():
-    return ["text1", "text2", "text3", "text4", "text5"]
+    return [f"text{i}" for i in range(NUM_IMAGES)]
 
 
 @pytest.fixture
 def features():
-    return np.random.rand(5, 10)  # 5 features of size 10
+    return np.random.rand(NUM_IMAGES, FEATURE_SIZE)
 
 
 @pytest.fixture
@@ -40,7 +44,8 @@ def test_batch_iterator(request_batch, process_batch):
     )
 
     batches = list(processor.batch_iterator(request_batch))
-    assert len(batches) == 3  # We expect 3 batches with a batch_size of 2 for 5 items
+    # We expect 3 batches with a batch_size of 2 for 5 items
+    assert len(batches) == NUM_IMAGES // batch_size + 1
     assert all(
         len(batch["texts"]) == batch_size for batch in batches[:-1]
     )  # All but the last should have batch_size items
@@ -59,11 +64,11 @@ async def test_process_batches(request_batch, process_batch):
 
     result = await processor.process(request_batch)
     # Ensure all texts are processed
-    assert len(result["texts"]) == 5
+    assert len(result["texts"]) == NUM_IMAGES
     # Check if images are concatenated properly
-    assert result["images"].shape == (5, 100, 100)
+    assert result["images"].shape == (NUM_IMAGES, IMAGE_SIZE, IMAGE_SIZE)
     # Check if features are concatenated properly
-    assert result["features"].shape == (5, 10)
+    assert result["features"].shape == (NUM_IMAGES, FEATURE_SIZE)
 
 
 def test_merge_outputs(request_batch, process_batch):
diff --git a/aana/tests/test_image.py b/aana/tests/test_image.py
index 5a88d0b8..e1cfc767 100644
--- a/aana/tests/test_image.py
+++ b/aana/tests/test_image.py
@@ -9,7 +9,7 @@ def load_numpy_from_image_bytes(content: bytes) -> np.ndarray:
     """
     Load a numpy array from image bytes.
     """
-    image = Image(content=content, save_on_disc=False)
+    image = Image(content=content, save_on_disk=False)
     return image.get_numpy()
 
 
@@ -32,7 +32,7 @@ def test_image(mock_download_file):
 
     try:
         path = resources.path("aana.tests.files.images", "Starry_Night.jpeg")
-        image = Image(path=path, save_on_disc=False)
+        image = Image(path=path, save_on_disk=False)
         assert image.path == path
         assert image.content is None
         assert image.numpy is None
@@ -49,7 +49,7 @@ def test_image(mock_download_file):
 
     try:
         url = "http://example.com/Starry_Night.jpeg"
-        image = Image(url=url, save_on_disc=False)
+        image = Image(url=url, save_on_disk=False)
         assert image.path is None
         assert image.content is None
         assert image.numpy is None
@@ -67,7 +67,7 @@ def test_image(mock_download_file):
     try:
         path = resources.path("aana.tests.files.images", "Starry_Night.jpeg")
         content = path.read_bytes()
-        image = Image(content=content, save_on_disc=False)
+        image = Image(content=content, save_on_disk=False)
         assert image.path is None
         assert image.content == content
         assert image.numpy is None
@@ -84,7 +84,7 @@ def test_image(mock_download_file):
 
     try:
         numpy = np.random.rand(100, 100, 3).astype(np.uint8)
-        image = Image(numpy=numpy, save_on_disc=False)
+        image = Image(numpy=numpy, save_on_disk=False)
         assert image.path is None
         assert image.content is None
         assert np.array_equal(image.numpy, numpy)
@@ -111,12 +111,12 @@ def test_image_path_not_exist():
 
 def test_save_image(mock_download_file):
     """
-    Test that save_on_disc works.
+    Test that save_on_disk works.
     """
 
     try:
         path = resources.path("aana.tests.files.images", "Starry_Night.jpeg")
-        image = Image(path=path, save_on_disc=True)
+        image = Image(path=path, save_on_disk=True)
         assert image.path == path
         assert image.content is None
         assert image.numpy is None
@@ -127,7 +127,7 @@ def test_save_image(mock_download_file):
 
     try:
         url = "http://example.com/Starry_Night.jpeg"
-        image = Image(url=url, save_on_disc=True)
+        image = Image(url=url, save_on_disk=True)
         assert image.content is None
         assert image.numpy is None
         assert image.url == url
@@ -138,7 +138,7 @@ def test_save_image(mock_download_file):
     try:
         path = resources.path("aana.tests.files.images", "Starry_Night.jpeg")
         content = path.read_bytes()
-        image = Image(content=content, save_on_disc=True)
+        image = Image(content=content, save_on_disk=True)
         assert image.content == content
         assert image.numpy is None
         assert image.url is None
@@ -148,7 +148,7 @@ def test_save_image(mock_download_file):
 
     try:
         numpy = np.random.rand(100, 100, 3).astype(np.uint8)
-        image = Image(numpy=numpy, save_on_disc=True)
+        image = Image(numpy=numpy, save_on_disk=True)
         assert image.content is None
         assert np.array_equal(image.numpy, numpy)
         assert image.url is None
@@ -164,7 +164,7 @@ def test_cleanup(mock_download_file):
 
     try:
         url = "http://example.com/Starry_Night.jpeg"
-        image = Image(url=url, save_on_disc=True)
+        image = Image(url=url, save_on_disk=True)
         assert image.path.exists()
     finally:
         image.cleanup()
@@ -176,7 +176,7 @@ def test_at_least_one_input():
     Test that at least one input is provided.
     """
     with pytest.raises(ValueError):
-        Image(save_on_disc=False)
+        Image(save_on_disk=False)
 
     with pytest.raises(ValueError):
-        Image(save_on_disc=True)
+        Image(save_on_disk=True)
diff --git a/aana/tests/test_settings.py b/aana/tests/test_settings.py
index d69a457a..cbd16f3c 100644
--- a/aana/tests/test_settings.py
+++ b/aana/tests/test_settings.py
@@ -5,7 +5,7 @@
 def test_default_tmp_data_dir():
     """Test that the default temporary data directory is set correctly."""
     settings = Settings()
-    assert settings.tmp_data_dir == Path("/tmp/aana/data")
+    assert settings.tmp_data_dir == Path("/tmp/aana_data")
 
 
 def test_custom_tmp_data_dir(monkeypatch):

From 8e93fded437ced09cf12260852c84a4cac632c94 Mon Sep 17 00:00:00 2001
From: Aleksandr Movchan <aleksandr@mobiuslabs.com>
Date: Fri, 10 Nov 2023 14:10:52 +0000
Subject: [PATCH 4/4] Changed methods of BLIP2 deployment to be consistent with
 VLLM deployment

---
 aana/configs/pipeline.py                      |  2 +-
 aana/deployments/hf_blip2_deployment.py       | 48 +++++++++++++++----
 .../deployments/test_hf_blip2_deployment.py   | 16 +++++--
 3 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/aana/configs/pipeline.py b/aana/configs/pipeline.py
index a8e39454..81ca9857 100644
--- a/aana/configs/pipeline.py
+++ b/aana/configs/pipeline.py
@@ -155,7 +155,7 @@
         "name": "hf_blip2_opt_2_7b",
         "type": "ray_deployment",
         "deployment_name": "hf_blip2_deployment_opt_2_7b",
-        "method": "generate_captions",
+        "method": "generate_batch",
         "inputs": [
             {
                 "name": "images",
diff --git a/aana/deployments/hf_blip2_deployment.py b/aana/deployments/hf_blip2_deployment.py
index 800bfe14..f5b5c677 100644
--- a/aana/deployments/hf_blip2_deployment.py
+++ b/aana/deployments/hf_blip2_deployment.py
@@ -52,6 +52,17 @@ class CaptioningOutput(TypedDict):
     """
     The output of the captioning model.
 
+    Attributes:
+        caption (str): the caption
+    """
+
+    caption: str
+
+
+class CaptioningBatchOutput(TypedDict):
+    """
+    The output of the captioning model.
+
     Attributes:
         captions (List[str]): the list of captions
     """
@@ -82,11 +93,11 @@ async def apply_config(self, config: Dict[str, Any]):
         # and process them in parallel
         self.batch_size = config_obj.batch_size
         self.num_processing_threads = config_obj.num_processing_threads
-        # The actual inference is done in _generate_captions()
+        # The actual inference is done in _generate()
         # We use lambda because BatchProcessor expects dict as input
-        # and we use **kwargs to unpack the dict into named arguments for _generate_captions()
+        # and we use **kwargs to unpack the dict into named arguments for _generate()
         self.batch_processor = BatchProcessor(
-            process_batch=lambda request: self._generate_captions(**request),
+            process_batch=lambda request: self._generate(**request),
             batch_size=self.batch_size,
             num_threads=self.num_processing_threads,
         )
@@ -103,7 +114,26 @@ async def apply_config(self, config: Dict[str, Any]):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model.to(self.device)
 
-    async def generate_captions(self, **kwargs) -> CaptioningOutput:
+    async def generate(self, image: Image) -> CaptioningOutput:
+        """
+        Generate captions for the given image.
+
+        Args:
+            image (Image): the image
+
+        Returns:
+            CaptioningOutput: the dictionary with one key "captions"
+                            and the list of captions for the image as value
+
+        Raises:
+            InferenceException: if the inference fails
+        """
+        captions: CaptioningBatchOutput = await self.batch_processor.process(
+            {"images": [image]}
+        )
+        return CaptioningOutput(caption=captions["captions"][0])
+
+    async def generate_batch(self, **kwargs) -> CaptioningBatchOutput:
         """
         Generate captions for the given images.
 
@@ -111,17 +141,17 @@ async def generate_captions(self, **kwargs) -> CaptioningOutput:
             images (List[Image]): the images
 
         Returns:
-            CaptioningOutput: the dictionary with one key "captions"
+            CaptioningBatchOutput: the dictionary with one key "captions"
                             and the list of captions for the images as value
 
         Raises:
             InferenceException: if the inference fails
         """
         # Call the batch processor to process the requests
-        # The actual inference is done in _generate_captions()
+        # The actual inference is done in _generate()
         return await self.batch_processor.process(kwargs)
 
-    def _generate_captions(self, images: List[Image]) -> CaptioningOutput:
+    def _generate(self, images: List[Image]) -> CaptioningBatchOutput:
         """
         Generate captions for the given images.
 
@@ -131,7 +161,7 @@ def _generate_captions(self, images: List[Image]) -> CaptioningOutput:
             images (List[Image]): the images
 
         Returns:
-            CaptioningOutput: the dictionary with one key "captions"
+            CaptioningBatchOutput: the dictionary with one key "captions"
                             and the list of captions for the images as value
 
         Raises:
@@ -152,6 +182,6 @@ def _generate_captions(self, images: List[Image]) -> CaptioningOutput:
             generated_texts = [
                 generated_text.strip() for generated_text in generated_texts
             ]
-            return CaptioningOutput(captions=generated_texts)
+            return CaptioningBatchOutput(captions=generated_texts)
         except Exception as e:
             raise InferenceException(self.model_id) from e
diff --git a/aana/tests/deployments/test_hf_blip2_deployment.py b/aana/tests/deployments/test_hf_blip2_deployment.py
index 3c501967..5fc94d11 100644
--- a/aana/tests/deployments/test_hf_blip2_deployment.py
+++ b/aana/tests/deployments/test_hf_blip2_deployment.py
@@ -21,7 +21,11 @@ def ray_setup(deployment):
 
 @pytest.mark.skipif(not is_gpu_available(), reason="GPU is not available")
 @pytest.mark.asyncio
-async def test_hf_blip2_deployments():
+@pytest.mark.parametrize(
+    "image_name, expected_text",
+    [("Starry_Night.jpeg", "the starry night by vincent van gogh")],
+)
+async def test_hf_blip2_deployments(image_name, expected_text):
     for name, deployment in deployments.items():
         # skip if not a VLLM deployment
         if deployment.name != "HFBlip2Deployment":
@@ -29,14 +33,18 @@ async def test_hf_blip2_deployments():
 
         handle = ray_setup(deployment)
 
-        path = resources.path("aana.tests.files.images", "Starry_Night.jpeg")
+        path = resources.path("aana.tests.files.images", image_name)
         image = Image(path=path, save_on_disk=False)
 
+        output = await handle.generate.remote(image=image)
+        caption = output["caption"]
+        compare_texts(expected_text, caption)
+
         images = [image] * 8
 
-        output = await handle.generate_captions.remote(images=images)
+        output = await handle.generate_batch.remote(images=images)
         captions = output["captions"]
 
         assert len(captions) == 8
         for caption in captions:
-            compare_texts("the starry night by vincent van gogh", caption)
+            compare_texts(expected_text, caption)