gimlet-api 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

gml/tensor.py

@@ -19,6 +19,7 @@ from typing import List, Literal, Optional, Tuple
 
 import gml.proto.src.api.corepb.v1.model_exec_pb2 as modelexecpb
 import google.protobuf.wrappers_pb2 as wrapperspb
+import numpy as np
 
 
 def box_format_str_to_proto(box_format: str):
@@ -122,17 +123,21 @@ class DetectionNumCandidatesDimension(DimensionSemantics):
 class DetectionOutputDimension(DimensionSemantics):
     def __init__(
         self,
-        coordinates_start_index: int,
-        box_format: BoundingBoxFormat,
+        coordinates_start_index: Optional[int] = None,
+        box_format: Optional[BoundingBoxFormat] = None,
         box_confidence_index: Optional[int] = None,
         class_index: Optional[int] = None,
         scores_range: Optional[Tuple[int, int]] = None,
+        scores_are_logits: bool = False,
     ):
-        self.coordinates_range = (coordinates_start_index, 4)
+        self.coordinates_range = None
+        if coordinates_start_index is not None:
+            self.coordinates_range = (coordinates_start_index, 4)
         self.box_format = box_format
         self.box_confidence_index = box_confidence_index
         self.class_index = class_index
         self.scores_range = scores_range
+        self.scores_are_logits = scores_are_logits
 
     def to_proto(self) -> modelexecpb.DimensionSemantics:
         scores_range = None
@@ -143,20 +148,29 @@ class DetectionOutputDimension(DimensionSemantics):
                     size=self.scores_range[1],
                 )
             )
-        box_confidence_index = -1
+        box_confidence_index = np.iinfo(np.int32).min
         if self.box_confidence_index is not None:
             box_confidence_index = self.box_confidence_index
+        box_format_proto = None
+        if self.box_format is not None:
+            box_format_proto = self.box_format.to_proto()
+        box_coordinate_range_proto = None
+        if self.coordinates_range is not None:
+            box_coordinate_range_proto = (
+                modelexecpb.DimensionSemantics.DetectionOutputParams.IndexRange(
+                    start=self.coordinates_range[0],
+                    size=self.coordinates_range[1],
+                )
+            )
         return modelexecpb.DimensionSemantics(
             kind=modelexecpb.DimensionSemantics.DIMENSION_SEMANTICS_KIND_DETECTION_OUTPUT,
             detection_output_params=modelexecpb.DimensionSemantics.DetectionOutputParams(
-                box_coordinate_range=modelexecpb.DimensionSemantics.DetectionOutputParams.IndexRange(
-                    start=self.coordinates_range[0],
-                    size=self.coordinates_range[1],
-                ),
-                box_format=self.box_format.to_proto(),
+                box_coordinate_range=box_coordinate_range_proto,
+                box_format=box_format_proto,
                 box_confidence_index=box_confidence_index,
                 class_index=self.class_index,
                 scores_range=scores_range,
+                scores_are_logits=self.scores_are_logits,
             ),
         )
 
@@ -171,16 +185,23 @@ def _segmentation_mask_kind_to_proto(kind: str):
             return (
                 modelexecpb.DimensionSemantics.SegmentationMaskParams.SEGMENTATION_MASK_KIND_CLASS_LABEL
             )
-        case "score_masks":
+        case "score_mask":
             return (
                 modelexecpb.DimensionSemantics.SegmentationMaskParams.SEGMENTATION_MASK_KIND_SCORE
             )
+        case "logits_mask":
+            return (
+                modelexecpb.DimensionSemantics.SegmentationMaskParams.SEGMENTATION_MASK_KIND_LOGITS
+            )
         case _:
             raise ValueError("Invalid segmentation mask kind: {}".format(kind))
 
 
 class SegmentationMaskChannel(DimensionSemantics):
-    def __init__(self, kind: Literal["bool_masks", "int_label_masks", "score_masks"]):
+    def __init__(
+        self,
+        kind: Literal["bool_masks", "int_label_masks", "score_mask", "logits_mask"],
+    ):
         self.kind = _segmentation_mask_kind_to_proto(kind)
 
     def to_proto(self) -> modelexecpb.DimensionSemantics:
@@ -224,13 +245,63 @@ class RegressionValueDimension(DimensionSemantics):
         )
 
 
