test: fix all the tests

pyproject.toml

@@ -43,6 +43,7 @@ gunicorn = ">=20.1.0"
 fastapi = { extras = ["all"], version = ">=0.88.0" }
 httpx = ">=0.23.1"
 huggingface-hub = ">=0.19.4"
+numpy = ">1.22"
 
 [tool.poetry.group.dev.dependencies]
 black = ">=22.3.0"

src/emma_common/datamodels/features.py

@@ -1,9 +1,8 @@
-from __future__ import annotations
-
-from typing import Any
+from typing import Any, Optional
 
 import torch
 from pydantic import BaseModel
+from typing_extensions import Self
 
 
 class EmmaExtractedFeatures(BaseModel):
@@ -14,7 +13,7 @@ class EmmaExtractedFeatures(BaseModel):
     bbox_probas: torch.Tensor
     cnn_features: torch.Tensor
     class_labels: list[str]
-    entity_labels: list[str] | None = None
+    entity_labels: Optional[list[str]] = None
     width: int
     height: int
 
@@ -24,7 +23,7 @@ class Config:
         arbitrary_types_allowed: bool = True
 
     @classmethod
-    def from_raw_response(cls, raw_response: dict[str, Any]) -> EmmaExtractedFeatures:
+    def from_raw_response(cls, raw_response: dict[str, Any]) -> Self:
         """Instantiate the features from a raw json response.
 
         Dict keys that are used below can be found within the emma-simbot/perception repo.

src/emma_common/datamodels/state.py

@@ -1,6 +1,5 @@
-from __future__ import annotations
-
 from enum import Enum
+from typing import Optional
 
 from pydantic import BaseModel
 
@@ -25,7 +24,7 @@ class EnvironmentStateTurn(BaseModel):
     """State of the environment at a single timestep."""
 
     features: list[EmmaExtractedFeatures]
-    output: str | None = None
+    output: Optional[str] = None
 
 
 class EmmaPolicyRequest(BaseModel):

src/emma_common/datamodels/torch.py

@@ -8,9 +8,7 @@ class TorchDataMixin:
     """A utility class containing useful functions for torch serialisation."""
 
     @staticmethod
-    def to_bytes(  # noqa: WPS602
-        raw_object: Any,
-    ) -> bytes:
+    def to_bytes(raw_object: Any) -> bytes:  # noqa: WPS602
         """Factory method that builds a byte representation from a raw object."""
         buffer = io.BytesIO()
         torch.save(raw_object, buffer)