gimlet-api 0.0.8__py3-none-any.whl → 0.0.10__py3-none-any.whl

gml/hf.py

@@ -24,8 +24,10 @@ from typing import Any, BinaryIO, Dict, List, Optional, TextIO, Tuple
 
 import torch
 import transformers
+from rich.progress import Console
 from transformers import (
     BaseImageProcessor,
+    DynamicCache,
     Pipeline,
     PreTrainedModel,
     PreTrainedTokenizer,
@@ -49,6 +51,7 @@ from gml.tensor import (
     DetectionNumCandidatesDimension,
     DetectionOutputDimension,
     DimensionSemantics,
+    EmbeddingDimension,
     ImageChannelDimension,
     ImageHeightDimension,
     ImageWidthDimension,
@@ -60,6 +63,11 @@ from gml.tensor import (
 
 FALLBACK_RESIZE_SIZE = 512
 
+# Set dynamic dimension max size to less than the int64 max, leaving leeway for the size to be ~4x by the model.
+MAX_DYNAMIC_VAL = 2**61
+
+console = Console()
+
 
 class HuggingFaceTokenizer(Model):
     def __init__(self, tokenizer: PreTrainedTokenizer, name: Optional[str] = None):
@@ -105,7 +113,6 @@ def flatten(items):
 
 
 class WrapWithFunctionalCache(torch.nn.Module):
-
     def __init__(self, model: transformers.PreTrainedModel):
         super().__init__()
         self.model = model
@@ -128,6 +135,8 @@ class HuggingFaceTextGenerationPipeline:
         name: Optional[str] = None,
         tokenizer_name: Optional[str] = None,
         dynamic_seqlen: bool = False,
+        dynamic_batch: bool = False,
+        export_predispatch: bool = False,
     ):
         self.pipeline = pipeline
         self.tokenizer_model = HuggingFaceTokenizer(pipeline.tokenizer, tokenizer_name)
@@ -139,13 +148,20 @@ class HuggingFaceTextGenerationPipeline:
         self.model = self.model.to(torch.float16)
         self.model = WrapWithFunctionalCache(pipeline.model)
 
+        self.dynamic_batch = dynamic_batch
+        self.batch_size = 1
+        if self.dynamic_batch:
+            # dynamic tracing fails for dimensions of size 1.
+            self.batch_size = 2
+
         self.language_model = TorchModel(
             name,
             torch_module=self.model,
+            export_predispatch=export_predispatch,
             **self._guess_model_spec(dynamic_seqlen),
         )
 
-    def _initialize_key_value_cache(self):
+    def _initialize_key_value_cache(self) -> DynamicCache:
         cache = []
         config = self.pipeline.model.config
         head_dim = (
@@ -158,7 +174,12 @@ class HuggingFaceTextGenerationPipeline:
             if config.num_key_value_heads is None
             else config.num_key_value_heads
         )
-        cache_shape = (1, num_key_value_heads, self._cache_length_for_tracing, head_dim)
+        cache_shape = (
+            self.batch_size,
+            num_key_value_heads,
+            self._cache_length_for_tracing,
+            head_dim,
+        )
         for _ in range(config.num_hidden_layers):
             cache.append(
                 [
@@ -166,7 +187,67 @@ class HuggingFaceTextGenerationPipeline:
                     torch.zeros(cache_shape).to(torch.float16),
                 ]
             )
-        return cache
+        return DynamicCache.from_legacy_cache(cache)
+
+    def _parse_transformer_config(
+        self, model: transformers.PreTrainedModel
+    ) -> modelexecpb.TransformerConfig:
+        # Only non-default rope config set the rope_scaling parameter
+        attention_head_size = getattr(
+            model.config,
+            "attention_head_size",
+            model.config.hidden_size // model.config.num_attention_heads,
+        )
+        partial_rotary_factor = getattr(model.config, "partial_rotary_factor", 1.0)
+        rotary_embedding_dim = getattr(
+            model.config,
+            "rotary_dim",
+            int(attention_head_size * partial_rotary_factor),
+        )
+        if (
+            hasattr(model.config, "rope_scaling")
+            and model.config.rope_scaling is not None
+        ):
+            rope_scaling = model.config.rope_scaling
+            rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))
+            if not rope_type == "llama3":
+                raise NotImplementedError(
+                    "rope scaling type {} is not supported".format(rope_type)
+                )
+            # LLAMA 3 example config: https://huggingface.co/meta-llama/Llama-3.2-1B/blob/main/config.json
+            llama3_config = modelexecpb.Llama3RopeConfig()
+            llama3_config.theta = model.config.rope_theta
+            llama3_config.rotary_embedding_dim = rotary_embedding_dim
+            llama3_config.max_position_embeddings = model.config.max_position_embeddings
+
+            llama3_config.factor = rope_scaling["factor"]
+            llama3_config.high_freq_factor = rope_scaling["high_freq_factor"]
+            llama3_config.low_freq_factor = rope_scaling["low_freq_factor"]
+            llama3_config.original_max_position_embeddings = rope_scaling[
+                "original_max_position_embeddings"
+            ]
+            return modelexecpb.TransformerConfig(
+                position_embedding_config=modelexecpb.PositionEmbeddingConfig(
+                    kind=modelexecpb.PositionEmbeddingKind.POSITION_EMBEDDING_KIND_ROPE_LLAMA3,
+                    llama3_rope_config=llama3_config,
+                ),
+            )
+        # Default rope configs:
+        # 1. Llama-2: https://huggingface.co/NousResearch/Llama-2-7b-hf/blob/main/config.json
+        # 2. Qwen2.5: https://huggingface.co/Qwen/Qwen2.5-14B-Instruct-1M/blob/main/config.json
+        # 3. Mixtral: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1/blob/main/config.json
+        default_rope_config = modelexecpb.DefaultRopeConfig()
+        default_rope_config.theta = model.config.rope_theta
+        default_rope_config.max_position_embeddings = (
+            model.config.max_position_embeddings
+        )
+        default_rope_config.rotary_embedding_dim = rotary_embedding_dim
+        return modelexecpb.TransformerConfig(
+            position_embedding_config=modelexecpb.PositionEmbeddingConfig(
+                kind=modelexecpb.PositionEmbeddingKind.POSITION_EMBEDDING_KIND_ROPE_DEFAULT,
+                default_rope_config=default_rope_config,
+            ),
+        )
 
     def _guess_model_spec(self, dynamic_seqlen: bool) -> Dict:
         input_dict = self.pipeline.preprocess("this is a prompt! Test test test?")