+class TokensDimension(DimensionSemantics):
+    def to_proto(self) -> modelexecpb.DimensionSemantics:
+        return modelexecpb.DimensionSemantics(
+            kind=modelexecpb.DimensionSemantics.DIMENSION_SEMANTICS_KIND_TOKENS,
+        )
+
+
+class AttentionMaskDimension(DimensionSemantics):
+    def to_proto(self) -> modelexecpb.DimensionSemantics:
+        return modelexecpb.DimensionSemantics(
+            kind=modelexecpb.DimensionSemantics.DIMENSION_SEMANTICS_KIND_ATTENTION_MASK,
+        )
+
+
+class VocabLogitsDimension(DimensionSemantics):
+    def to_proto(self) -> modelexecpb.DimensionSemantics:
+        return modelexecpb.DimensionSemantics(
+            kind=modelexecpb.DimensionSemantics.DIMENSION_SEMANTICS_KIND_VOCAB_LOGITS,
+        )
+
+
+class EmbeddingDimension(DimensionSemantics):
+    def to_proto(self) -> modelexecpb.DimensionSemantics:
+        return modelexecpb.DimensionSemantics(
+            kind=modelexecpb.DimensionSemantics.DIMENSION_SEMANTICS_KIND_EMBEDDING,
+        )
+
+
 class TensorSemantics:
-    def __init__(self, dimensions: List[DimensionSemantics]):
+    def __init__(
+        self,
+        dimensions: List[DimensionSemantics],
+        kind: modelexecpb.TensorSemantics.TensorSemanticsKind = modelexecpb.TensorSemantics.TENSOR_SEMANTICS_KIND_DIMENSION,
+    ):
         self.dimensions = dimensions
+        self.kind = kind
 
     def to_proto(self) -> modelexecpb.TensorSemantics:
         return modelexecpb.TensorSemantics(
             dimensions=[dim.to_proto() for dim in self.dimensions],
+            kind=self.kind,
+        )
+
+
+class UnusedTensorSemantics(TensorSemantics):
+    def __init__(self):
+        super().__init__(
+            dimensions=[],
+            kind=modelexecpb.TensorSemantics.TENSOR_SEMANTICS_KIND_UNUSED,
+        )
+
+
+class AttentionKeyValueCacheTensorSemantics(TensorSemantics):
+    def __init__(self):
+        super().__init__(
+            dimensions=[],
+            kind=modelexecpb.TensorSemantics.TENSOR_SEMANTICS_KIND_ATTENTION_KEY_VALUE_CACHE,
         )
 
 
@@ -297,37 +368,47 @@ class BinarySegmentationMasks(TensorSemantics):
 class YOLOOutput(TensorSemantics):
     """YOLOOutput represents a detection output from a YOLO model.
 
-    The YOLO model should output a tensor of shape [B, N_BOXES, (4 or 5) + NUM_CLASSES].
+    If `has_box_conf=True` then the YOLO model should output a tensor of shape [B, NUM_BOXES, 5 + NUM_CLASSES].
+    Otherwise, it should output a tensor of shape [B, NUM_BOXES, 4 + NUM_CLASSES].
     """
 
-    def __init__(self, has_box_conf=True):
+    def __init__(self, version="v5"):
+        if version != "v5" and version != "v8":
+            raise ValueError(
+                "gml.tensor.YOLOOutput alias currently only supports YOLO versions v5 and v8"
+            )
         dimensions = [
             BatchDimension(),
-            DetectionNumCandidatesDimension(is_nms=False),
         ]
 
-        if has_box_conf:
-            dimensions.append(
-                DetectionOutputDimension(
-                    coordinates_start_index=0,
-                    box_format=BoundingBoxFormat(
-                        box_format="cxcywh",
-                        is_normalized=False,
+        if version == "v5":
+            dimensions.extend(
+                [
+                    DetectionNumCandidatesDimension(is_nms=False),
+                    DetectionOutputDimension(
+                        coordinates_start_index=0,
+                        box_format=BoundingBoxFormat(
+                            box_format="cxcywh",
+                            is_normalized=False,
+                        ),
+                        box_confidence_index=4,
+                        scores_range=(5, -1),
                     ),
-                    box_confidence_index=4,
-                    scores_range=(5, -1),
-                )
+                ]
             )
-        else:
-            dimensions.append(
-                DetectionOutputDimension(
-                    coordinates_start_index=0,
-                    box_format=BoundingBoxFormat(
-                        box_format="cxcywh",
-                        is_normalized=False,
+        elif version == "v8":
+            dimensions.extend(
+                [
+                    DetectionOutputDimension(
+                        coordinates_start_index=0,
+                        box_format=BoundingBoxFormat(
+                            box_format="cxcywh",
+                            is_normalized=False,
+                        ),
+                        scores_range=(4, -1),
                     ),
-                    scores_range=(4, -1),
-                )
+                    DetectionNumCandidatesDimension(is_nms=False),
+                ]
             )
 
         super().__init__(dimensions)