From df5016971011bf107f6ecb90f7242bd7e044749d Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Mon, 12 Feb 2024 10:24:39 -0500 Subject: [PATCH] Support both str and Path jsonl_path inputs in ImageCaptionJsonlDataset. --- .../_shared/data/datasets/image_caption_jsonl_dataset.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/invoke_training/_shared/data/datasets/image_caption_jsonl_dataset.py b/src/invoke_training/_shared/data/datasets/image_caption_jsonl_dataset.py index 5cd6a3ed..5ef3405b 100644 --- a/src/invoke_training/_shared/data/datasets/image_caption_jsonl_dataset.py +++ b/src/invoke_training/_shared/data/datasets/image_caption_jsonl_dataset.py @@ -21,10 +21,13 @@ class ImageCaptionJsonlDataset(torch.utils.data.Dataset): """A dataset that loads images and captions from a directory of image files and .txt files.""" def __init__( - self, jsonl_path: Path, image_column: str = IMAGE_COLUMN_DEFAULT, caption_column: str = CAPTION_COLUMN_DEFAULT + self, + jsonl_path: Path | str, + image_column: str = IMAGE_COLUMN_DEFAULT, + caption_column: str = CAPTION_COLUMN_DEFAULT, ): super().__init__() - self._jsonl_path = jsonl_path + self._jsonl_path = Path(jsonl_path) self._image_column = image_column self._caption_column = caption_column