[PTDT-2372] Added support for MMC tasks annotations

Author: real-mjozwiak

libs/labelbox/src/labelbox/data/annotation_types/__init__.py

@@ -63,4 +63,6 @@
 from .data.tiled_image import TileLayer
 
 from .llm_prompt_response.prompt import PromptText
-from .llm_prompt_response.prompt import PromptClassificationAnnotation
+from .llm_prompt_response.prompt import PromptClassificationAnnotation
+
+from .mmc import MessageInfo, OrderedMessageInfo, MessageSingleSelection, MessageMultiSelection, MessageRanking

libs/labelbox/src/labelbox/data/annotation_types/annotation.py

@@ -7,6 +7,7 @@
 
 from labelbox.data.annotation_types.classification.classification import ClassificationAnnotation
 from .ner import DocumentEntity, TextEntity, ConversationEntity
+from .mmc import BaseMMCAnnotation
 
 
 class ObjectAnnotation(BaseAnnotation, ConfidenceMixin, CustomMetricsMixin):
@@ -28,5 +29,5 @@ class ObjectAnnotation(BaseAnnotation, ConfidenceMixin, CustomMetricsMixin):
         extra (Dict[str, Any])
     """
 
-    value: Union[TextEntity, ConversationEntity, DocumentEntity, Geometry]
+    value: Union[TextEntity, ConversationEntity, DocumentEntity, Geometry, BaseMMCAnnotation]
     classifications: List[ClassificationAnnotation] = []

libs/labelbox/src/labelbox/data/annotation_types/mmc.py

@@ -0,0 +1,39 @@
+from abc import ABC
+from typing import ClassVar, List
+
+from labelbox import pydantic_compat
+from labelbox.utils import _CamelCaseMixin
+
+
+class MessageInfo(_CamelCaseMixin):
+    message_id: str
+    model_config_name: str
+
+
+class OrderedMessageInfo(MessageInfo):
+    order: int
+
+
+class BaseMMCAnnotation(_CamelCaseMixin, ABC):
+    format: ClassVar[str]
+    parent_message_id: str
+
+
+class MessageSingleSelection(BaseMMCAnnotation, MessageInfo):
+    format: ClassVar[str] = "message-single-selection"
+
+
+class MessageMultiSelection(BaseMMCAnnotation):
+    format: ClassVar[str] = "message-multi-selection"
+    selected_messages: List[MessageInfo]
+
+
+class MessageRanking(BaseMMCAnnotation):
+    format: ClassVar[str] = "message-ranking"
+    ranked_messages: List[OrderedMessageInfo]
+
+    @pydantic_compat.validator("ranked_messages")
+    def _validate_ranked_messages(cls, v: List[OrderedMessageInfo]):
+        if not {msg.order for msg in v} == set(range(1, len(v) + 1)):
+            raise ValueError("Messages must be ordered by unique and consecutive natural numbers starting from 1")
+        return v

libs/labelbox/src/labelbox/data/serialization/ndjson/converter.py

@@ -72,7 +72,7 @@ def serialize(
                 ConfusionMatrixMetric,
                 RelationshipAnnotation,
             ]] = []
-            # First pass to get all RelatiohnshipAnnotaitons
+            # First pass to get all RelationshipAnnotaitons
             # and update the UUIDs of the source and target annotations
             for annotation in label.annotations:
                 if isinstance(annotation, RelationshipAnnotation):

libs/labelbox/src/labelbox/data/serialization/ndjson/objects.py

@@ -5,6 +5,7 @@
 from labelbox.data.annotation_types.ner.conversation_entity import ConversationEntity
 from labelbox.data.annotation_types.video import VideoObjectAnnotation, DICOMObjectAnnotation
 from labelbox.data.mixins import ConfidenceMixin, CustomMetricsMixin, CustomMetric, CustomMetricsNotSupportedMixin
+from labelbox.utils import _CamelCaseMixin
 import numpy as np
 
 from labelbox import pydantic_compat
@@ -19,6 +20,7 @@
 from ...annotation_types.geometry import DocumentRectangle, Rectangle, Polygon, Line, Point, Mask
 from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation
 from ...annotation_types.video import VideoMaskAnnotation, DICOMMaskAnnotation, MaskFrame, MaskInstance
+from ...annotation_types.mmc import BaseMMCAnnotation
 from .classification import NDClassification, NDSubclassification, NDSubclassificationType
 from .base import DataRow, NDAnnotation, NDJsonBase
 
@@ -666,8 +668,53 @@ def from_common(
                    custom_metrics=custom_metrics)
 
 
+class NDMessageTask(NDAnnotation):
+
+    class MessageTaskData(_CamelCaseMixin):
+        format: str
+        data: BaseMMCAnnotation
+
+    message_evaluation_task: MessageTaskData
+
+    @classmethod
+    def from_common(
+        cls,
+        uuid: str,
+        annotation_value: BaseMMCAnnotation,
+        classifications: List[ClassificationAnnotation],
+        name: str,
+        feature_schema_id: Cuid,
+        extra: Dict[str, Any],
+        data: Union[ImageData, TextData],
+        confidence: Optional[float] = None,
+        custom_metrics: Optional[List[CustomMetric]] = None
+    ) -> "NDMessageTask":
+        return cls(
+            uuid=uuid,
+            name=name,
+            dataRow=DataRow(id=data.uid, global_key=data.global_key),
+            message_evaluation_task=cls.MessageTaskData(
+                format=annotation_value.format,
+                data=annotation_value
+            )
+        )
+
+
 class NDObject:
 
+    _ANNOTATION_TO_SERIALIZER = {
+        Line: NDLine,
+        Point: NDPoint,
+        Polygon: NDPolygon,
+        Rectangle: NDRectangle,
+        DocumentRectangle: NDDocumentRectangle,
+        Mask: NDMask,
+        TextEntity: NDTextEntity,
+        DocumentEntity: NDDocumentEntity,
+        ConversationEntity: NDConversationEntity,
+        BaseMMCAnnotation: NDMessageTask
+    }
+
     @staticmethod
     def to_common(annotation: "NDObjectType") -> ObjectAnnotation:
         common_annotation = annotation.to_common()
@@ -755,17 +802,11 @@ def lookup_object(
             else:
                 result = NDSegments
         else:
-            result = {
-                Line: NDLine,
-                Point: NDPoint,
-                Polygon: NDPolygon,
-                Rectangle: NDRectangle,
-                DocumentRectangle: NDDocumentRectangle,
-                Mask: NDMask,
-                TextEntity: NDTextEntity,
-                DocumentEntity: NDDocumentEntity,
-                ConversationEntity: NDConversationEntity,
-            }.get(type(annotation.value))
+            result = next((
+                serializer_class
+                for annotation_class, serializer_class in NDObject._ANNOTATION_TO_SERIALIZER.items()
+                if isinstance(annotation.value, annotation_class)
+            ), None)
         if result is None:
             raise TypeError(
                 f"Unable to convert object to MAL format. `{type(annotation.value)}`"