Add ConversationEntity annotation class(es)

Author: Val Brodsky

labelbox/data/annotation_types/__init__.py

@@ -10,9 +10,10 @@
 from .annotation import ObjectAnnotation
 from .annotation import VideoObjectAnnotation
 
-from .ner import TextEntity
+from .ner import ConversationEntity
 from .ner import DocumentEntity
 from .ner import DocumentTextSelection
+from .ner import TextEntity
 
 from .classification import Checklist
 from .classification import ClassificationAnswer

labelbox/data/annotation_types/annotation.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, Type
 
 from labelbox.data.mixins import ConfidenceNotSupportedMixin, ConfidenceMixin
 
@@ -51,7 +51,7 @@ class ObjectAnnotation(BaseAnnotation, ConfidenceMixin):
         classifications (Optional[List[ClassificationAnnotation]]): Optional sub classification of the annotation
         extra (Dict[str, Any])
     """
-    value: Union[TextEntity, DocumentEntity, Geometry]
+    value: Union[Type[TextEntity], DocumentEntity, Geometry]
     classifications: List[ClassificationAnnotation] = []
 
 

labelbox/data/annotation_types/label.py

@@ -6,7 +6,6 @@
 
 import labelbox
 from labelbox.data.annotation_types.data.tiled_image import TiledImageData
-from labelbox.data.annotation_types.ner import DocumentEntity
 from labelbox.schema import ontology
 from .annotation import (ClassificationAnnotation, ObjectAnnotation,
                          VideoClassificationAnnotation, VideoObjectAnnotation)

labelbox/data/annotation_types/ner/__init__.py

@@ -1,2 +1,3 @@
-from .text_entity import TextEntity
+from .conversation_entity import ConversationEntity
 from .document_entity import DocumentEntity, DocumentTextSelection
+from .text_entity import TextEntity

labelbox/data/annotation_types/ner/conversation_entity.py

@@ -0,0 +1,11 @@
+from typing import Dict, Any
+
+from pydantic import BaseModel, root_validator
+
+from labelbox.data.annotation_types.ner.text_entity import TextEntity
+from labelbox.utils import _CamelCaseMixin
+
+
+class ConversationEntity(TextEntity, _CamelCaseMixin):
+    """ Represents a text entity """
+    message_id: str

labelbox/data/serialization/ndjson/objects.py

@@ -1,7 +1,8 @@
 from ast import Bytes
 from io import BytesIO
-from typing import Any, Dict, List, Tuple, Union, Optional
+from typing import Any, Dict, List, Tuple, Type, Union, Optional
 import base64
+from labelbox.data.annotation_types.ner.conversation_entity import ConversationEntity
 from labelbox.data.mixins import ConfidenceMixin
 import numpy as np
 
@@ -399,6 +400,33 @@ def from_common(cls,
                    confidence=confidence)
 
 
+class NDConversationEntity(NDTextEntity):
+    message_id: str
+
+    def to_common(self) -> ConversationEntity:
+        return ConversationEntity(start=self.location.start, end=self.location.end, message_id=self.message_id)
+
+    @classmethod
+    def from_common(cls,
+                    conversation_entity: ConversationEntity,
+                    classifications: List[ClassificationAnnotation],
+                    name: str,
+                    feature_schema_id: Cuid,
+                    extra: Dict[str, Any],
+                    data: Union[ImageData, TextData],
+                    confidence: Optional[float] = None) -> "NDConversationEntity":
+        return cls(location=Location(
+                    start=conversation_entity.start,
+                    end=conversation_entity.end),
+                    message_id=conversation_entity.message_id,
+                   dataRow=DataRow(id=data.uid),
+                   name=name,
+                   schema_id=feature_schema_id,
+                   uuid=extra.get('uuid'),
+                   classifications=classifications,
+                   confidence=confidence)
+
+
 class NDObject:
 
     @staticmethod
@@ -463,6 +491,7 @@ def lookup_object(
                 Mask: NDMask,
                 TextEntity: NDTextEntity,
                 DocumentEntity: NDDocumentEntity,
+                ConversationEntity: NDConversationEntity,
             }.get(type(annotation.value))
         if result is None:
             raise TypeError(
@@ -472,6 +501,6 @@ def lookup_object(
 
 
 NDObjectType = Union[NDLine, NDPolygon, NDPoint, NDRectangle, NDMask,
-                     NDTextEntity, NDDocumentEntity]
+                     Type[NDTextEntity], NDDocumentEntity]
 
 NDFrameObjectType = NDFrameRectangle, NDFramePoint, NDFrameLine

tests/data/annotation_types/test_ner.py

@@ -1,4 +1,5 @@
 from labelbox.data.annotation_types import TextEntity, DocumentEntity, DocumentTextSelection
+from labelbox.data.annotation_types.ner.conversation_entity import ConversationEntity
 
 
 def test_ner():
@@ -22,3 +23,17 @@ def test_document_entity():
     assert document_entity.text_selections[0].token_ids == ["1", "2"]
     assert document_entity.text_selections[0].group_id == "1"
     assert document_entity.text_selections[0].page == 1
+
+
+def test_conversation_entity():
+    document_entity = ConversationEntity(name="tool_name",
+                                         message_id=1,
+                                         start=0,
+                                         end=1,
+                                         confidence=0.5)
+
+    assert document_entity.name == "tool_name"
+    assert document_entity.message_id == "1"
+    assert document_entity.start == 0
+    assert document_entity.end == 1
+    assert document_entity.confidence == 0.5