Update tests

Val Brodsky · Val Brodsky · commit d53098b3a6bd · 2023-03-14T12:11:32.000-07:00
Add DocumentEntity to tests

Added integration test
diff --git a/examples/annotation_import/pdf.ipynb b/examples/annotation_import/pdf.ipynb
diff --git a/labelbox/data/annotation_types/__init__.py b/labelbox/data/annotation_types/__init__.py
@@ -11,6 +11,8 @@
 from .annotation import VideoObjectAnnotation
 
 from .ner import TextEntity
+from .ner import DocumentEntity
+from .ner import DocumentTextSelection
 
 from .classification import Checklist
 from .classification import ClassificationAnswer
diff --git a/labelbox/data/annotation_types/annotation.py b/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 TextEntity
+from .ner import DocumentEntity, TextEntity
 
 
 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[TextEntity, Geometry]
+    value: Union[TextEntity, DocumentEntity, Geometry]
     classifications: List[ClassificationAnnotation] = []
 
 
diff --git a/labelbox/data/annotation_types/label.py b/labelbox/data/annotation_types/label.py
@@ -6,6 +6,7 @@
 
 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)
diff --git a/labelbox/data/annotation_types/ner.py b/labelbox/data/annotation_types/ner.py
@@ -33,4 +33,4 @@ def validate_page(cls, v):
 class DocumentEntity(BaseModel):
     """ Represents a text entity """
     name: str
