Dealing with pydantic deserialization of subclasses

Author: Val Brodsky

labelbox/data/annotation_types/annotation.py

@@ -6,7 +6,7 @@
 from .classification import Checklist, Dropdown, Radio, Text
 from .feature import FeatureSchema
 from .geometry import Geometry, Rectangle, Point
-from .ner import DocumentEntity, TextEntity
+from .ner import DocumentEntity, TextEntity, ConversationEntity
 
 
 class BaseAnnotation(FeatureSchema, abc.ABC):
@@ -51,7 +51,7 @@ class ObjectAnnotation(BaseAnnotation, ConfidenceMixin):
         classifications (Optional[List[ClassificationAnnotation]]): Optional sub classification of the annotation
         extra (Dict[str, Any])
     """
-    value: Union[Type[TextEntity], DocumentEntity, Geometry]
+    value: Union[TextEntity, ConversationEntity, DocumentEntity, Geometry]
     classifications: List[ClassificationAnnotation] = []
 
 

labelbox/data/serialization/ndjson/label.py

@@ -10,7 +10,7 @@
 from ...annotation_types.collection import LabelCollection, LabelGenerator
 from ...annotation_types.data import ImageData, TextData, VideoData
 from ...annotation_types.label import Label
-from ...annotation_types.ner import TextEntity
+from ...annotation_types.ner import TextEntity, ConversationEntity
 from ...annotation_types.classification import Dropdown
 from ...annotation_types.metrics import ScalarMetric, ConfusionMatrixMetric
 
@@ -52,7 +52,6 @@ def _generate_annotations(
             annots = []
             data_row = annotations[0].data_row
             for annotation in annotations:
-
                 if isinstance(annotation, NDSegments):
                     annots.extend(
                         NDSegments.to_common(annotation, annotation.name,
@@ -72,7 +71,7 @@ def _generate_annotations(
 
     def _infer_media_type(
         self, data_row: DataRow,
-        annotations: List[Union[TextEntity, VideoClassificationAnnotation,
+        annotations: List[Union[TextEntity, ConversationEntity, VideoClassificationAnnotation,
                                 VideoObjectAnnotation, ObjectAnnotation,
                                 ClassificationAnnotation, ScalarMetric,
                                 ConfusionMatrixMetric]]
@@ -82,7 +81,7 @@ def _infer_media_type(
 
         types = {type(annotation) for annotation in annotations}
         data = ImageData
-        if TextEntity in types:
+        if TextEntity in types or ConversationEntity:
             data = TextData
         elif VideoClassificationAnnotation in types or VideoObjectAnnotation in types:
             data = VideoData

labelbox/data/serialization/ndjson/objects.py

@@ -500,7 +500,8 @@ def lookup_object(
         return result
 
 
+NERTextType = Union[NDConversationEntity, NDTextEntity]
 NDObjectType = Union[NDLine, NDPolygon, NDPoint, NDRectangle, NDMask,
-                     Type[NDTextEntity], NDDocumentEntity]
+                     NERTextType, NDDocumentEntity]
 
 NDFrameObjectType = NDFrameRectangle, NDFramePoint, NDFrameLine

tests/data/serialization/ndjson/test_text_entity.py

@@ -12,6 +12,7 @@
 def test_text_entity_import(filename: str):
     with open(filename, 'r') as file:
         data = json.load(file)
+
     res = list(NDJsonConverter.deserialize(data))
     res = list(NDJsonConverter.serialize(res))
     assert res == data