Updated pydantic models to support lists

base/ocr.py

@@ -2,7 +2,7 @@
 Custom types for dealing with the Google Vision API JSON output.
 """
 from enum import IntEnum
-from typing import Any, Optional
+from typing import Any, Optional, List
 
 from pydantic import BaseModel
 
@@ -36,7 +36,7 @@ class DetectedLanguage(BaseModel):
 
 
 class TextProperty(BaseModel):
-    detectedLanguages: list[DetectedLanguage]
+    detectedLanguages: List[DetectedLanguage]
     detectedBreak: Optional[DetectedBreak]
 
 
@@ -50,21 +50,21 @@ class Symbol(BaseModel):
 class Word(BaseModel):
     property: Optional[TextProperty]
     boundingBox: Any
-    symbols: list[Symbol]
+    symbols: List[Symbol]
     confidence: float
 
 
 class Paragraph(BaseModel):
     property: Optional[TextProperty]
     boundingBox: Any
-    words: list[Word]
+    words: List[Word]
     confidence: float
 
 
 class Block(BaseModel):
     property: Optional[TextProperty]
     boundingBox: Any
-    paragraphs: list[Paragraph]
+    paragraphs: List[Paragraph]
     blockType: BlockType
     confidence: float
 
@@ -73,12 +73,12 @@ class Page(BaseModel):
     property: Optional[TextProperty]
     width: int
     height: int
-    blocks: list[Block]
+    blocks: List[Block]
     confidence: float
 
 
 class TextAnnotation(BaseModel):
-    pages: list[Page]
+    pages: List[Page]
     text: str
 
 

data_models.py

@@ -4,7 +4,7 @@ Pydantic models used throughout the codebase.
 In particular, these are the types that are used as input and output of each step of the pipeline.
 """
 import json
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, List
 
 from pydantic import BaseModel, Field
 
@@ -65,7 +65,7 @@ class NutrientTable(BaseModel):
     nutrientBasisQuantityValue: Optional[str]
     nutrientBasisQuantityMeasurementUnitCode: Optional[str]
     preperationStateCode: Optional[str]
-    values: list[NutrientTableElement]
+    values: List[NutrientTableElement]
 
     def __str__(self):
         top = "Nutrients per " + " ".join(
@@ -83,7 +83,7 @@ class Attribute(BaseModel):
     coordinates: str
     entity: str
     probability: float
-    value: Union[str, list[NutrientTable]]
+    value: Union[str, List[NutrientTable]]
     model: str
 
 
@@ -116,13 +116,13 @@ class NetContentAttribute(BaseModel):
 
 class AllergensOut(BaseModel):
     entity: str
-    values: list[AttributeAllergen]
+    values: List[AttributeAllergen]
     model: str
 
 
 class CommunicationChannelsOut(BaseModel):
     entity: str
-    values: list[AttributeCommunicationChannel]
+    values: List[AttributeCommunicationChannel]
     model: str
 
 
@@ -131,7 +131,7 @@ class PipelineInput(BaseModel):
 
 
 class PipelineOutput(BaseModel):
-    attributes: list[
+    attributes: List[
         Union[Attribute, CommunicationChannelsOut, AllergensOut, NetContentAttribute]
     ]
     job_id: str = Field(alias="job-id")
@@ -147,13 +147,13 @@ class TextWithLanguage(BaseModel):
 
 
 class OCRTextOut(BaseModel):
-    blocks: list[str]
+    blocks: List[str]
     full_text: str
-    sentences: list[TextWithLanguage]
+    sentences: List[TextWithLanguage]
 
 
 class OCRTableOut(BaseModel):
-    tables: list[list[list[str]]]
+    tables: List[List[List[str]]]
 
 
 class OCROut(BaseModel):
@@ -162,7 +162,7 @@ class OCROut(BaseModel):
 
 
 class OCROutList(BaseModel):
-    __root__: list[OCROut]
+    __root__: List[OCROut]
 
     def __iter__(self):
         return iter(self.__root__)
@@ -172,11 +172,11 @@ class OCROutList(BaseModel):
 
 
 class OCRWrapperOut(BaseModel):
-    blocks: list[str]
+    blocks: List[str]
     full_text: str
     job_id: str
-    sentences: list[TextWithLanguage]
-    tables: list[list[list[str]]]
+    sentences: List[TextWithLanguage]
+    tables: List[List[List[str]]]
 
 
 class ClassifiedText(BaseModel):
@@ -210,8 +210,8 @@ class NetContent(BaseModel):
 
 
 class ModelOut(BaseModel):
-    blocks: list[Union[NetContent, Allergen, CommunicationChannels, ClassifiedText]]
-    tables: Optional[list[NutrientTable]]
+    blocks: List[Union[NetContent, Allergen, CommunicationChannels, ClassifiedText]]
+    tables: Optional[List[NutrientTable]]
     job_id: str
     model: str
     full_text: str
@@ -221,7 +221,7 @@ class ModelOut(BaseModel):
 
 
 class ModelOutList(BaseModel):
-    __root__: list[ModelOut]
+    __root__: List[ModelOut]
 
     def __iter__(self):
         return iter(self.__root__)
@@ -231,6 +231,6 @@ class ModelOutList(BaseModel):
 
 
 class TrainModelOut(BaseModel):
-    # To be defined later when we have a list of accepted formats
+    # To be defined later when we have a List of accepted formats
     model: Optional[Any] = None
     artifacts: Optional[Any] = None