-    text_selections: List[DocumentTextSelection]
+    textSelections: List[DocumentTextSelection]
diff --git a/labelbox/data/serialization/ndjson/converter.py b/labelbox/data/serialization/ndjson/converter.py
@@ -40,6 +40,7 @@ def serialize(
         Returns:
             A generator for accessing the ndjson representation of the data
         """
+
         for example in NDLabel.from_common(labels):
             res = example.dict(by_alias=True)
             for k, v in list(res.items()):
diff --git a/labelbox/data/serialization/ndjson/objects.py b/labelbox/data/serialization/ndjson/objects.py
@@ -376,12 +376,15 @@ def from_common(cls,
 
 class NDDocumentEntity(NDBaseObject, ConfidenceMixin):
     name: str
-    text_selections: List[DocumentTextSelection]
+    textSelections: List[DocumentTextSelection]
 
 
     def to_common(self) -> DocumentEntity:
-        return TextEntity(name=self.name, text_selections=self.text_selections)
+        return TextEntity(name=self.name, textSelections=self.textSelections)
 
+        return obj.from_common(annotation.value, subclasses, annotation.name,
+                               annotation.feature_schema_id, annotation.extra,
+                               data, **optional_kwargs)
     @classmethod
     def from_common(cls,
                     document_entity: DocumentEntity,
@@ -391,7 +394,8 @@ def from_common(cls,
                     extra: Dict[str, Any],
                     data: Union[ImageData, TextData],
                     confidence: Optional[float] = None) -> "NDDocumentEntity":
-        return cls(text_selections=document_entity.text_selections,
+
+        return cls(textSelections=document_entity.textSelections,
                    dataRow=DataRow(id=data.uid),
                    name=name,
                    schema_id=feature_schema_id,
@@ -446,7 +450,8 @@ def from_common(
         optional_kwargs = {}
         if (annotation.confidence):
             optional_kwargs['confidence'] = annotation.confidence
-        return obj.from_common(annotation.value, subclasses, annotation.name,
+        name = annotation.name or annotation.value.name
+        return obj.from_common(annotation.value, subclasses, name,
                                annotation.feature_schema_id, annotation.extra,
                                data, **optional_kwargs)
 
diff --git a/tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483-lb-textlayer.json b/tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483-lb-textlayer.json
diff --git a/tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483.pdf b/tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483.pdf
diff --git a/tests/data/annotation_types/test_ner.py b/tests/data/annotation_types/test_ner.py
@@ -1,4 +1,4 @@
-from labelbox.data.annotation_types import TextEntity
+from labelbox.data.annotation_types import TextEntity, DocumentEntity, DocumentTextSelection
 
 
 def test_ner():
@@ -7,3 +7,11 @@ def test_ner():
     text_entity = TextEntity(start=start, end=end)
     assert text_entity.start == start
     assert text_entity.end == end
+
+def test_document_entity():
+    document_entity = DocumentEntity(name="tool_name", textSelections=[DocumentTextSelection(tokenIds=["1", "2"], groupId="1", page=1)])
+
+    assert document_entity.name == "tool_name"
+    assert document_entity.textSelections[0].tokenIds == ["1", "2"]
+    assert document_entity.textSelections[0].groupId == "1"
+    assert document_entity.textSelections[0].page == 1
diff --git a/tests/integration/annotation_import/conftest.py b/tests/integration/annotation_import/conftest.py
@@ -114,10 +114,11 @@ def ontology():
             'options': []
         }]
     }
+    named_entity = {'tool': 'named-entity', 'name': 'named-entity', 'required': False, 'color': '#A30059', 'classifications': [], }
 
     tools = [
         bbox_tool, polygon_tool, polyline_tool, point_tool, entity_tool,
-        segmentation_tool
+        segmentation_tool, named_entity
     ]
     classifications = [checklist, free_form_text, radio]
     return {"tools": tools, "classifications": classifications}
@@ -152,13 +153,36 @@ def configured_project_pdf(client, ontology, rand_gen, pdf_url):
             where=LabelingFrontend.name == "editor"))[0]
     project.setup(editor, ontology)
     data_row_ids = []
-    data_row_ids.append(dataset.create_data_row(row_data=pdf_url).uid)
+    data_row_ids.append(dataset.create_data_row(pdf_url).uid)
     project.datasets.connect(dataset)
     project.data_row_ids = data_row_ids
     yield project
     project.delete()
     dataset.delete()
 
+@pytest.fixture
+def dataset_pdf_entity(client, rand_gen, pdf_entity_row_data):
+    dataset = client.create_dataset(name=rand_gen(str))
+    data_row_ids = []
+    data_row_ids.append(dataset.create_data_row(pdf_entity_row_data).uid)
+    yield dataset, data_row_ids
+    dataset.delete()
+
+@pytest.fixture
+def configured_project_pdf_entity(client, ontology, rand_gen, dataset_pdf_entity):
+    project = client.create_project(name=rand_gen(str),
+                                    queue_mode=QueueMode.Dataset)
+
+    editor = list(
+        client.get_labeling_frontends(
+            where=LabelingFrontend.name == "editor"))[0]
+    project.setup(editor, ontology)
+
+    dataset, data_row_ids = dataset_pdf_entity
+    project.datasets.connect(dataset)
+    project.data_row_ids = data_row_ids
+    yield project
+    project.delete()
 
 @pytest.fixture
 def configured_project_without_data_rows(client, configured_project, rand_gen):
diff --git a/tests/integration/annotation_import/test_bulk_import_request.py b/tests/integration/annotation_import/test_bulk_import_request.py
@@ -2,6 +2,10 @@
 import ndjson
 import pytest
 import random
+from labelbox.data.annotation_types.annotation import ObjectAnnotation
+from labelbox.data.annotation_types.label import Label
+from labelbox.data.annotation_types.data.text import TextData
+from labelbox.data.annotation_types.ner import DocumentEntity, DocumentTextSelection
 
 from labelbox.data.serialization import NDJsonConverter
 from labelbox.exceptions import MALValidationError, UuidError
@@ -287,3 +291,44 @@ def test_pdf_mal_bbox(client, configured_project_pdf):
     import_annotations.wait_until_done()
 
     assert import_annotations.errors == []
+
+
+def test_pdf_document_entity(client, configured_project_pdf_entity, dataset_pdf_entity):
+    # for content "Metal-insulator (MI) transitions have been one of the" in OCR JSON extract tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483-lb-textlayer.json
+    document_text_selection = DocumentTextSelection(groupId= "2f4336f4-a07e-4e0a-a9e1-5629b03b719b",
+                                                    tokenIds=["3f984bf3-1d61-44f5-b59a-9658a2e3440f",
+                                                              "3bf00b56-ff12-4e52-8cc1-08dbddb3c3b8",
+                                                              "6e1c3420-d4b7-4c5a-8fd6-ead43bf73d80",
+                                                              "87a43d32-af76-4a1d-b262-5c5f4d5ace3a",
+                                                              "e8606e8a-dfd9-4c49-a635-ad5c879c75d0",
+                                                              "67c7c19e-4654-425d-bf17-2adb8cf02c30",
+                                                              "149c5e80-3e07-49a7-ab2d-29ddfe6a38fa",
+                                                              "b0e94071-2187-461e-8e76-96c58738a52c"],
+                                                    page=1)
+
+    entities_annotation_document_entity = DocumentEntity(name="named_entity", 
+                                                         textSelections = [document_text_selection])
+    entities_annotation = ObjectAnnotation(name="named-entity",
+                                           value=entities_annotation_document_entity)
+
+    labels = []
+    _, data_row_uids = dataset_pdf_entity
+    for data_row_uid in data_row_uids:
+        labels.append(Label(
+            data= TextData(
+                uid=data_row_uid),
+            annotations = [
+                entities_annotation,
+            ]
+            )
+        )
+        
+
+    import_annotations = MALPredictionImport.create_from_objects(
+        client=client,
+        project_id=configured_project_pdf_entity.uid,
+        name=f"import {str(uuid.uuid4())}",
+        predictions=labels)
+    import_annotations.wait_until_done()
+
+    assert import_annotations.errors == []
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -175,6 +175,18 @@ def image_url(client):
 def pdf_url(client):
     return client.upload_file('tests/assets/loremipsum.pdf')
 
+@pytest.fixture(scope="session")
+def pdf_entity_row_data(client):
+    pdf_url = client.upload_file('tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483.pdf')
+    text_layer_url = client.upload_file('tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483-lb-textlayer.json')
+
+    return {
+        "row_data": {
+            "pdf_url": pdf_url,
+            "text_layer_url": text_layer_url
+        },
+        "global_key": str(uuid.uuid4())
+    }    
 
 @pytest.fixture
 def project(client, rand_gen):
