Add DocumentEntity class

Author: Val Brodsky

labelbox/data/annotation_types/ner.py

@@ -1,6 +1,6 @@
-from typing import Dict, Any
+from typing import Dict, Any, List
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, root_validator, validator
 
 
 class TextEntity(BaseModel):
@@ -17,3 +17,20 @@ 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
+    text_selections: List[DocumentTextSelection]