From 562728f5618090e4da47f88f158c91f0aaed6067 Mon Sep 17 00:00:00 2001
From: movchan74 <aleksandr@mobius.ml>
Date: Tue, 7 Nov 2023 15:31:11 +0000
Subject: [PATCH] Pydantic model for captions

---
 aana/configs/pipeline.py         |  3 +++
 aana/models/pydantic/captions.py | 40 ++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)
 create mode 100644 aana/models/pydantic/captions.py

diff --git a/aana/configs/pipeline.py b/aana/configs/pipeline.py
index ff0eb821..966479ad 100644
--- a/aana/configs/pipeline.py
+++ b/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,
             }
         ],
     },
diff --git a/aana/models/pydantic/captions.py b/aana/models/pydantic/captions.py
new file mode 100644
index 00000000..4de79851
--- /dev/null
+++ b/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"
+            )
+        }