Pydantic model for captions

aana/configs/pipeline.py

@@ -3,6 +3,7 @@
 It is used to generate the pipeline and the API endpoints.
 """
 
+from aana.models.pydantic.captions import CaptionsList, VideoCaptionsList
 from aana.models.pydantic.image_input import ImageInputList
 from aana.models.pydantic.prompt import Prompt
 from aana.models.pydantic.sampling_params import SamplingParams
@@ -185,6 +186,7 @@
                 "name": "captions_hf_blip2_opt_2_7b",
                 "key": "captions",
                 "path": "image_batch.images.[*].caption_hf_blip2_opt_2_7b",
+                "data_model": CaptionsList,
             }
         ],
     },
@@ -260,6 +262,7 @@
                 "name": "video_captions_hf_blip2_opt_2_7b",
                 "key": "captions",
                 "path": "video_batch.videos.[*].frames.[*].caption_hf_blip2_opt_2_7b",
+                "data_model": VideoCaptionsList,
             }
         ],
     },

aana/models/pydantic/captions.py

@@ -0,0 +1,40 @@
+from typing import List
+from pydantic import BaseModel
+
+from aana.models.pydantic.base import BaseListModel
+
+
+class Caption(BaseModel):
+    """A model for a caption."""
+
+    __root__: str
+
+    def __str__(self):
+        return self.__root__
+
+    class Config:
+        schema_extra = {"description": "A caption."}
+
+
+class CaptionsList(BaseListModel):
+    """A model for a list of captions."""
+
+    __root__: List[Caption]
+
+    class Config:
+        schema_extra = {"description": "A list of captions."}
+
+
+class VideoCaptionsList(BaseListModel):
+    """A model for a list of captions for a list of videos."""
+
+    __root__: List[CaptionsList]
+
+    class Config:
+        schema_extra = {
+            "description": (
+                "A list of a list of captions. "
+                "For a list of videos and a list of captions for each video "
+                "and for each video we have a list of captions"
+            )
+        }