Refactor ner classes each in a separate file

Author: Val Brodsky

labelbox/data/annotation_types/ner/__init__.py

@@ -0,0 +1,2 @@
+from .text_entity import TextEntity
+from .document_entity import DocumentEntity, DocumentTextSelection

labelbox/data/annotation_types/ner/document_entity.py

@@ -0,0 +1,20 @@
+from typing import List
+
+from pydantic import BaseModel, validator
+
+
+class DocumentTextSelection(BaseModel):
+    tokenIds: List[str]
+    groupId: str
+    page: int
+
+    @validator("page")
+    def validate_page(cls, v):
+        if v < 1:
+            raise ValueError("Page must be greater than 1")
+        return v
+
+class DocumentEntity(BaseModel):
+    """ Represents a text entity """
+    name: str
+    textSelections: List[DocumentTextSelection]

labelbox/data/annotation_types/ner.py renamed to labelbox/data/annotation_types/ner/text_entity.py

@@ -1,6 +1,6 @@
-from typing import Dict, Any, List
+from typing import Dict, Any
 
-from pydantic import BaseModel, root_validator, validator
+from pydantic import BaseModel, root_validator
 
 
 class TextEntity(BaseModel):
@@ -17,20 +17,3 @@ def validate_start_end(cls, values):
                 raise ValueError(
                     "Location end must be greater or equal to start")
         return values
-
-
-class DocumentTextSelection(BaseModel):
-    tokenIds: List[str]
-    groupId: str
-    page: int
-
-    @validator("page")
-    def validate_page(cls, v):
-        if v < 1:
-            raise ValueError("Page must be greater than 1")
-        return v
-
-class DocumentEntity(BaseModel):
-    """ Represents a text entity """
-    name: str
-    textSelections: List[DocumentTextSelection]

labelbox/data/serialization/ndjson/objects.py

@@ -450,8 +450,7 @@ def from_common(
         optional_kwargs = {}
         if (annotation.confidence):
             optional_kwargs['confidence'] = annotation.confidence
-        name = annotation.name or annotation.value.name
-        return obj.from_common(annotation.value, subclasses, name,
+        return obj.from_common(annotation.value, subclasses, annotation.name,
                                annotation.feature_schema_id, annotation.extra,
                                data, **optional_kwargs)
 

tests/integration/annotation_import/conftest.py

@@ -164,7 +164,8 @@ def configured_project_pdf(client, ontology, rand_gen, pdf_url):
 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)
+    data_row = dataset.create_data_row(pdf_entity_row_data)
+    data_row_ids.append(data_row.uid)
     yield dataset, data_row_ids
     dataset.delete()
 

tests/integration/conftest.py

@@ -188,6 +188,7 @@ def pdf_entity_row_data(client):
         "global_key": str(uuid.uuid4())
     }    
 
+
 @pytest.fixture
 def project(client, rand_gen):
     project = client.create_project(name=rand_gen(str),