From a441dade284aa9e682c652430a739d7298c3e82e Mon Sep 17 00:00:00 2001
From: Alexandre Marques <alexandre@neuralmagic.com>
Date: Mon, 4 Nov 2024 20:01:26 +0000
Subject: [PATCH] Quality fixes

---
 src/guidellm/request/emulated.py | 4 +++-
 src/guidellm/utils/__init__.py   | 2 ++
 src/guidellm/utils/images.py     | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/guidellm/request/emulated.py b/src/guidellm/request/emulated.py
index 9dc3825..f15387e 100644
--- a/src/guidellm/request/emulated.py
+++ b/src/guidellm/request/emulated.py
@@ -405,6 +405,8 @@ def sample_prompt(self, tokens: int) -> str:
 
 
     def sample_images(self):
-        image_indices = self._rng.choice(len(self._images), size=self._config.images, replace=False)
+        image_indices = self._rng.choice(
+            len(self._images), size=self._config.images, replace=False,
+        )
 
         return [self._images[i] for i in image_indices]
diff --git a/src/guidellm/utils/__init__.py b/src/guidellm/utils/__init__.py
index 1e51f22..eb4931b 100644
--- a/src/guidellm/utils/__init__.py
+++ b/src/guidellm/utils/__init__.py
@@ -38,4 +38,6 @@
     "resolve_transformers_dataset_split",
     "split_lines_by_punctuation",
     "split_text",
+    "ImageDescriptor",
+    "load_images",
 ]
diff --git a/src/guidellm/utils/images.py b/src/guidellm/utils/images.py
index 5d73bc0..5c5a727 100644
--- a/src/guidellm/utils/images.py
+++ b/src/guidellm/utils/images.py
@@ -65,4 +65,4 @@ def load_images(data: str) -> List[ImageDescriptor]:
                     )
                 )
 
-        return images
+    return images