@@ -179,7 +260,7 @@ class HuggingFaceTextGenerationPipeline:
         input_tensor_semantics = []
 
         # This currently assumes that all HF language models have inputs that are [B, NUM_TOKENS].
-        inputs.append(input_dict["input_ids"])
+        inputs.append(torch.tile(input_dict["input_ids"], [self.batch_size, 1]))
         input_tensor_semantics.append(
             TensorSemantics(
                 dimensions=[
@@ -192,7 +273,7 @@ class HuggingFaceTextGenerationPipeline:
         # Assume that the model supports a KeyValue cache.
         cache_values = self._initialize_key_value_cache()
         inputs.append(cache_values)
-        for _ in cache_values:
+        for _ in range(len(cache_values)):
             input_tensor_semantics.append(AttentionKeyValueCacheTensorSemantics())
             input_tensor_semantics.append(AttentionKeyValueCacheTensorSemantics())
 
@@ -209,7 +290,7 @@ class HuggingFaceTextGenerationPipeline:
             if (
                 not found_logits
                 and len(tensor.shape) == 3
-                and tensor.shape[0] == 1
+                and tensor.shape[0] == self.batch_size
                 and tensor.shape[1] == seqlen
             ):
                 # This should be the logits tensor.
@@ -226,14 +307,38 @@ class HuggingFaceTextGenerationPipeline:
             else:
                 output_tensor_semantics.append(AttentionKeyValueCacheTensorSemantics())
 
+        if not found_logits:
+            raise ValueError(
+                "could not determine output logits tensor for text generation model"
+            )
+
+        num_experts_per_tok = (
+            1
+            if not hasattr(self.pipeline.model.config, "num_experts_per_tok")
+            else self.pipeline.model.config.num_experts_per_tok
+        )
+
         dynamic_shapes = None
-        seqlen = torch.export.Dim("seqlen", min=2, max=9223372036854775096)
+        # Set range to half of seqlen to account for # of tokens per expert.
+        # pytorch export creates a constraint on the number of possible tokens
+        # sent to each expert. That value is num_experts * seqlen. If we don't divide
+        # by number of experts, the tracing creates an integer value that exceeds the valid int64
+        # range and will throw a hard to decipher error message.
+        seqlen = torch.export.Dim(
+            "seqlen", min=2, max=MAX_DYNAMIC_VAL // num_experts_per_tok
+        )
 
-        cache_length = torch.export.Dim("cache_length", min=2, max=9223372036854775096)
+        cache_length = torch.export.Dim("cache_length", min=2, max=MAX_DYNAMIC_VAL)
         dynamic_shapes = [
             {1: seqlen},
-            [[{2: cache_length}, {2: cache_length}] for _ in cache_values],
+            [[{2: cache_length}, {2: cache_length}] for _ in range(len(cache_values))],
         ]
+        if self.dynamic_batch:
+            batch = torch.export.Dim("batch")
+            dynamic_shapes[0][0] = batch
+            for i in range(len(cache_values)):
+                dynamic_shapes[1][i][0][0] = batch
+                dynamic_shapes[1][i][1][0] = batch
 
         return {
             "example_inputs": inputs,
@@ -241,6 +346,7 @@ class HuggingFaceTextGenerationPipeline:
             "input_tensor_semantics": input_tensor_semantics,
             "output_tensor_semantics": output_tensor_semantics,
             "generation_config": HuggingFaceGenerationConfig(self.pipeline.model),
+            "transformer_config": self._parse_transformer_config(self.pipeline.model),
         }
 
     def models(self) -> List[Model]:
@@ -334,6 +440,14 @@ class HuggingFaceImageProcessor:
             raise NotImplementedError(
                 "only semantic segmentation is currently supported"
             )
+        # TODO(philkuz): Support panoptic segmentation models. Multiple outputs come from panoptic segmentation models.
+        # We need to decide whether we should invest in converting the panoptic segmentation output to semantic segmentation
+        # format or if we should directly support panoptic segmentation output.
+        if hasattr(self.processor, "post_process_panoptic_segmentation"):
+            raise NotImplementedError(
+                "panoptic segmentation models are not supported yet"
+            )
+
         dimensions = [
             BatchDimension(),
             # TODO(james): verify all semantic segmentation in hugging face output a logits mask.
@@ -687,8 +801,8 @@ class HuggingFaceZeroShotObjectDetectionPipeline:
 
         spec["dynamic_shapes"].extend(
             [
-                {0: "num_labels"},
-                {0: "num_labels"},
+                {0: torch.export.Dim("num_labels", max=MAX_DYNAMIC_VAL)},
+                {0: torch.export.Dim("num_labels", max=MAX_DYNAMIC_VAL)},
             ]
         )
 
@@ -754,33 +868,121 @@ class HuggingFaceDepthEstimationPipeline:
         return [self.model]
 
 
-def import_huggingface_pipeline(pipeline: Pipeline, **kwargs) -> List[Model]:
-    if pipeline.framework != "pt":
-        raise ValueError(
-            "unimplemented: hugging face pipeline framework: {}".format(
-                pipeline.framework
+class HuggingFaceFeatureExtractionPipeline:
+    def __init__(self, pipeline: Pipeline, name: Optional[str] = None):
+        self.pipeline = pipeline
+        if name is None:
+            name = pipeline.model.name_or_path
+
+        self.tokenizer_model = HuggingFaceTokenizer(self.pipeline.tokenizer)
+
+        self.model = TorchModel(
+            name=name,
+            torch_module=self.pipeline.model,
+            **self._guess_model_spec(),
+        )
+
+    def _guess_model_spec(self) -> Dict:
+        spec = {
+            "example_inputs": [],
+            "input_tensor_semantics": [],
+            "output_tensor_semantics": [],
+            "dynamic_shapes": [],
+        }
+
+        input_dict = self.pipeline.preprocess("this is a prompt! Test test test?")
+        if "input_ids" not in input_dict:
+            raise ValueError(
+                'HuggingFaceFeatureExtractionPipeline expects preprocessed inputs to have an "input_ids" tensor'
             )
+
+        spec["example_inputs"].append(input_dict["input_ids"])
+        spec["input_tensor_semantics"].extend(
+            [
+                TensorSemantics(
+                    dimensions=[
+                        BatchDimension(),
+                        TokensDimension(),
+                    ]
+                ),
+            ]
         )
 
-    if pipeline.task == "text-generation":
-        return HuggingFaceTextGenerationPipeline(pipeline, **kwargs).models()
-    elif pipeline.task == "image-segmentation":
-        return HuggingFaceImageSegmentationPipeline(pipeline, **kwargs).models()
-    elif pipeline.task == "object-detection":
-        return HuggingFaceObjectDetectionPipeline(pipeline, **kwargs).models()
-    elif pipeline.task == "zero-shot-object-detection":
-        return HuggingFaceZeroShotObjectDetectionPipeline(pipeline, **kwargs).models()
-    elif pipeline.task == "depth-estimation":
-        return HuggingFaceDepthEstimationPipeline(pipeline, **kwargs).models()
-    raise ValueError(
-        "unimplemented: hugging face pipeline task: {} (supported tasks: [{}])".format(
-            pipeline.task,
+        spec["output_tensor_semantics"].extend(
             [
-                "text-generation",
-                "image-segmentation",
-                "object-detection",
-                "zero-shot-object-detection",
-                "depth-estimation",
-            ],
-        )
-    )
+                TensorSemantics(
+                    dimensions=[
+                        BatchDimension(),
+                        TokensDimension(),
+                        EmbeddingDimension(),
+                    ],
+                ),
+                TensorSemantics(
+                    dimensions=[
+                        BatchDimension(),
+                        EmbeddingDimension(),
+                    ],
+                ),
+            ]
+        )
+
+        max_seqlen = (
+            getattr(self.pipeline.model.config, "max_position_embeddings", 500) - 1
+        )
+        spec["dynamic_shapes"].extend(
+            [
+                {
+                    1: torch.export.Dim(
+                        "seqlen",
+                        max=max_seqlen,
+                    )
+                },
+            ]
+        )
+        return spec
+
+    def models(self) -> List[Model]:
+        return [self.model, self.tokenizer_model]
+
+
+def import_huggingface_pipeline(pipeline: Pipeline, **kwargs) -> List[Model]:
+    with console.status(
+        f'Importing HuggingFace pipeline: "{pipeline.model.name_or_path}"'
+    ):
+        if pipeline.framework != "pt":
+            raise ValueError(
+                "unimplemented: hugging face pipeline framework: {}".format(
+                    pipeline.framework
+                )
+            )
+
+        if pipeline.task == "text-generation":
+            result = HuggingFaceTextGenerationPipeline(pipeline, **kwargs).models()
+        elif pipeline.task == "image-segmentation":
+            result = HuggingFaceImageSegmentationPipeline(pipeline, **kwargs).models()
+        elif pipeline.task == "object-detection":
+            result = HuggingFaceObjectDetectionPipeline(pipeline, **kwargs).models()
+        elif pipeline.task == "zero-shot-object-detection":
+            result = HuggingFaceZeroShotObjectDetectionPipeline(
+                pipeline, **kwargs
+            ).models()
+        elif pipeline.task == "depth-estimation":
+            result = HuggingFaceDepthEstimationPipeline(pipeline, **kwargs).models()
+        elif pipeline.task == "feature-extraction":
+            result = HuggingFaceFeatureExtractionPipeline(pipeline, **kwargs).models()
+        else:
+            raise ValueError(
+                "unimplemented: hugging face pipeline task: {} (supported tasks: [{}])".format(
+                    pipeline.task,
+                    [
+                        "text-generation",
+                        "image-segmentation",
+                        "object-detection",
+                        "zero-shot-object-detection",
+                        "depth-estimation",
+                        "feature-extraction",
+                    ],
+                )
+            )
+    console.print(f'Imported HuggingFace pipeline: "{pipeline.model.name_or_path}".')
+    return result

gml/model.py

@@ -18,6 +18,7 @@ from __future__ import annotations
 import abc
 import contextlib
 import io
+from importlib.metadata import distributions
 from pathlib import Path
 from typing import BinaryIO, Dict, List, Literal, Optional, Sequence, TextIO, Tuple
 
@@ -29,6 +30,33 @@ from gml.compile import to_torch_mlir
 from gml.preprocessing import ImagePreprocessingStep
 from gml.tensor import TensorSemantics
 
+MODELING_PACKAGES = set(
+    {
+        "accelerate",
+        "compressed-tensors",
+        "gimlet-api",
+        "mlir-gml",
+        "safetensors",
+        "safetensors-mlir",
+        "torch",
+        "torch-mlir-gml",
+        "torchvision",
+        "transformers",
+        "tokenizers",
+    }
+)
+
+
+def get_modeling_packages() -> Sequence[modelexecpb.PackageInfo]:
+    packages = []
+    for dist in filter(
+        lambda d: d.metadata["Name"] in MODELING_PACKAGES, distributions()
+    ):
+        packages.append(
+            modelexecpb.PackageInfo(name=dist.metadata["Name"], version=dist.version)
+        )
+    return sorted(packages, key=lambda p: p.name)
+
 
 class GenerationConfig:
     def __init__(self, eos_token_ids: List[int]):
@@ -52,6 +80,7 @@ class Model(abc.ABC):
         class_labels_file: Optional[Path] = None,
         image_preprocessing_steps: Optional[List[ImagePreprocessingStep]] = None,
         generation_config: Optional[GenerationConfig] = None,
+        transformer_config: Optional[modelexecpb.TransformerConfig] = None,
     ):
         self.name = name
         self.kind = kind
@@ -66,6 +95,7 @@ class Model(abc.ABC):
         self.output_tensor_semantics = output_tensor_semantics
         self.image_preprocessing_steps = image_preprocessing_steps
         self.generation_config = generation_config
+        self.transformer_config = transformer_config
 
     def to_proto(self) -> modelexecpb.ModelInfo:
         image_preprocessing_steps = None
@@ -89,6 +119,10 @@ class Model(abc.ABC):
                 semantics.to_proto() for semantics in self.output_tensor_semantics
             ],
             generation_config=generation_config,
+            tracing_metadata=modelexecpb.TracingMetadata(
+                package_info=get_modeling_packages(),
+            ),
+            transformer_config=self.transformer_config,
         )
 
     @abc.abstractmethod
@@ -111,6 +145,7 @@ class TorchModel(Model):
         input_shapes: Optional[List[List[int]]] = None,
         input_dtypes: Optional[List[torch.dtype]] = None,
         dynamic_shapes: Optional[Sequence[Dict[int, str | "torch.export.Dim"]]] = None,
+        export_predispatch: bool = False,
         **kwargs,
     ):
         super().__init__(
@@ -133,6 +168,7 @@ class TorchModel(Model):
                 torch.rand(shape, dtype=dtype)
                 for shape, dtype in zip(self.input_shapes, self.input_dtypes)
             ]
+        self.export_predispatch = export_predispatch
 
     def _convert_to_torch_mlir(self, weight_manager: Optional[AssetManager] = None):
         return to_torch_mlir(
@@ -140,6 +176,7 @@ class TorchModel(Model):
             self.example_inputs,
             self.dynamic_shapes,
             weight_manager=weight_manager,
+            export_predispatch=self.export_predispatch,
         )
 
     def _collect_assets(

gml/pipelines.py

@@ -173,6 +173,16 @@ nodes:
 
 
 class LiveChatPipeline(Pipeline):
+    def __init__(
+        self,
+        message_template: str = r'''"{% for message in messages %}{% if message['role'] == 'system' %}{{message['content']}}{% else %}{{'<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n' }}{% endif %}"''',
+        preset_system_prompt: str = r"'<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful assistant<|eot_id|>'",
+        add_generation_prompt: bool = True,
+    ):
+        self.add_generation_prompt = add_generation_prompt
+        self.message_template = message_template
+        self.preset_system_prompt = preset_system_prompt
+
     def to_yaml(self, models: List[Model], org_name: str) -> str:
         if len(models) != 2:
             raise ValueError(
@@ -195,6 +205,16 @@ nodes:
   kind: TextStreamSource
   outputs:
   - prompt
+- name: query_template
+  kind: TemplateChatMessage
+  attributes:
+    add_generation_prompt: {self.add_generation_prompt}
+    message_template: {self.message_template}
+    preset_system_prompt: {self.preset_system_prompt}
+  inputs:
+    query: .text_source.prompt
+  outputs:
+  - chat_message
 - name: tokenize
   kind: Tokenize
   attributes:
@@ -203,7 +223,7 @@ nodes:
         name: {tokenizer.name}
         org: {org_name}
   inputs:
-    text: .text_source.prompt
+    text: .query_template.chat_message
   outputs:
   - tokens
 - name: generate

gml/preprocessing.py

@@ -28,7 +28,21 @@ class ImagePreprocessingStep(abc.ABC):
         pass
 
 
+class ResizeImage(ImagePreprocessingStep):
+    """ResizeImage resizes the image to the target size without preserving aspect ratio."""
+
+    def to_proto(self) -> modelexecpb.ImagePreprocessingStep:
+        return modelexecpb.ImagePreprocessingStep(
+            kind=modelexecpb.ImagePreprocessingStep.IMAGE_PREPROCESSING_KIND_RESIZE,
+            resize_params=modelexecpb.ImagePreprocessingStep.ImageResizeParams(
+                kind=modelexecpb.ImagePreprocessingStep.ImageResizeParams.IMAGE_RESIZE_KIND_STRETCH,
+            ),
+        )
+
+
 class LetterboxImage(ImagePreprocessingStep):
+    """LetterboxImage resizes the image to the target size while preserving aspect ratio by introducting letterbox padding."""
+
     def to_proto(self) -> modelexecpb.ImagePreprocessingStep:
         return modelexecpb.ImagePreprocessingStep(
             kind=modelexecpb.ImagePreprocessingStep.IMAGE_PREPROCESSING_KIND_RESIZE,
@@ -38,14 +52,14 @@ class LetterboxImage(ImagePreprocessingStep):
         )
 
 
-class ResizeImage(ImagePreprocessingStep):
-    """ResizeImage resizes the image to the target size without preserving aspect ratio."""
+class CenterCropImage(ImagePreprocessingStep):
+    """CenterCropImage resizes the image to the target size while preserving aspect ratio by center cropping along one dimension."""
 
     def to_proto(self) -> modelexecpb.ImagePreprocessingStep:
         return modelexecpb.ImagePreprocessingStep(
             kind=modelexecpb.ImagePreprocessingStep.IMAGE_PREPROCESSING_KIND_RESIZE,
             resize_params=modelexecpb.ImagePreprocessingStep.ImageResizeParams(
-                kind=modelexecpb.ImagePreprocessingStep.ImageResizeParams.IMAGE_RESIZE_KIND_STRETCH,
+                kind=modelexecpb.ImagePreprocessingStep.ImageResizeParams.IMAGE_RESIZE_KIND_CENTERCROP,
             ),
         )