Initial upload: aesthetic-immich-starter

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MediaPipe GenAI Python API."""

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MediaPipe Tasks GenAI Bundler API."""
+
+import mediapipe.tasks.python.genai.bundler.llm_bundler
+
+BundleConfig = llm_bundler.BundleConfig
+create_bundle = llm_bundler.create_bundle
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del llm_bundler

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/llm_bundler.py

@@ -0,0 +1,236 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Functions to perform llm packing."""
+
+import ctypes
+import dataclasses
+import enum
+from typing import List, Optional
+
+from mediapipe.tasks.python.core import mediapipe_c_bindings
+from mediapipe.tasks.python.core import mediapipe_c_utils
+from mediapipe.tasks.python.core import serial_dispatcher
+from mediapipe.tasks.python.genai.bundler import llm_bundler_metadata_options
+from mediapipe.tasks.python.genai.bundler import llm_bundler_metadata_options_c
+from mediapipe.tasks.python.genai.converter import external_dependencies
+from mediapipe.tasks.python.metadata.metadata_writers import model_asset_bundle_utils
+
+sentencepiece = external_dependencies.sentencepiece
+
+_CTYPES_SIGNATURES = (
+    mediapipe_c_utils.CFunction(
+        func_name="MpLlmBundlerGenerateMetadata",
+        argtypes=[
+            ctypes.POINTER(
+                llm_bundler_metadata_options_c.LlmBundlerMetadataOptionsC
+            ),
+            ctypes.POINTER(ctypes.c_int),
+        ],
+        restype=ctypes.POINTER(ctypes.c_char),
+    ),
+    mediapipe_c_utils.CFunction(
+        func_name="MpLlmBundlerFreeMetadata",
+        argtypes=[
+            ctypes.POINTER(ctypes.c_char),
+        ],
+        restype=None,
+    ),
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class BundleConfig:
+  """Config for LLM Bundler.
+
+  Attributes:
+    tflite_model: Path to the multi-signature tflite model with "prefill" and
+      "decode" signatures converted using ODML Transformers APIs.
+    tokenizer_model: Path to the tokenizer model. Currently only SentencePience
+      tokenizer is supported. As such, tokenizer.model proto is expected to be
+      passed here.
+    start_token: Token that will be used to signify the beginning of a sequence.
+    stop_tokens: Tokens that will be used to signify the end of a sequence.
+    output_filename: Name of the generated `.task` file containing the Bundle.
+    enable_bytes_to_unicode_mapping: Enables GPT-2 style bytes to unicode
+      mapping. For more details see:
+      https://github.com/openai/gpt-2/blob/master/src/encoder.py#L9
+    system_prompt: The starting text would be feed to the model on the start of
+      each session, commonly called the System prompt. This is useful for
+      instruction tuned models and pre-conditioning the model behavior. This
+      field is mutually exclusive with prompt_prefix_system and
+      prompt_suffix_system.
+    prompt_prefix_user: The prefix that should be prepended to each user portion
+      of prompts passed to the model.
+    prompt_suffix_user: The suffix that should be appended to each user portion
+      of prompts passed to the model.
+    prompt_prefix_model: The prefix that should be prepended to each model
+      portion of prompts passed to the model.
+    prompt_suffix_model: The suffix that should be appended to each model
+      portion of prompts passsed to the model.
+    prompt_prefix_system: The prefix that should be prepended to each system
+      portion of prompts passed to the model.
+    prompt_suffix_system: The suffix that should be appended to each system
+      portion of prompts passsed to the model.
+    user_role_token: The token that indicates the user's turn.
+    system_role_token: The token that indicates the system's turn.
+    model_role_token: The token that indicates the model's turn.
+    end_role_token: The token that indicates the end of the prompt.
+    tflite_embedder: Path to the embedding model if the embedding lookup is done
+      externally.
+    tflite_per_layer_embedder: Path to the per layer embedding model if the per
+      layer embedding lookup is done externally.
+    tflite_vision_encoder: Path to the vision encoder model.
+    tflite_vision_adapter: Path to the vision adapter model.
+  """
+
+  tflite_model: str
+  tokenizer_model: str
+  start_token: str
+  stop_tokens: List[str]
+  output_filename: str
+  enable_bytes_to_unicode_mapping: bool = False
+  system_prompt: Optional[str] = None
+  prompt_prefix_user: Optional[str] = None
+  prompt_suffix_user: Optional[str] = None
+  prompt_prefix_model: Optional[str] = None
+  prompt_suffix_model: Optional[str] = None
+  prompt_prefix_system: Optional[str] = None
+  prompt_suffix_system: Optional[str] = None
+  user_role_token: Optional[str] = None
+  system_role_token: Optional[str] = None
+  model_role_token: Optional[str] = None
+  end_role_token: Optional[str] = None
+  tflite_embedder: Optional[str] = None
+  tflite_per_layer_embedder: Optional[str] = None
+  tflite_vision_encoder: Optional[str] = None
+  tflite_vision_adapter: Optional[str] = None
+
+
+class _BundleTags(enum.Enum):
+  """Bundle tags."""
+
+  TF_LITE_PREFILL_DECODE = 1
+  TOKENIZER_MODEL = 2
+  METADATA = 3
+  TF_LITE_EMBEDDER = 4
+  TF_LITE_PER_LAYER_EMBEDDER = 5
+  TF_LITE_VISION_ENCODER = 6
+  TF_LITE_VISION_ADAPTER = 7
+
+
+def _validate_config(config: BundleConfig):
+  """Validates the given config.
+
+  Args:
+    config: The config to validate.
+
+  Raises:
+    ValueError if the config is invalid. Invalid configs can be:
+    - tokenizer_model is not a valid SentencePiece model.
+    - start_token is not a valid token in the tokenizer_model.
+    - stop_tokens is not a list of valid tokens in the tokenizer_model.
+    - system_prompt and prompt_*_system are both set.
+  """
+  if not isinstance(config.stop_tokens, list):
+    raise ValueError("stop_tokens must be a list of strings.")
+  if not config.stop_tokens:
+    raise ValueError("stop_tokens must be non-empty.")
+
+  try:
+    sp = sentencepiece.SentencePieceProcessor()
+    sp.Load(config.tokenizer_model)
+  except Exception as e:
+    raise ValueError(
+        f"Failed to load tokenizer model from {config.tokenizer_model}. "
+        "Please ensure you are passing a valid SentencePiece model."
+    ) from e
+
+  if config.system_prompt and (
+      config.prompt_prefix_system or config.prompt_suffix_system
+  ):
+    raise ValueError(
+        "system_prompt and prompt_*_system are mutually exclusive."
+    )
+
+
+def create_bundle(config: BundleConfig):
+  """Creates a bundle from the given config."""
+  _validate_config(config)
+
+  artifacts = {}
+  with open(config.tflite_model, "rb") as f:
+    artifacts[_BundleTags.TF_LITE_PREFILL_DECODE.name] = f.read()
+
+  with open(config.tokenizer_model, "rb") as f:
+    artifacts[_BundleTags.TOKENIZER_MODEL.name] = f.read()
+
+  lib = mediapipe_c_bindings.load_raw_library(_CTYPES_SIGNATURES)
+  with serial_dispatcher.SerialDispatcher(
+      lib, _CTYPES_SIGNATURES
+  ) as bundler_lib:
+    metadata_options = llm_bundler_metadata_options.LlmBundlerMetadataOptions(
+        start_token=config.start_token,
+        stop_tokens=config.stop_tokens,
+        enable_bytes_to_unicode_mapping=config.enable_bytes_to_unicode_mapping,
+        system_prompt=config.system_prompt,
+        prompt_prefix_user=config.prompt_prefix_user,
+        prompt_suffix_user=config.prompt_suffix_user,
+        prompt_prefix_model=config.prompt_prefix_model,
+        prompt_suffix_model=config.prompt_suffix_model,
+        prompt_prefix_system=config.prompt_prefix_system,
+        prompt_suffix_system=config.prompt_suffix_system,
+        user_role_token=config.user_role_token,
+        system_role_token=config.system_role_token,
+        model_role_token=config.model_role_token,
+        end_role_token=config.end_role_token,
+    )
+    metadata_options_c = metadata_options.to_ctypes()
+    buffer_size = ctypes.c_int()
+    metadata_buffer = None
+    try:
+      metadata_buffer = bundler_lib.MpLlmBundlerGenerateMetadata(
+          ctypes.byref(metadata_options_c), ctypes.byref(buffer_size)
+      )
+      artifacts[_BundleTags.METADATA.name] = ctypes.string_at(
+          metadata_buffer, buffer_size.value
+      )
+    finally:
+      if metadata_buffer:
+        bundler_lib.MpLlmBundlerFreeMetadata(metadata_buffer)
+
+  if config.tflite_embedder:
+    with open(config.tflite_embedder, "rb") as f:
+      artifacts[_BundleTags.TF_LITE_EMBEDDER.name] = f.read()
+
+  if config.tflite_per_layer_embedder:
+    with open(config.tflite_per_layer_embedder, "rb") as f:
+      artifacts[_BundleTags.TF_LITE_PER_LAYER_EMBEDDER.name] = f.read()
+
+  if config.tflite_vision_encoder:
+    with open(config.tflite_vision_encoder, "rb") as f:
+      artifacts[_BundleTags.TF_LITE_VISION_ENCODER.name] = f.read()
+
+  if config.tflite_vision_adapter:
+    with open(config.tflite_vision_adapter, "rb") as f:
+      artifacts[_BundleTags.TF_LITE_VISION_ADAPTER.name] = f.read()
+
+  output_filename = config.output_filename
+  if not output_filename.endswith(".task"):
+    output_filename = config.output_filename + ".task"
+
+  model_asset_bundle_utils.create_model_asset_bundle(
+      artifacts,
+      output_filename,
+  )

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/llm_bundler_metadata_options.py

@@ -0,0 +1,142 @@
+# Copyright 2026 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""LlmBundlerMetadataOptions Python dataclass."""
+
+import ctypes
+import dataclasses
+from typing import Optional
+
+from mediapipe.tasks.python.genai.bundler import llm_bundler_metadata_options_c
+
+
+@dataclasses.dataclass
+class LlmBundlerMetadataOptions:
+  """Options for bundling LLM metadata.
+
+  Attributes:
+    start_token: Token that will be used to signify the beginning of a sequence.
+    stop_tokens: Tokens that will be used to signify the end of a sequence.
+    enable_bytes_to_unicode_mapping: Enables GPT-2 style bytes to unicode
+      mapping. For more details see:
+      https://github.com/openai/gpt-2/blob/master/src/encoder.py#L9
+    system_prompt: The starting text would be feed to the model on the start of
+      each session, commonly called the System prompt. This is useful for
+      instruction tuned models and pre-conditioning the model behavior. This
+      field is mutually exclusive with prompt_prefix_system and
+      prompt_suffix_system.
+    prompt_prefix_user: The prefix that should be prepended to each user portion
+      of prompts passed to the model.
+    prompt_suffix_user: The suffix that should be appended to each user portion
+      of prompts passed to the model.
+    prompt_prefix_model: The prefix that should be prepended to each model
+      portion of prompts passed to the model.
+    prompt_suffix_model: The suffix that should be appended to each model
+      portion of prompts passsed to the model.
+    prompt_prefix_system: The prefix that should be prepended to each system
+      portion of prompts passed to the model.
+    prompt_suffix_system: The suffix that should be appended to each system
+      portion of prompts passsed to the model.
+    user_role_token: The token that indicates the user's turn.
+    system_role_token: The token that indicates the system's turn.
+    model_role_token: The token that indicates the model's turn.
+    end_role_token: The token that indicates the end of the prompt.
+  """
+
+  start_token: str
+  stop_tokens: list[str]
+  enable_bytes_to_unicode_mapping: bool = False
+  system_prompt: Optional[str] = None
+  prompt_prefix_user: Optional[str] = None
+  prompt_suffix_user: Optional[str] = None
+  prompt_prefix_model: Optional[str] = None
+  prompt_suffix_model: Optional[str] = None
+  prompt_prefix_system: Optional[str] = None
+  prompt_suffix_system: Optional[str] = None
+  user_role_token: Optional[str] = None
+  system_role_token: Optional[str] = None
+  model_role_token: Optional[str] = None
+  end_role_token: Optional[str] = None
+
+  def to_ctypes(
+      self,
+  ) -> llm_bundler_metadata_options_c.LlmBundlerMetadataOptionsC:
+    """Converts the object to a ctypes structure."""
+    stop_tokens_c = (ctypes.c_char_p * len(self.stop_tokens))()
+    for i, token in enumerate(self.stop_tokens):
+      stop_tokens_c[i] = token.encode("utf-8")
+    start_token_c = self.start_token.encode("utf-8")
+    system_prompt_c = (
+        self.system_prompt.encode("utf-8") if self.system_prompt else None
+    )
+    prompt_prefix_user_c = (
+        self.prompt_prefix_user.encode("utf-8")
+        if self.prompt_prefix_user
+        else None
+    )
+    prompt_suffix_user_c = (
+        self.prompt_suffix_user.encode("utf-8")
+        if self.prompt_suffix_user
+        else None
+    )
+    prompt_prefix_model_c = (
+        self.prompt_prefix_model.encode("utf-8")
+        if self.prompt_prefix_model
+        else None
+    )
+    prompt_suffix_model_c = (
+        self.prompt_suffix_model.encode("utf-8")
+        if self.prompt_suffix_model
+        else None
+    )
+    prompt_prefix_system_c = (
+        self.prompt_prefix_system.encode("utf-8")
+        if self.prompt_prefix_system
+        else None
+    )
+    prompt_suffix_system_c = (
+        self.prompt_suffix_system.encode("utf-8")
+        if self.prompt_suffix_system
+        else None
+    )
+    user_role_token_c = (
+        self.user_role_token.encode("utf-8") if self.user_role_token else None
+    )
+    system_role_token_c = (
+        self.system_role_token.encode("utf-8")
+        if self.system_role_token
+        else None
+    )
+    model_role_token_c = (
+        self.model_role_token.encode("utf-8") if self.model_role_token else None
+    )
+    end_role_token_c = (
+        self.end_role_token.encode("utf-8") if self.end_role_token else None
+    )
+    return llm_bundler_metadata_options_c.LlmBundlerMetadataOptionsC(
+        start_token=start_token_c,
+        stop_tokens=stop_tokens_c,
+        num_stop_tokens=len(self.stop_tokens),
+        enable_bytes_to_unicode_mapping=self.enable_bytes_to_unicode_mapping,
+        system_prompt=system_prompt_c,
+        prompt_prefix_user=prompt_prefix_user_c,
+        prompt_suffix_user=prompt_suffix_user_c,
+        prompt_prefix_model=prompt_prefix_model_c,
+        prompt_suffix_model=prompt_suffix_model_c,
+        prompt_prefix_system=prompt_prefix_system_c,
+        prompt_suffix_system=prompt_suffix_system_c,
+        user_role_token=user_role_token_c,
+        system_role_token=system_role_token_c,
+        model_role_token=model_role_token_c,
+        end_role_token=end_role_token_c,
+    )

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/llm_bundler_metadata_options_c.py

@@ -0,0 +1,38 @@
+# Copyright 2026 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Ctypes wrapper for LlmBundlerMetadataOptions."""
+
+import ctypes
+
+
+class LlmBundlerMetadataOptionsC(ctypes.Structure):
+  """Wrapper for LlmBundlerMetadataOptions in C."""
+
+  _fields_ = [
+      ("start_token", ctypes.c_char_p),
+      ("stop_tokens", ctypes.POINTER(ctypes.c_char_p)),
+      ("num_stop_tokens", ctypes.c_int),
+      ("enable_bytes_to_unicode_mapping", ctypes.c_bool),
+      ("system_prompt", ctypes.c_char_p),
+      ("prompt_prefix_user", ctypes.c_char_p),
+      ("prompt_suffix_user", ctypes.c_char_p),
+      ("prompt_prefix_model", ctypes.c_char_p),
+      ("prompt_suffix_model", ctypes.c_char_p),
+      ("prompt_prefix_system", ctypes.c_char_p),
+      ("prompt_suffix_system", ctypes.c_char_p),
+      ("user_role_token", ctypes.c_char_p),
+      ("system_role_token", ctypes.c_char_p),
+      ("model_role_token", ctypes.c_char_p),
+      ("end_role_token", ctypes.c_char_p),
+  ]

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/llm_bundler_metadata_options_test.py

@@ -0,0 +1,60 @@
+# Copyright 2026 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from absl.testing import absltest
+from mediapipe.tasks.python.genai.bundler import llm_bundler_metadata_options
+
+
+class LlmBundlerMetadataOptionsTest(absltest.TestCase):
+
+  def test_to_ctypes(self):
+    options = llm_bundler_metadata_options.LlmBundlerMetadataOptions(
+        start_token="start",
+        stop_tokens=["stop1", "stop2"],
+        enable_bytes_to_unicode_mapping=True,
+        system_prompt="system",
+        prompt_prefix_user="user_prefix",
+        prompt_suffix_user="user_suffix",
+        prompt_prefix_model="model_prefix",
+        prompt_suffix_model="model_suffix",
+        prompt_prefix_system=None,
+        prompt_suffix_system=None,
+        user_role_token="user_role",
+        system_role_token=None,
+        model_role_token="model_role",
+        end_role_token="end_role",
+    )
+
+    options_c = options.to_ctypes()
+
+    self.assertEqual(options_c.start_token, b"start")
+    self.assertEqual(options_c.stop_tokens[0], b"stop1")
+    self.assertEqual(options_c.stop_tokens[1], b"stop2")
+    self.assertEqual(options_c.num_stop_tokens, 2)
+    self.assertTrue(options_c.enable_bytes_to_unicode_mapping)
+    self.assertEqual(options_c.system_prompt, b"system")
+    self.assertEqual(options_c.prompt_prefix_user, b"user_prefix")
+    self.assertEqual(options_c.prompt_suffix_user, b"user_suffix")
+    self.assertEqual(options_c.prompt_prefix_model, b"model_prefix")
+    self.assertEqual(options_c.prompt_suffix_model, b"model_suffix")
+    self.assertIsNone(options_c.prompt_prefix_system)
+    self.assertIsNone(options_c.prompt_suffix_system)
+    self.assertEqual(options_c.user_role_token, b"user_role")
+    self.assertIsNone(options_c.system_role_token)
+    self.assertEqual(options_c.model_role_token, b"model_role")
+    self.assertEqual(options_c.end_role_token, b"end_role")
+
+
+if __name__ == "__main__":
+  absltest.main()

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/__init__.py

@@ -0,0 +1,24 @@
+# Copyright 2022 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MediaPipe Tasks GenAI Converter API."""
+
+import mediapipe.tasks.python.genai.converter.llm_converter
+
+ConversionConfig = llm_converter.ConversionConfig
+convert_checkpoint = llm_converter.convert_checkpoint
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del mediapipe
+del llm_converter

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/converter_base.py

@@ -0,0 +1,184 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Defines a couple base classes for the conversion/quantization process."""
+
+from typing import Iterator
+import os
+from typing import Dict, List, Optional, Tuple
+import numpy as np
+
+
+class QuantizationAction:
+  """Container of the tensor values and its corresponding quantization settings.
+
+  The container is responsible for hosting all of the information that is
+  required to execute the weight-only quantization.
+
+  Attributes:
+    tensor_name: A string that represents the input tensor name.
+    tensor_value: A numpy array that contains the unquantized tensor values.
+    target_name: A string that represents the updated tensor name.
+    quantize_axis: A list of integers representing the dimensions to be
+      quantized along. For example, if an input tensor has shape [128, 256] and
+      the quantize_axis==[0], it means the quantization happens along the 0-th
+      dimension, resulting in [256] scaling factors.
+    quantize_bits: An integer that specifies the target quantization bits. It
+      currently only supports either 8 or 4 bits.
+    pack_dim: An integer specifying which dimension to pack the quantized bits.
+      This is only applicable when the quantize_bits == 4.
+    is_quantized: A boolean indicating whether the tensor is already quantized
+      in the checkpoint. Default to False.
+  """
+
+  def __init__(
+      self,
+      tensor_name: str,
+      tensor_value: Optional[np.ndarray] = None,
+      target_name: Optional[str] = None,
+      quantize_axis: Optional[List[int]] = None,
+      quantize_bits: Optional[int] = None,
+      pack_dim: Optional[int] = 0,
+      is_quantized: bool = False,
+  ):
+    """Initializes the model attributes."""
+    self.tensor_name = tensor_name
+    self.tensor_value = tensor_value
+    self.target_name = target_name
+    self.quantize_axis = quantize_axis
+    self.quantize_bits = quantize_bits
+    self.pack_dim = pack_dim
+    self.is_quantized = is_quantized
+
+  def __str__(self) -> str:
+    output_string = "QuantizationAction(\n"
+    output_string += f"  tensor_name: {self.tensor_name}\n"
+    output_string += f"  target_name: {self.target_name}\n"
+    output_string += f"  quantize_axis: {self.quantize_axis}\n"
+    output_string += f"  quantize_bits: {self.quantize_bits}\n"
+    output_string += f"  pack_dim: {self.pack_dim}\n"
+    output_string += f"  is_quantized: {self.is_quantized}\n"
+    if self.tensor_value is not None:
+      output_string += f"  tensor_value: {self.tensor_value.shape}\n"
+    output_string += ")\n"
+    return output_string
+
+
+class CkptLoaderBase:
+  """Base class for loading the checkpoint.
+
+  This class is responsible for loading the checkpoint files into the layer
+  weight tensors (as numpy arrays) + quantization setting information (8/4
+  bits). The returned data should be a list of QuantizationAction that describes
+  how to quantize each layer weights.
+  """
+
+  def __init__(
+      self,
+      ckpt_path: str,
+      is_symmetric: bool,
+      attention_quant_bits: int,
+      feedforward_quant_bits: int,
+      embedding_quant_bits: int,
+  ):
+    """Initializes the loader.
+
+    Args:
+      ckpt_path: The filepath to the checkpoint.
+      is_symmetric: Whether to apply symmetric or asymmetric quantization.
+      attention_quant_bits: An integer that specify the target quantization bits
+        (support 8 or 4) for the attention layers.
+      feedforward_quant_bits: An integer that specify the target quantization
+        bits (support 8 or 4) for the feedforward layers in each Transformer
+        blocks.
+      embedding_quant_bits: An integer that specify the target quantization bits
+        (support 8 or 4) for the embedding (and the final projection) layers.
+    """
+    self._ckpt_path = ckpt_path
+    self._is_symmetric = is_symmetric
+    self._attention_quant_bits = attention_quant_bits
+    self._feedforward_quant_bits = feedforward_quant_bits
+    self._embedding_quant_bits = embedding_quant_bits
+
+  def load_to_actions(
+      self,
+  ) -> Iterator[Optional[List[QuantizationAction]]]:
+    """Loads the checkpoint and returns the quantization actions."""
+    raise NotImplementedError("The load_to_actions method is not implemented.")
+
+
+class LayerActionMapperBase:
+  """Base class for mapping the layer weights to quantization actions.
+
+  This class is responsible for mapping from each layer to its corresponding
+  quantization information (e.g. target quantization bits / updated tensor
+  name...).
+  """
+
+  def __init__(
+      self,
+      is_symmetric: bool,
+      attention_quant_bits: int,
+      feedforward_quant_bits: int,
+      embedding_quant_bits: int,
+      backend: str,
+  ):
+    self._is_symmetric = is_symmetric
+    self._attention_quant_bits = attention_quant_bits
+    self._feedforward_quant_bits = feedforward_quant_bits
+    self._embedding_quant_bits = embedding_quant_bits
+    self._backend = backend
+
+  def map_to_actions(
+      self, layer_name: str
+  ) -> Optional[List[QuantizationAction]]:
+    """Maps the layer weights to quantization actions.
+
+    Args:
+      layer_name: A string representing the name of the layer weight. Note that
+        it is expected the layer information is contained in the name which is
+        enough to determine the target quantization information. Any child class
+        is expected to implement this function.
+    """
+    raise NotImplementedError("The map_to_actions method is not implemented.")
+
+
+class ModelWriterBase:
+  """Base class for writing the quantized model.
+
+  This class is responsible for taking a dictionary of the quantized
+  tensors/names and writing them into the format that can be loaded by the
+  on-device inference engine.
+  """
+
+  def __init__(self, output_dir: str, backend: str):
+    """Initializes the class.
+
+    Args:
+      output_dir: A string that represents the output directory to write the
+        resulting file(s).
+      backend: A string that represents the target backend to run the output
+        file(s).
+    """
+    self._output_dir = output_dir
+    if not os.path.exists(self._output_dir):
+      os.mkdir(self._output_dir)
+    self._backend = backend
+
+  def write_variables(
+      self,
+      variables: Dict[str, Tuple[np.ndarray, bool]],
+      use_fake_values: bool = False,
+  ):
+    raise NotImplementedError("The write_variables method is not implemented.")

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/converter_factory.py

@@ -0,0 +1,79 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility library that helps create the converter instances."""
+from mediapipe.tasks.python.genai.converter import converter_base
+from mediapipe.tasks.python.genai.converter import pytorch_converter
+from mediapipe.tasks.python.genai.converter import safetensors_converter
+from mediapipe.tasks.python.genai.converter import weight_bins_writer
+
+
+def create_ckpt_loader(
+    ckpt_format: str, *args, **kwargs
+) -> converter_base.CkptLoaderBase:
+  """Creates the checkpoint loader.
+
+  Args:
+    ckpt_format: A string that indicates which input checkpoint format is.
+    *args: Additional arguments to be passed into the loader.
+    **kwargs: Additional arguments to be passed into the loader.
+
+  Returns:
+    A created CkptLoader instance.
+  """
+  del args
+  if ckpt_format == "pytorch":
+    return pytorch_converter.PytorchCkptLoader(
+        ckpt_path=kwargs["ckpt_path"],
+        is_symmetric=kwargs["is_symmetric"],
+        attention_quant_bits=kwargs["attention_quant_bits"],
+        feedforward_quant_bits=kwargs["feedforward_quant_bits"],
+        embedding_quant_bits=kwargs["embedding_quant_bits"],
+        special_model=kwargs["special_model"],
+        backend=kwargs["backend"],
+    )
+  elif ckpt_format == "safetensors":
+    return safetensors_converter.SafetensorsCkptLoader(
+        ckpt_path=kwargs["ckpt_path"],
+        is_symmetric=kwargs["is_symmetric"],
+        attention_quant_bits=kwargs["attention_quant_bits"],
+        feedforward_quant_bits=kwargs["feedforward_quant_bits"],
+        embedding_quant_bits=kwargs["embedding_quant_bits"],
+        special_model=kwargs["special_model"],
+        backend=kwargs["backend"],
+    )
+  else:
+    raise ValueError(f"Unknown checkpoint format: {ckpt_format}")
+
+
+def create_writer(
+    writer_type: str, *args, **kwargs
+) -> converter_base.ModelWriterBase:
+  """Creates the model writer.
+
+  Args:
+    writer_type: A string the indicates which model writer to create.
+    *args: Additional arguments to be passed into the loader.
+    **kwargs: Additional arguments to be passed into the loader.
+
+  Returns:
+    A created ModelWriter instance.
+  """
+  del args
+  if writer_type == "weight_bins":
+    return weight_bins_writer.WeightBinsWriter(
+        output_dir=kwargs["output_dir"], backend=kwargs["backend"]
+    )
+  else:
+    raise ValueError(f"Unknown writer type: {writer_type}")

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/external_dependencies.py

@@ -0,0 +1,26 @@
+# Copyright 2025 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MediaPipe LLM Converter Dependencies."""
+try:
+  import jax
+  from jax import lax
+  from jax import numpy as jnp
+  import torch
+
+  import sentencepiece
+except ImportError as e:
+  raise ImportError(
+      "JAX, Torch, and SentencePiece are required for the LLM conversions."
+      " Please run `pip install torch jax sensetencepiece`."
+  ) from e

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/llm_converter.py

@@ -0,0 +1,644 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Functions to perform the checkpoint conversion."""
+
+import contextlib
+import ctypes
+import os
+from typing import Any, List, Optional
+
+from absl import logging
+import numpy as np
+
+from mediapipe.tasks.python.core import mediapipe_c_bindings
+from mediapipe.tasks.python.core import mediapipe_c_utils
+from mediapipe.tasks.python.core import serial_dispatcher
+from mediapipe.tasks.python.genai.converter import converter_base
+from mediapipe.tasks.python.genai.converter import converter_factory
+from mediapipe.tasks.python.genai.converter import external_dependencies
+from mediapipe.tasks.python.genai.converter import quantization_util
+
+jnp = external_dependencies.jnp
+
+_CTYPES_SIGNATURES = (
+    mediapipe_c_utils.CStatusFunction(
+        func_name='MpLlmConverterGenerateCpuTfLite',
+        core_argtypes=(
+            ctypes.c_char_p,  # model_type
+            ctypes.c_char_p,  # weight_path
+            ctypes.c_char_p,  # vocab_model_file
+            ctypes.c_bool,  # is_quantized
+            ctypes.c_char_p,  # output_tflite_file
+        ),
+    ),
+    mediapipe_c_utils.CStatusFunction(
+        func_name='MpLlmConverterGenerateGpuTfLite',
+        core_argtypes=(
+            ctypes.c_char_p,  # model_type
+            ctypes.c_char_p,  # weight_path
+            ctypes.c_char_p,  # vocab_model_file
+            ctypes.c_bool,  # is_quantized
+            ctypes.c_bool,  # obfuscate
+            ctypes.c_char_p,  # output_tflite_file
+            ctypes.c_int,  # lora_rank
+            ctypes.c_char_p,  # lora_weight_path
+            ctypes.c_char_p,  # lora_output_tflite_file
+            ctypes.c_char_p,  # lora_main_model_type
+            ctypes.c_char_p,  # image_encoder_file
+            ctypes.c_char_p,  # image_adapter_file
+            ctypes.c_char_p,  # submodel_type
+            ctypes.c_bool,  # use_dynamic_ple
+            ctypes.c_bool,  # apply_srq
+            ctypes.c_int,  # block_size
+        ),
+    ),
+    mediapipe_c_utils.CStatusFunction(
+        func_name='MpLlmConverterConvertHfTokenizer',
+        core_argtypes=(
+            ctypes.c_char_p,  # vocab_model_file
+            ctypes.c_char_p,  # output_vocab_file
+        ),
+    ),
+)
+
+
+def _safe_encode_str(value: str | None) -> bytes:
+  """Encodes a string to bytes, returning an empty byte string if None."""
+  return value.encode('utf-8') if value is not None else b''
+
+
+class ConversionConfig(object):
+  """Config for checkpoint conversion.
+
+  Attributes:
+    input_ckpt: Directory or path for the input checkpoint.
+    ckpt_format: Checkpoint format, e.g. 'safetensors', 'pytorch'.
+    model_type: Name of the model, e.g. GEMMA_2B.
+    backend: Target backend to run the model. Can be either "cpu" or "gpu".
+    output_dir: Where the output file(s) to be stored.
+    is_quantized: Whether the checkpoint is already quantized. If the checkpoint
+      is already quantized, the converter will not quantize it again and it will
+      ignore the quantization parameters.
+    is_symmetric: Whether to quantize symmetrically.
+    attention_quant_bits: Target quantization bits for the attention layers.
+    feedforward_quant_bits: Target quantization bits for the feedforward layers.
+    embedding_quant_bits: Target quantization bits for the embedding layers.
+    combine_file_only: Whether to combine the weight files only (assuming the
+      weight files are already existed).
+    vocab_model_file: The file path to the 1) SentencePiece vocab model; 2)
+      Hugging Face BPE tokenizer files; 1) is applicable for the Gemma model and
+      2) is applicable for other models. When 2) is used, the provided path is
+      expected to point to a directory that contains both tokenizer.json and
+      tokenizer_config.json files.
+    obfuscate: Whether to obfuscate the model.
+    output_tflite_file: (optional) the output tflite filename. If not provided,
+      the output will be `model.tflite` stored in the output_dir.
+    fp16_scale: A scalar value between [0, 1]. Some models can run into
+      activation overflow issue when running in 16-bit floating point mode. To
+      solve this, we need to scale down the weights of certain layers. See
+      go/llm-on-device-fp16 for more detailed explanation.
+    lora_ckpt: The directory or path for the lora checkpoint. Required in order
+      to convert the lora weights.
+    lora_rank: An integer representing the rank of LoRA. Required in order to
+      convert the lora weights.If not provided, then the converter assumes there
+      is no LoRA weights. Note that only the GPU backend supports LoRA.
+    lora_alpha: A float representing the scale of LoRA weights.
+    lora_output_tflite_file: A string indicating the name of the generated
+      tflite file for the LoRA weight. Only applicable when the lora_rank is not
+      zero.
+    lora_main_model_type: If the LoRA will be applied to a submodel packaged
+      with a main model, what the main model type is.
+    image_encoder_file: A string with the name of the image encoder tflite file.
+    image_adapter_file: A string with the name of the image adapter tflite file.
+    submodel_type: Name of submodel, e.g. GEMMA_2B.
+    use_fake_weights: Whether to use fake weights. If set to True, the weights
+      will be filled with zeros.
+    use_dynamic_ple: Whether any PLE embeddings should be loaded dynamically.
+      Default is true, which will cause embeddings to only be loaded into VRAM
+      on demand.
+    use_mse_quant: Whether to use MSE quantization for recomputing scales.
+    block_size: Default of 0 is off. Can also be 32 or 128, for using blockwise
+      Q4_0 compression on the models and layers which support this.
+    apply_srq: Default of None will apply SRQ for pre-quantized models only.
+      Will override this behavior if set explicitly.
+  """
+
+  def __init__(
+      self,
+      input_ckpt: str,
+      ckpt_format: str,
+      model_type: str,
+      backend: str,
+      output_dir: str,
+      is_quantized: bool = False,
+      is_symmetric: bool = True,
+      attention_quant_bits: int = 8,
+      feedforward_quant_bits: int = 8,
+      embedding_quant_bits: int = 8,
+      combine_file_only: bool = False,
+      vocab_model_file: str = '',
+      obfuscate: bool = False,
+      output_tflite_file: Optional[str] = None,
+      fp16_scale: Optional[float] = None,
+      block_size: int = 0,
+      lora_ckpt: Optional[str] = None,
+      lora_rank: Optional[int] = None,
+      lora_alpha: Optional[float] = None,
+      lora_output_tflite_file: Optional[str] = None,
+      lora_main_model_type: Optional[str] = None,
+      image_encoder_file: Optional[str] = None,
+      image_adapter_file: Optional[str] = None,
+      submodel_type: Optional[str] = None,
+      use_fake_weights: bool = False,
+      use_dynamic_ple: bool = True,
+      use_mse_quant: bool = False,
+      apply_srq: Optional[bool] = None,
+  ):
+    self.input_ckpt = input_ckpt
+    self.ckpt_format = ckpt_format
+    self.model_type = model_type
+    self.backend = backend
+    if os.path.isfile(output_dir):
+      raise ValueError('Output directory mush not point to an existing file.')
+    if not os.path.isdir(output_dir):
+      logging.info('Creating output directory: %s', output_dir)
+      os.makedirs(output_dir, exist_ok=True)
+    self.output_dir = output_dir
+    self.is_quantized = is_quantized
+    self.is_symmetric = is_symmetric
+    self.attention_quant_bits = attention_quant_bits
+    self.feedforward_quant_bits = feedforward_quant_bits
+    self.embedding_quant_bits = embedding_quant_bits
+    self.combine_file_only = combine_file_only
+    self.vocab_model_file = vocab_model_file
+    self.obfuscate = obfuscate
+    self.image_encoder_file = image_encoder_file
+    self.image_adapter_file = image_adapter_file
+    self.submodel_type = submodel_type
+    self.use_fake_weights = use_fake_weights
+    self.use_dynamic_ple = use_dynamic_ple
+    self.use_mse_quant = use_mse_quant
+    self.apply_srq = apply_srq
+    if output_tflite_file:
+      parent_dir = os.path.dirname(output_tflite_file)
+      if not os.path.isdir(parent_dir):
+        logging.info('Creating tflite parent directory: %s', parent_dir)
+        os.makedirs(parent_dir, exist_ok=True)
+      self.output_tflite_file = output_tflite_file
+    else:
+      self.output_tflite_file = os.path.join(output_dir, 'model.tflite')
+
+    self.fp16_scale = None
+    if block_size not in (0, 32, 128):
+      raise ValueError('Block size must be one of (0 (off), 32, 128)')
+    else:
+      self.block_size = block_size
+    self.lora_ckpt = lora_ckpt
+    self.lora_rank = lora_rank
+    self.lora_alpha = lora_alpha
+    self.lora_output_tflite_file = lora_output_tflite_file
+    self.lora_main_model_type = lora_main_model_type
+    if (self.lora_ckpt is None) ^ (self.lora_rank is None):
+      raise ValueError(
+          'lora_ckpt and lora_rank must be either both provided or both not'
+          ' provided.'
+      )
+    if self.lora_rank is not None:
+      if backend == 'cpu':
+        raise ValueError('LoRA is not supported for CPU backend.')
+      lora_applicable_models = [
+          'GEMMA_2B',
+          'GEMMA2_2B',
+          'PHI_2',
+          'GEMMA3_1B',
+          'GEMMA3_4B',
+          'GEMMA3_12B',
+          'GEMMA3_27B',
+          'GEMMA3_300M',
+      ]
+      if model_type not in lora_applicable_models:
+        raise ValueError(
+            'LoRA is only applicable for the model_type:'
+            f' {", ".join(lora_applicable_models)}, but get model_type:'
+            f' {model_type}.'
+        )
+
+
+@contextlib.contextmanager
+def filemanager(filename: str, mode: str):
+  try:
+    with open(filename, mode) as f:
+      yield f
+  finally:
+    pass
+
+
+class _LlmConverter:
+  """Bundles all conversion logic for LLM models."""
+
+  def __init__(self, lib: Any):
+    self._lib = serial_dispatcher.SerialDispatcher(lib, _CTYPES_SIGNATURES)
+
+  def quantize_by_actions(
+      self,
+      actions: List[converter_base.QuantizationAction],
+      backend: str,
+      is_symmetric: bool,
+      use_mse_quant: bool = False,
+      block_size: int = 0,
+  ):
+    """Quantizes the weights by actions.
+
+    Args:
+      actions: A list of QuantizationAction that contains the information and
+        tensor values to be quantized.
+      backend: Target backend to run the model. Can be either "cpu" or "gpu".
+      is_symmetric: Whether to quantize symmetrically.
+      use_mse_quant: Whether to use MSE quantization for recomputing scales.
+      block_size: 0 to disable, and 32 or 128 to use blockwise Q4_0 for
+        recomputing scales.
+
+    Returns:
+      A dictionary that maps from the updated tensor names to the quantized
+      tensor values + a boolean that indicates whether the tensor values need to
+      be packed (only applicable for the 4-bit quantized weights).
+    """
+    output_tensors = {}
+    qvalue_suffix = '_quantized_value'
+    scale_suffix = '_quantized_scale'
+    zp_suffix = '_quantized_zp'
+    for action in actions:
+      if action.tensor_value is None:
+        continue
+      # The dtype needs to be compared in string as it is a custom numpy dtype.
+      # Explicitly cast the bfloat16 and float16 dtype to float32 to make sure
+      # its value is converted and serialized correctly.
+      if (
+          str(action.tensor_value.dtype) == 'bfloat16'
+          or action.tensor_value.dtype == np.float16
+      ):
+        action.tensor_value = action.tensor_value.astype(np.float32)
+      if (
+          (not action.is_quantized)
+          and action.tensor_value.dtype != np.float32
+          and action.tensor_value.dtype != np.int8
+      ):
+        raise ValueError(
+            'All tensors should be casted to either float32 or int8, but got: '
+            '%s'
+            % action.tensor_value.dtype
+        )
+      if action.is_quantized:
+        pack = action.tensor_value.dtype == jnp.int4
+        if qvalue_suffix in action.target_name:
+          target_name = action.target_name[: -len(qvalue_suffix)]
+          # Stores the quantized value in int8 for 4-bit quantization.
+          if pack:
+            action.tensor_value = action.tensor_value.astype(jnp.int8)
+          output_tensors[target_name] = (action.tensor_value, pack)
+        elif (
+            scale_suffix in action.target_name
+            or zp_suffix in action.target_name
+        ):
+          output_tensors[action.target_name] = (
+              action.tensor_value,
+              False,
+          )
+        else:
+          output_tensors[action.target_name] = (action.tensor_value, False)
+      if action.quantize_axis:
+        quant_bits = action.quantize_bits
+        quant_axis = action.quantize_axis.copy()
+        if (
+            block_size > 0
+            and quant_bits == 4
+            and quant_axis != [0]
+            and quant_axis != [1]
+            and quant_axis != [0, 1]
+        ):
+          # Turn off 4bit blockwise quant for unsupported quantization axis
+          # configurations
+          print(
+              'Disabling blockwise 4bit quant for unsupported quantization'
+              ' axis: ',
+              quant_axis,
+              ' on tensor ',
+              action.target_name,
+          )
+          quant_bits = 8
+        pack = quant_bits == 4
+
+        if action.tensor_value.dtype == np.int8:
+          if backend == 'cpu' and pack:
+            raise ValueError(
+                'Converting pre-quantized checkpoint into 4-bit is not'
+                ' supported for CPU backend.'
+            )
+          output_tensors[action.target_name] = (action.tensor_value, pack)
+        else:
+          if is_symmetric:
+            bs = block_size if pack else 0
+            target_var, scale = quantization_util.quantize_tensor(
+                var=action.tensor_value,
+                axis=action.quantize_axis,
+                sym=is_symmetric,
+                number_bits=quant_bits,
+                use_mse_quant=use_mse_quant,
+                # blockwise quant applies to symmetric 4bit only
+                block_size=bs,
+            )
+            if bs > 0:
+              if len(scale.shape) > 3:
+                raise ValueError(
+                    'Blockwise int4 quantization not yet supported for'
+                    ' 4-dimensional or higher scales.'
+                )
+              # Reshape and transpose output from blockwise quant.
+              # We want per-block scale values as the last dimension.
+              if quant_axis == [0] or quant_axis == [0, 1]:
+                if len(scale.shape) <= 2:
+                  scale = scale.transpose()
+                else:
+                  scale = scale.transpose([1, 2, 0])
+              elif quant_axis == [1]:
+                if len(scale.shape) == 3:
+                  # Only tested in one case, so not sure if this generalizes.
+                  # Seems to behave identically to our experimental
+                  # quant_axis = [0, 1] support so far.
+                  scale = scale.transpose([2, 0, 1])
+
+              # Reshape to fold in the per-block dimension
+              target_var = target_var.reshape(action.tensor_value.shape)
+            output_tensors[action.target_name] = (target_var, pack)
+            output_tensors[action.target_name + scale_suffix] = (
+                scale,
+                False,
+            )
+            zp = None
+          else:
+            target_var, scale, zp = quantization_util.quantize_tensor(
+                var=action.tensor_value,
+                axis=action.quantize_axis,
+                sym=is_symmetric,
+                number_bits=action.quantize_bits,
+                use_mse_quant=use_mse_quant,
+            )
+          if backend == 'cpu' and pack:
+            target_var, scale, zp = quantization_util.update_to_uint4(
+                target_var, scale, zp
+            )
+          output_tensors[action.target_name] = (target_var, pack)
+          output_tensors[action.target_name + scale_suffix] = (scale, False)
+          if zp is not None:
+            output_tensors[action.target_name + zp_suffix] = (zp, False)
+      else:
+        output_tensors[action.target_name] = (action.tensor_value, False)
+    return output_tensors
+
+  def combined_weight_bins_to_tflite(
+      self,
+      model_type: str,
+      backend: str,
+      weight_path: str,
+      output_tflite_file: str,
+      obfuscate: bool,
+      vocab_model_file: str,
+      lora_rank: Optional[int] = None,
+      lora_weight_path: Optional[str] = None,
+      lora_output_tflite_file: Optional[str] = None,
+      lora_main_model_type: Optional[str] = None,
+      image_encoder_file: Optional[str] = None,
+      image_adapter_file: Optional[str] = None,
+      submodel_type: Optional[str] = None,
+      use_dynamic_ple: Optional[bool] = None,
+      apply_srq: Optional[bool] = None,
+      block_size: int = 0,
+  ):
+    """Combines weight files to tflite file."""
+    if backend == 'cpu':
+      if lora_rank is not None:
+        logging.fatal('LoRA is not supported for CPU backend.')
+      self._lib.MpLlmConverterGenerateCpuTfLite(
+          model_type.encode('utf-8'),
+          weight_path.encode('utf-8'),
+          vocab_model_file.encode('utf-8'),
+          True,
+          output_tflite_file.encode('utf-8'),
+      )
+    elif backend == 'gpu':
+      self._lib.MpLlmConverterGenerateGpuTfLite(
+          model_type.encode('utf-8'),
+          weight_path.encode('utf-8'),
+          vocab_model_file.encode('utf-8'),
+          True,
+          obfuscate,
+          output_tflite_file.encode('utf-8'),
+          0 if lora_rank is None else lora_rank,
+          _safe_encode_str(lora_weight_path),
+          _safe_encode_str(lora_output_tflite_file),
+          _safe_encode_str(lora_main_model_type),
+          _safe_encode_str(image_encoder_file),
+          _safe_encode_str(image_adapter_file),
+          _safe_encode_str(submodel_type),
+          True if use_dynamic_ple is None else use_dynamic_ple,
+          False if apply_srq is None else apply_srq,
+          block_size,
+      )
+    else:
+      raise ValueError(f'Unsupported backend: {backend}')
+
+  def convert_bpe_vocab(self, vocab_model_file: str, output_dir: str) -> str:
+    """Converts the BPE vocab model file to SPM format.
+
+    Args:
+      vocab_model_file: The input BPE vocab model file path.
+      output_dir: The output directory to store the SPM vocab model file.
+
+    Returns:
+      The path to the output SPM vocab model file.
+    """
+    if not os.path.isdir(vocab_model_file):
+      raise ValueError(
+          'The input BPE vocab model file path is expected to be a directory'
+          ' that contains both tokenizer.json and tokenizer_config.json files.'
+      )
+    output_vocab_file = os.path.join(output_dir, 'spm.model')
+    self._lib.MpLlmConverterConvertHfTokenizer(
+        vocab_model_file.encode('utf-8'),
+        output_vocab_file.encode('utf-8'),
+    )
+    return output_vocab_file
+
+  def sort_layer_info(self, layer_info_file: str) -> None:
+    """Loads and sorts the layer info file."""
+    layer_info = []
+    with filemanager(layer_info_file, 'r') as finfo:
+      for line in finfo:
+        line = line.strip()
+        if line:
+          layer_info.append(line)
+    layer_info = list(set(layer_info))
+    layer_info.sort()
+    with filemanager(layer_info_file, 'w') as finfo:
+      for line in layer_info:
+        finfo.write(line + '\n')
+        finfo.write('\n')
+
+  def maybe_quantize_and_write_tensors_to_bins(
+      self,
+      ckpt_loader: converter_base.CkptLoaderBase,
+      config: ConversionConfig,
+  ) -> None:
+    """Quantizes the weight tensors according to the loader and writes them to bins."""
+    actions = ckpt_loader.load_to_actions()
+
+    for action in actions:
+      # Quantize the weight
+      quantized_tensors = self.quantize_by_actions(
+          action,
+          config.backend,
+          config.is_symmetric,
+          config.use_mse_quant,
+          block_size=config.block_size,
+      )
+      del action
+      # Write the tensors into file(s).
+      writer = converter_factory.create_writer(
+          writer_type='weight_bins',
+          output_dir=config.output_dir,
+          backend=config.backend,
+      )
+      writer.write_variables(quantized_tensors, config.use_fake_weights)
+      del quantized_tensors
+      del writer
+
+  def convert_checkpoint(self, config: ConversionConfig) -> None:
+    """Converts the checkpoint to tflite file."""
+    logging.info('input folder: %s', config.input_ckpt)
+
+    if os.path.isdir(config.vocab_model_file):
+      vocab_model_path = self.convert_bpe_vocab(
+          config.vocab_model_file, config.output_dir
+      )
+    else:
+      vocab_model_path = config.vocab_model_file
+
+    if not config.combine_file_only:
+      # Load the layer weights and prepare the quantization configurations.
+      loader = converter_factory.create_ckpt_loader(
+          config.ckpt_format,
+          ckpt_path=config.input_ckpt,
+          is_quantized=config.is_quantized,
+          is_symmetric=config.is_symmetric,
+          backend=config.backend,
+          attention_quant_bits=config.attention_quant_bits,
+          feedforward_quant_bits=config.feedforward_quant_bits,
+          embedding_quant_bits=config.embedding_quant_bits,
+          special_model=config.model_type,
+          fp16_scale=config.fp16_scale,
+      )
+      self.maybe_quantize_and_write_tensors_to_bins(loader, config)
+
+      if config.lora_ckpt is not None and config.lora_ckpt != config.input_ckpt:
+        # If lora ckpt and the input ckpt is the same. The lora conversion is
+        # handled in the previous loader.
+        lora_loader = converter_factory.create_ckpt_loader(
+            config.ckpt_format,
+            ckpt_path=config.lora_ckpt,
+            is_quantized=config.is_quantized,
+            is_symmetric=config.is_symmetric,
+            backend=config.backend,
+            attention_quant_bits=config.attention_quant_bits,
+            feedforward_quant_bits=config.feedforward_quant_bits,
+            embedding_quant_bits=config.embedding_quant_bits,
+            special_model=config.model_type,
+        )
+        self.maybe_quantize_and_write_tensors_to_bins(lora_loader, config)
+
+      self.sort_layer_info(os.path.join(config.output_dir, 'layer_info.txt'))
+
+    self.combined_weight_bins_to_tflite(
+        config.model_type,
+        config.backend,
+        weight_path=config.output_dir,
+        output_tflite_file=config.output_tflite_file,
+        obfuscate=config.obfuscate,
+        vocab_model_file=vocab_model_path,
+        lora_rank=config.lora_rank,
+        lora_weight_path=config.output_dir,
+        lora_output_tflite_file=config.lora_output_tflite_file,
+        lora_main_model_type=config.lora_main_model_type,
+        image_encoder_file=config.image_encoder_file,
+        image_adapter_file=config.image_adapter_file,
+        submodel_type=config.submodel_type,
+        use_dynamic_ple=config.use_dynamic_ple,
+        # Default SRQ behavior is to match is_quantized
+        apply_srq=config.is_quantized
+        if config.apply_srq is None
+        else config.apply_srq,
+        block_size=config.block_size,
+    )
+
+  def __enter__(self) -> '_LlmConverter':
+    """Returns `self` upon entering the runtime context."""
+    return self
+
+  def __exit__(self, *args) -> None:
+    """Shuts down the LlmConverter on exit of the context manager.
+
+    Args:
+      *args: Unused.
+    Raises:
+      RuntimeError: If the LLM converter failed to close.
+    """
+    del args  # Unused.
+    self._lib.close()
+
+
+def convert_checkpoint(config: ConversionConfig) -> None:
+  """Converts the checkpoint to tflite file."""
+  lib = mediapipe_c_bindings.load_raw_library(_CTYPES_SIGNATURES)
+  with _LlmConverter(lib) as converter:
+    converter.convert_checkpoint(config)
+
+
+def quantize_by_actions(
+    actions: List[converter_base.QuantizationAction],
+    backend: str,
+    is_symmetric: bool,
+    use_mse_quant: bool = False,
+    block_size: int = 0,
+):
+  """Quantizes the weights by actions.
+
+  Args:
+    actions: A list of QuantizationAction that contains the information and
+      tensor values to be quantized.
+    backend: Target backend to run the model. Can be either "cpu" or "gpu".
+    is_symmetric: Whether to quantize symmetrically.
+    use_mse_quant: Whether to use MSE quantization for recomputing scales.
+    block_size: 0 to disable, and 32 or 128 to use blockwise quantization for
+      recomputing scales.
+
+  Returns:
+    A dictionary that maps from the updated tensor names to the quantized
+    tensor values + a boolean that indicates whether the tensor values need to
+    be packed (only applicable for the 4-bit quantized weights).
+  """
+  lib = mediapipe_c_bindings.load_raw_library(_CTYPES_SIGNATURES)
+  with _LlmConverter(lib) as converter:
+    return converter.quantize_by_actions(
+        actions, backend, is_symmetric, use_mse_quant, block_size
+    )

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/llm_converter_test.py

@@ -0,0 +1,264 @@
+# Copyright 2025 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+import unittest
+from mediapipe.tasks.python.genai.converter import converter_base
+from mediapipe.tasks.python.genai.converter import llm_converter
+
+_LLM_CONVERTER_FUNCTIONS = (
+    'MpLlmConverterGenerateCpuTfLite',
+    'MpLlmConverterGenerateGpuTfLite',
+    'MpLlmConverterConvertHfTokenizer',
+)
+
+
+class LlmConverterTest(googletest.TestCase, parameterized.TestCase):
+
+  def get_fake_action(self, input_dtype):
+    if input_dtype == 'bfloat16':
+      # Create a TensorFlow bfloat16 tensor
+      bfloat16_tensor = tf.constant([1.0, -1.0, 2.0, -2.0], dtype=tf.bfloat16)
+      # Convert the TensorFlow tensor to a NumPy array
+      tensor_value = bfloat16_tensor.numpy()
+    else:
+      tensor_value = np.array(
+          [1.0, -1.0, 2.0, -2.0], dtype=np.dtype(input_dtype)
+      )
+    return converter_base.QuantizationAction(
+        tensor_name='mdl_vars.params.lm.softmax.logits_ffn.w',
+        target_name='params.lm.softmax.logits_ffn.w',
+        quantize_axis=[0],
+        quantize_bits=8,
+        pack_dim=0,
+        tensor_value=tensor_value,
+    )
+
+  @parameterized.parameters(
+      {'input_dtype': 'float32'},
+      {'input_dtype': 'float16'},
+      {'input_dtype': 'bfloat16'},
+      {'input_dtype': 'int8'},
+  )
+  def test_quantize_by_actions(self, input_dtype):
+    out = llm_converter.quantize_by_actions(
+        [self.get_fake_action(input_dtype)], backend='gpu', is_symmetric=True
+    )
+
+    if input_dtype == 'int8':
+      # The values are pre-quantized and should be the same.
+      np.testing.assert_allclose(
+          out['params.lm.softmax.logits_ffn.w'][0],
+          np.array([1, -1, 2, -2], dtype=np.int8),
+      )
+    else:
+      np.testing.assert_allclose(
+          out['params.lm.softmax.logits_ffn.w'][0],
+          np.array([64, -64, 127, -127], dtype=np.int8),
+      )
+      np.testing.assert_allclose(
+          out['params.lm.softmax.logits_ffn.w_quantized_scale'][0],
+          np.array(0.015748, dtype=np.float32),
+          rtol=1e-03,
+      )
+
+  def test_generate_cpu_tflite_receives_correct_args(self):
+    c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
+    c_lib.MpLlmConverterGenerateCpuTfLite.return_value = 0
+
+    config = llm_converter.ConversionConfig(
+        model_type='GEMMA_2B',
+        backend='cpu',
+        input_ckpt='/tmp/ckpt',
+        ckpt_format='fake',
+        output_dir='/tmp/',
+        vocab_model_file='/tmp/vocab.model',
+        output_tflite_file='/tmp/output.tflite',
+    )
+    llm_converter_lib = llm_converter._LlmConverter(c_lib)
+
+    llm_converter_lib.combined_weight_bins_to_tflite(
+        model_type=config.model_type,
+        backend=config.backend,
+        weight_path=config.output_dir,
+        output_tflite_file=config.output_tflite_file,
+        obfuscate=config.obfuscate,
+        vocab_model_file=config.vocab_model_file,
+    )
+
+    c_lib.MpLlmConverterGenerateCpuTfLite.assert_called_once_with(
+        b'GEMMA_2B',
+        b'/tmp/',
+        b'/tmp/vocab.model',
+        True,
+        b'/tmp/output.tflite',
+        mock.ANY,  # error_message
+    )
+
+  def test_generate_cpu_tflite_propagates_failure(self):
+    c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
+    c_lib.MpLlmConverterGenerateCpuTfLite.return_value = 13  # Simulate failure
+
+    config = llm_converter.ConversionConfig(
+        model_type='GEMMA_2B',
+        backend='cpu',
+        input_ckpt='/tmp/ckpt',
+        ckpt_format='fake',
+        output_dir='/tmp/',
+        vocab_model_file='/tmp/vocab.model',
+        output_tflite_file='/tmp/output.tflite',
+    )
+    llm_converter_lib = llm_converter._LlmConverter(c_lib)
+
+    with self.assertRaises(RuntimeError):
+      llm_converter_lib.combined_weight_bins_to_tflite(
+          model_type=config.model_type,
+          backend=config.backend,
+          weight_path=config.output_dir,
+          output_tflite_file=config.output_tflite_file,
+          obfuscate=config.obfuscate,
+          vocab_model_file=config.vocab_model_file,
+      )
+
+  def test_generate_gpu_tflite_receives_correct_args(self):
+    c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
+    c_lib.MpLlmConverterGenerateGpuTfLite.return_value = 0
+
+    config = llm_converter.ConversionConfig(
+        model_type='GEMMA_2B',
+        backend='gpu',
+        input_ckpt='/tmp/ckpt',
+        ckpt_format='fake',
+        output_dir='/tmp/',
+        vocab_model_file='/tmp/vocab.model',
+        output_tflite_file='/tmp/output.tflite',
+        use_dynamic_ple=True,
+    )
+    llm_converter_lib = llm_converter._LlmConverter(c_lib)
+
+    llm_converter_lib.combined_weight_bins_to_tflite(
+        model_type=config.model_type,
+        backend=config.backend,
+        weight_path=config.output_dir,
+        output_tflite_file=config.output_tflite_file,
+        obfuscate=config.obfuscate,
+        vocab_model_file=config.vocab_model_file,
+        lora_rank=4,
+        lora_weight_path='/tmp/lora_weights',
+        lora_output_tflite_file='/tmp/lora_output.tflite',
+        lora_main_model_type='GEMMA_2B',
+        image_encoder_file='/tmp/image_encoder.tflite',
+        image_adapter_file='/tmp/image_adapter.tflite',
+        submodel_type='GEMMA_2B',
+        use_dynamic_ple=config.use_dynamic_ple,
+        apply_srq=False,
+    )
+
+    c_lib.MpLlmConverterGenerateGpuTfLite.assert_called_once_with(
+        b'GEMMA_2B',
+        b'/tmp/',
+        b'/tmp/vocab.model',
+        True,
+        False,
+        b'/tmp/output.tflite',
+        4,
+        b'/tmp/lora_weights',
+        b'/tmp/lora_output.tflite',
+        b'GEMMA_2B',
+        b'/tmp/image_encoder.tflite',
+        b'/tmp/image_adapter.tflite',
+        b'GEMMA_2B',
+        True,
+        False,
+        0,
+        mock.ANY,  # error_message
+    )
+
+  def test_generate_gpu_tflite_propagates_failure(self):
+    c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
+    c_lib.MpLlmConverterGenerateGpuTfLite.return_value = 13  # Simulate failure
+
+    config = llm_converter.ConversionConfig(
+        model_type='GEMMA_2B',
+        backend='gpu',
+        input_ckpt='/tmp/ckpt',
+        ckpt_format='fake',
+        output_dir='/tmp/',
+        vocab_model_file='/tmp/vocab.model',
+        output_tflite_file='/tmp/output.tflite',
+        use_dynamic_ple=True,
+    )
+    llm_converter_lib = llm_converter._LlmConverter(c_lib)
+
+    with self.assertRaises(RuntimeError):
+      llm_converter_lib.combined_weight_bins_to_tflite(
+          model_type=config.model_type,
+          backend=config.backend,
+          weight_path=config.output_dir,
+          output_tflite_file=config.output_tflite_file,
+          obfuscate=config.obfuscate,
+          vocab_model_file=config.vocab_model_file,
+          lora_rank=4,
+          lora_weight_path='/tmp/lora_weights',
+          lora_output_tflite_file='/tmp/lora_output.tflite',
+          lora_main_model_type='GEMMA_2B',
+          image_encoder_file='/tmp/image_encoder.tflite',
+          image_adapter_file='/tmp/image_adapter.tflite',
+          submodel_type='GEMMA_2B',
+          use_dynamic_ple=config.use_dynamic_ple,
+          apply_srq=False,
+      )
+
+  @mock.patch('os.path.isdir', return_value=True)
+  @mock.patch('os.path.join', return_value='/tmp/spm.model')
+  def test_convert_hf_tokenizer_receives_correct_args(
+      self, mock_join, mock_isdir
+  ):
+    del mock_join, mock_isdir
+
+    c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
+    c_lib.MpLlmConverterConvertHfTokenizer.return_value = 0
+
+    llm_converter_lib = llm_converter._LlmConverter(c_lib)
+    output_file = llm_converter_lib.convert_bpe_vocab(
+        vocab_model_file='/tmp/hf_tokenizer', output_dir='/tmp'
+    )
+
+    c_lib.MpLlmConverterConvertHfTokenizer.assert_called_once_with(
+        b'/tmp/hf_tokenizer', b'/tmp/spm.model', mock.ANY
+    )
+    self.assertEqual(output_file, '/tmp/spm.model')
+
+  @mock.patch('os.path.isdir', return_value=True)
+  @mock.patch('os.path.join', return_value='/tmp/spm.model')
+  def test_convert_hf_tokenizer_propagates_failure(self, mock_join, mock_isdir):
+    del mock_join, mock_isdir
+
+    c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
+    c_lib.MpLlmConverterConvertHfTokenizer.return_value = 13  # Simulate failure
+
+    llm_converter_lib = llm_converter._LlmConverter(c_lib)
+    with self.assertRaises(RuntimeError):
+      llm_converter_lib.convert_bpe_vocab(
+          vocab_model_file='/tmp/hf_tokenizer', output_dir='/tmp'
+      )
+
+
+if __name__ == '__main__':
+  googletest.main()

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/pytorch_converter.py

@@ -0,0 +1,320 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CkptLoader implementation for loading the Pytorch file."""
+
+from typing import Iterator
+import enum
+import os
+from typing import List, Optional
+
+import numpy as np
+
+from mediapipe.tasks.python.genai.converter import converter_base
+from mediapipe.tasks.python.genai.converter import external_dependencies
+
+torch = external_dependencies.torch
+
+
+class _PytorchReader:
+  """Pytorch reader."""
+
+  def __init__(self, model_path: str):
+    if not os.path.exists(model_path):
+      raise ValueError(f"{model_path} does not exists.")
+    self._model = torch.load(model_path, map_location=torch.device("cpu"))
+
+  def read_tensor_as_numpy(self, tensor_name) -> np.ndarray:
+    tensor = (
+        self._model[tensor_name]
+        .to(torch.float32)
+        .t()
+        .contiguous()
+        .detach()
+        .cpu()
+        .numpy()
+    )
+    return tensor
+
+  def get_tensor_names(self) -> List[str]:
+    names = list(self._model.keys())
+    return names
+
+
+class LayerType(enum.Enum):
+  """Enum for layer type."""
+
+  NONE = 0
+  ATTENTION = 1  # Layer is part of the attention module.
+  FEEDFORWARD = 2  # Layer is part of the feedforward module in the Transformer.
+  EMBEDDING = 3  # Layer is the embedding lookup or final projection layer.
+  LAYER_NORM = (
+      4  # Layer is layer normalization before and after attention layer.
+  )
+
+  @classmethod
+  def get_layer_type(cls, layer_name: str):
+    """Gets the layer type of the given layer name."""
+    ffn_layers = [
+        "mlp",
+    ]
+    attn_layers = [
+        "self_attention",
+    ]
+    emb_layers = [
+        "word_embeddings",
+        "lm_head",
+    ]
+    layer_norms = [
+        "input_layernorm",
+        "post_attention_layernorm",
+        "ln_f",
+    ]
+    if any(sub_name in layer_name for sub_name in attn_layers):
+      return LayerType.ATTENTION
+    if any(sub_name in layer_name for sub_name in ffn_layers):
+      return LayerType.FEEDFORWARD
+    if any(sub_name in layer_name for sub_name in emb_layers):
+      return LayerType.EMBEDDING
+    if any(sub_name in layer_name for sub_name in layer_norms):
+      return LayerType.LAYER_NORM
+    else:
+      return LayerType.NONE
+
+
+class FalconMapper(converter_base.LayerActionMapperBase):
+  """LayerActionMapper for handling the Falcon-rw-1b model."""
+
+  def __init__(
+      self,
+      is_symmetric: bool,
+      attention_quant_bits: int,
+      feedforward_quant_bits: int,
+      embedding_quant_bits: int,
+      backend: str,
+      reader: _PytorchReader,
+  ):
+    super().__init__(
+        is_symmetric=is_symmetric,
+        attention_quant_bits=attention_quant_bits,
+        feedforward_quant_bits=feedforward_quant_bits,
+        embedding_quant_bits=embedding_quant_bits,
+        backend=backend,
+    )
+    self._reader = reader
+
+  def map_to_actions(
+      self, layer_name: str
+  ) -> Optional[List[converter_base.QuantizationAction]]:
+    """Map the given layer name to actions."""
+    actions = []
+    tensor_value = self._reader.read_tensor_as_numpy(layer_name)
+    if "query_key_value" in layer_name:
+      qkv_tensors = self._decompose_falcon_qkv(tensor_value)
+      for tensor, name in zip(qkv_tensors, ["q", "k", "v"]):
+        decomposed_name = layer_name.replace("query_key_value", name)
+        action = self._map_to_action_helper(tensor, decomposed_name)
+        actions.append(action)
+    else:
+      actions.append(self._map_to_action_helper(tensor_value, layer_name))
+    return actions
+
+  def _map_to_action_helper(
+      self, tensor_value: np.ndarray, layer_name: str
+  ) -> converter_base.QuantizationAction:
+    quantize_axis = None
+    quantize_bits = None
+    layer_type = LayerType.get_layer_type(layer_name)
+
+    if layer_type != LayerType.LAYER_NORM and layer_name.endswith(".weight"):
+      layer_type = LayerType.get_layer_type(layer_name)
+      quantize_axis = [0]
+      if layer_type == LayerType.FEEDFORWARD:
+        quantize_bits = self._feedforward_quant_bits
+      elif layer_type == LayerType.ATTENTION:
+        quantize_bits = self._attention_quant_bits
+        if self._backend == "cpu" and ".dense." in layer_name:
+          tensor_value = np.transpose(tensor_value)
+          quantize_axis = [1]
+      elif layer_type == LayerType.EMBEDDING:
+        quantize_bits = self._embedding_quant_bits
+        if self._backend == "cpu" and "word_embeddings" in layer_name:
+          tensor_value = np.transpose(tensor_value)
+          quantize_axis = [1]
+    target_name = self.update_target_name(layer_name)
+
+    return converter_base.QuantizationAction(
+        tensor_name=layer_name,
+        tensor_value=tensor_value,
+        target_name=target_name,
+        quantize_axis=quantize_axis,
+        quantize_bits=quantize_bits,
+        pack_dim=0,
+    )
+
+  def _decompose_falcon_qkv(self, tensor_value: np.ndarray) -> List[np.ndarray]:
+    """Decomposes combined qkv tensor used in falcon model into separate q, k and v tensors."""
+    chunk_size = 64
+    hidden_size = 2048
+
+    tensor_value = tensor_value.transpose()
+
+    q_tensor = np.zeros(
+        (hidden_size,)
+        + ((hidden_size,) if len(tensor_value.shape) == 2 else ()),
+        dtype=tensor_value.dtype,
+    )
+    k_tensor = np.zeros_like(q_tensor, dtype=tensor_value.dtype)
+    v_tensor = np.zeros_like(k_tensor, dtype=tensor_value.dtype)
+
+    j = 0
+    for i in range(0 * chunk_size, hidden_size * 3, chunk_size * 3):
+      q_tensor[j : j + chunk_size] = tensor_value[i : i + chunk_size]
+      j += chunk_size
+
+    j = 0
+    for i in range(1 * chunk_size, hidden_size * 3, chunk_size * 3):
+      k_tensor[j : j + chunk_size] = tensor_value[i : i + chunk_size]
+      j += chunk_size
+
+    j = 0
+    for i in range(2 * chunk_size, hidden_size * 3, chunk_size * 3):
+      v_tensor[j : j + chunk_size] = tensor_value[i : i + chunk_size]
+      j += chunk_size
+
+    return [
+        np.ascontiguousarray(q_tensor.transpose()),
+        np.ascontiguousarray(k_tensor.transpose()),
+        np.ascontiguousarray(v_tensor.transpose()),
+    ]
+
+  def update_target_name(self, target_name: str) -> str:
+    """Updates the target name to match the tensor name convention."""
+    layer_type = LayerType.get_layer_type(target_name)
+
+    target_name = target_name.replace(
+        "transformer.h.", "params.lm.transformer.x_layers_"
+    )
+
+    if layer_type == LayerType.FEEDFORWARD:
+      target_name = target_name.replace(".weight", ".linear.w")
+      target_name = target_name.replace(".bias", ".bias.b")
+      target_name = target_name.replace(
+          "mlp.dense_h_to_4h", "ff_layer.ffn_layer1"
+      )
+      target_name = target_name.replace(
+          "mlp.dense_4h_to_h", "ff_layer.ffn_layer2"
+      )
+    elif layer_type == LayerType.ATTENTION:
+      target_name = target_name.replace("dense", "post")
+      target_name = target_name.replace(".weight", ".linear.w")
+      target_name = target_name.replace(".bias", ".bias.b")
+    elif layer_type == LayerType.EMBEDDING:
+      target_name = target_name.replace(
+          "transformer.word_embeddings", "params.lm.token_embedding"
+      )
+      target_name = target_name.replace(
+          "lm_head", "params.lm.softmax.logits_ffn"
+      )
+      target_name = target_name.replace(".weight", ".w")
+    elif layer_type == LayerType.LAYER_NORM:
+      target_name = target_name.replace("input_layernorm", "pre_layer_norm")
+      target_name = target_name.replace(
+          "pre_layer_norm.weight", "pre_layer_norm.scale"
+      )
+      if self._backend == "cpu":
+        target_name = target_name.replace(
+            "post_attention_layernorm", "ff_layer.pre_layer_norm"
+        )
+        target_name = target_name.replace(
+            "ff_layer.pre_layer_norm.weight", "ff_layer.pre_layer_norm.scale"
+        )
+      else:
+        target_name = target_name.replace(
+            "post_attention_layernorm", "post_layer_norm"
+        )
+        target_name = target_name.replace(
+            "post_layer_norm.weight", "post_layer_norm.scale"
+        )
+      target_name = target_name.replace(
+          "transformer.ln_f.weight", "params.lm.final_ln.scale"
+      )
+      target_name = target_name.replace(
+          "transformer.ln_f.bias", "params.lm.final_ln.bias"
+      )
+
+    return target_name
+
+
+class PytorchCkptLoader(converter_base.CkptLoaderBase):
+  """CkptLoader implementation for loading the Pytorch model."""
+
+  def __init__(
+      self,
+      ckpt_path: str,
+      is_symmetric: bool,
+      attention_quant_bits: int,
+      feedforward_quant_bits: int,
+      embedding_quant_bits: int,
+      special_model: str,
+      backend: str,
+  ):
+    """Initializes the loader.
+
+    Args:
+      ckpt_path: The filepath to the safetensors file.
+      is_symmetric: Whether to apply symmetric or asymmetric quantization.
+      attention_quant_bits: An integer that specify the target quantization bits
+        (support 8 or 4) for the attention layers.
+      feedforward_quant_bits: An integer that specify the target quantization
+        bits (support 8 or 4) for the feedforward layers in each Transformer
+        blocks.
+      embedding_quant_bits: An integer that specify the target quantization bits
+        (support 8 or 4) for the embedding (and the final projection) layers.
+      special_model: A string that indicates which input model is and whether
+        any special treatment is needed.
+      backend: A string indicating the backend used when converting this model.
+        Valid options are "cpu" and "gpu".
+    """
+    super().__init__(
+        ckpt_path,
+        is_symmetric,
+        attention_quant_bits,
+        feedforward_quant_bits,
+        embedding_quant_bits,
+    )
+
+    self._special_model = special_model
+    self._reader = _PytorchReader(ckpt_path)
+    if special_model in ["FALCON_RW_1B"]:
+      self.mapper = FalconMapper(
+          is_symmetric,
+          attention_quant_bits,
+          feedforward_quant_bits,
+          embedding_quant_bits,
+          backend,
+          self._reader,
+      )
+    else:
+      raise ValueError(f"Unknown special model: {special_model}")
+
+  def load_to_actions(
+      self,
+  ) -> Iterator[List[converter_base.QuantizationAction]]:
+    tensor_names = self._reader.get_tensor_names()
+    for tensor_name in tensor_names:
+      tensor_actions = self.mapper.map_to_actions(tensor_name)
+      if tensor_actions is None:
+        continue
+      yield tensor_actions

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/pytorch_converter_test.py

@@ -0,0 +1,86 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for pytorch_converter."""
+
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from mediapipe.tasks.python.genai.converter import pytorch_converter
+from mediapipe.tasks.python.test import test_utils
+
+_TEST_DATA_DIR = 'mediapipe/tasks/testdata/text'
+_PYTORCH_FILE = test_utils.get_test_data_path(
+    os.path.join(_TEST_DATA_DIR, 'falcon_rw_1b_test_weight.pt')
+)
+
+
+class PytorchConverterTest(parameterized.TestCase):
+  VARIABLE_NAMES = [
+      'transformer.word_embeddings.weight',
+      'transformer.h.0.input_layernorm.weight',
+      'transformer.h.0.input_layernorm.bias',
+      'transformer.h.0.self_attention.query_key_value.weight',
+      'transformer.h.0.self_attention.query_key_value.bias',
+      'transformer.h.0.self_attention.dense.weight',
+      'transformer.h.0.self_attention.dense.bias',
+      'transformer.h.0.post_attention_layernorm.weight',
+      'transformer.h.0.post_attention_layernorm.bias',
+      'transformer.h.0.mlp.dense_h_to_4h.weight',
+      'transformer.h.0.mlp.dense_h_to_4h.bias',
+      'transformer.h.0.mlp.dense_4h_to_h.weight',
+      'transformer.h.0.mlp.dense_4h_to_h.bias',
+      'transformer.ln_f.weight',
+      'transformer.ln_f.bias',
+      'lm_head.weight',
+  ]
+
+  def test_init(self):
+    loader = pytorch_converter.PytorchCkptLoader(
+        ckpt_path=_PYTORCH_FILE,
+        is_symmetric=True,
+        attention_quant_bits=8,
+        feedforward_quant_bits=8,
+        embedding_quant_bits=8,
+        special_model='FALCON_RW_1B',
+        backend='cpu',
+    )
+    self.assertEqual(loader._ckpt_path, _PYTORCH_FILE)
+    self.assertEqual(loader._is_symmetric, True)
+    self.assertEqual(loader._attention_quant_bits, 8)
+    self.assertEqual(loader._feedforward_quant_bits, 8)
+
+  @parameterized.product(
+      quant_bits=(4, 8),
+  )
+  def test_load_to_actions(self, quant_bits):
+    loader = pytorch_converter.PytorchCkptLoader(
+        ckpt_path=_PYTORCH_FILE,
+        is_symmetric=True,
+        attention_quant_bits=8,
+        feedforward_quant_bits=quant_bits,
+        embedding_quant_bits=8,
+        special_model='FALCON_RW_1B',
+        backend='cpu',
+    )
+    actions = loader.load_to_actions()
+    # There are 16 layers in the model, but qkv weight and bias would be
+    # decomposed to q, k, v tensors, so there would be 20 quantization actions.
+    self.assertEqual(sum(len(action) for action in actions), 20)
+
+
+if __name__ == '__main__':
+  absltest.main()

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/quantization_util.py

@@ -0,0 +1,694 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for quantizing tensors.
+
+Note that this is a reduced fork version of the praxis libraries to provide a
+self-contained library for packaging.
+"""
+
+from typing import Any, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+
+from mediapipe.tasks.python.genai.converter import external_dependencies
+
+jax = external_dependencies.jax
+jnp = external_dependencies.jnp
+lax = external_dependencies.lax
+
+JTensor = jax.Array
+_UINT4_ZP = 8  # Default zero point for unsigned 4-bit.
+
+# The following values are taken from fake_quantizer.py in Gemax.
+_MSE_QUANT_MULS = {
+    8: 0.05408,
+    4: 0.37755,
+}
+
+
+def _get_scan_range() -> np.ndarray:
+  # Produce candidate scan values.
+  return np.linspace(1.0, 0.5, num=11)
+
+
+def _get_mean_error(bound, t, min_value, max_value, p_value):
+  scale = bound / max_value
+  candidate = jnp.divide(t, scale)
+  candidate = jnp.clip(jnp.round(candidate), min_value, max_value)
+  candidate = jnp.multiply(candidate, scale)
+  pmean_error = jnp.mean(jnp.abs(jnp.subtract(candidate, t)) ** p_value)
+  return pmean_error
+
+
+def _get_best_bound_per_tensor(
+    t: JTensor,
+    bound: JTensor,
+    min_value: float,
+    max_value: float,
+    p_value: float = 1.0,
+) -> JTensor:
+  """Scan around [0.5, 1] * hard max value to get bound value for whole tensor.
+
+  This does a scan to get bound value(s) that minimize mean absolute error (MAE)
+  between original tensor 't' and quantized tensor. It's (almost) equivalent to
+  maximizing entropy.
+
+  Args:
+    t: The input float tensor.
+    bound: The hard max value for tensor 't'. It has the same length as shape.
+    min_value: Minimal value for the quantization bound.
+    max_value: Maximal value for the quantization bound.
+    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
+
+  Returns:
+    The best bound values for 't', that minimize p-mean error.
+  """
+
+  def _quant(scaling_factors):
+    return _get_mean_error(
+        bound * scaling_factors, t, min_value, max_value, p_value
+    )
+
+  scaling_factors = _get_scan_range()
+  diffs = jax.vmap(_quant)(scaling_factors)
+  best_scaling = scaling_factors[jnp.argmin(diffs)].astype(bound.dtype)
+  return bound * best_scaling
+
+
+def _quantrow(
+    vec: JTensor,
+    bound: JTensor,
+    min_value: float,
+    max_value: float,
+    p_value: float,
+    factors: np.ndarray,
+) -> JTensor:
+  """Get best rescaling factor from a list of factors applied a channel.
+
+  Args:
+    vec: The vector in a channel.
+    bound: The hard bound (max(abs(vec))) of the vector.
+    min_value: The target min value.
+    max_value: The target max value.
+    p_value: Exponent of the p-mean error metric.
+    factors: The values to be applied on top of bound.
+
+  Returns:
+    adjusted bound value out of the list of factors applied to bound.
+  """
+
+  def _quant(bounds):
+    return _get_mean_error(bounds, vec, min_value, max_value, p_value)
+
+  diffs = jax.vmap(_quant)(bound * factors)
+  best_scaling = factors[jnp.argmin(diffs)]
+  return bound * best_scaling
+
+
+def _get_best_bound_per_channel(
+    t: JTensor,
+    bound: JTensor,
+    min_value: float,
+    max_value: float,
+    p_value: float = 1.0,
+) -> JTensor:
+  """Scan around [0.5, 1] * hard max value to get bound value for each channel.
+
+  This does a scan to get bound value(s) that minimize mean absolute error (MAE)
+  between original tensor 't' and quantized tensor. It's (almost) equivalent to
+  maximizing entropy.
+
+  Args:
+    t: The input float tensor.
+    bound: The hard max value for tensor 't'. It has the same length as shape.
+    min_value: Minimal value for the quantization bound.
+    max_value: Maximal value for the quantization bound.
+    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
+
+  Returns:
+    The best bound values for 't', that minimize p-mean error.
+  """
+  assert len(t.shape) == 2
+  assert len(bound.shape) == 2
+  assert t.shape[1] == bound.shape[1]
+  assert bound.shape[0] == 1
+  scans = _get_scan_range()
+
+  def _quant(tensor, bound, min_value, max_value, p_value, factors):
+    ret = np.zeros(bound.shape)
+    for i in range(len(tensor)):
+      best = _quantrow(
+          tensor[i], bound[i], min_value, max_value, p_value, factors
+      )
+      ret[i] = best
+    return ret
+
+  t = t.transpose()
+  t_split = list(t)
+  res = _quant(t_split, bound[0, :], min_value, max_value, p_value, scans)
+  res = res.reshape(bound.shape)
+  return res
+
+
+def get_best_bound(
+    t: JTensor,
+    bound: JTensor,
+    min_value: float,
+    max_value: float,
+    p_value: float = 1.0,
+    per_channel: bool = False,
+) -> JTensor:
+  """Scan multiple factors on max value to get best bound value.
+
+  This does a scan to get bound value(s) that minimize mean absolute error (MAE)
+  between original tensor 't' and quantized tensor. It's (almost) equivalent to
+  maximizing entropy.
+
+  Args:
+    t: The input float tensor.
+    bound: The hard max value for tensor 't'. It has the same length as shape.
+    min_value: Minimal value for the quantization bound.
+    max_value: Maximal value for the quantization bound.
+    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
+    per_channel: if get best bound for entire tensor or per channel.
+
+  Returns:
+    The best bound values for 't', that minimize p-mean error.
+  """
+  if per_channel:
+    return _get_best_bound_per_channel(t, bound, min_value, max_value, p_value)
+  else:
+    return _get_best_bound_per_tensor(t, bound, min_value, max_value, p_value)
+
+
+def get_min_max(
+    bits: int = 8,
+    unsigned: bool = False,
+    use_fp: bool = False,
+) -> Tuple[float, float]:
+  """Gets the min/max range for a given number of bits.
+
+  Args:
+    bits: Target number of bits for quantization.
+    unsigned: If True compute min and max for unsigned number, else for signed.
+    use_fp: in floating point.
+
+  Returns:
+    min/max values for the provide number of bits.
+  """
+  if use_fp:
+    # TODO: support other fp types.
+    return -448.0, 448.0
+  # Calculation instead of jax.iinfo is used to support bits beside 4 and 8.
+  if unsigned:
+    # For unsigned 8 bits precision it is [0, 255]
+    return 0, 2**bits - 1
+  else:
+    # For signed 8 bits precision it is [-128, 127]
+    return -1 * 2 ** (bits - 1), 2 ** (bits - 1) - 1
+
+
+def pass_through(x: JTensor, fn: Any) -> JTensor:
+  # Create an exactly-zero expression with Sterbenz lemma that has an
+  # exactly-one gradient.
+  return x - jax.lax.stop_gradient(x) + jax.lax.stop_gradient(fn(x))
+
+
+def mse_reduce_precision(
+    t: JTensor,
+    contract_dims: Optional[Sequence[int]],
+    need_gradient: bool = False,
+    bits: int = 8,
+    use_symmetric: bool = True,
+    use_fp: bool = False,
+) -> Tuple[JTensor, JTensor, Optional[JTensor]]:
+  """Reduce the precision of a tensor using MSE quantization.
+
+  Args:
+    t: Input tensor.
+    contract_dims: Specifies contracting dimensions of the input tensor.
+    need_gradient: If gradient is needed out of this function.
+    bits: Target number of bits.
+    use_symmetric: If the input tensor is quantized symmetrically.
+    use_fp: Use floating point.
+
+  Returns:
+    A tuple of quantized tensor, quantization scale
+      and quantization zero point (optional).
+  """
+  if need_gradient:
+    raise ValueError('need_gradient is not supported for mse_reduce_precision.')
+  if use_fp:
+    raise ValueError('use_fp is not supported for mse_reduce_precision.')
+  if not use_symmetric:
+    raise ValueError('Gemax mse_quant is typically symmetric.')
+
+  scale_multiplier = _MSE_QUANT_MULS.get(bits)
+  if scale_multiplier is None:
+    raise ValueError(f'MSE_QUANT_MULS not defined for {bits} bits.')
+  rms = jnp.sqrt(jnp.mean(t**2, axis=contract_dims, keepdims=True))
+  scale = scale_multiplier * rms + 1e-12  # Add epsilon as in Gemax
+
+  # Calculate the zero point used in fake_quantizer.mse_quant with narrow range.
+  maxq = 2**bits - 1
+  zero_mse_quant = (maxq - 1) / 2.0  # This is 7 for 4 bits
+
+  # Simulate the quantization steps from fake_quantizer.mse_quant:
+  # 1. Scale and shift: t / scale + zero_mse_quant
+  t_shifted = jnp.divide(t, scale) + zero_mse_quant
+  # 2. Rounding
+  q_rounded = jnp.round(t_shifted)
+  # 3. Clipping to [0, maxq - 1]
+  q_clipped = jnp.clip(q_rounded, 0, maxq - 1)  # [0, 14] for 4 bits
+
+  # The quantized value (qvar) to be stored should be
+  # `q_clipped - zero_mse_quant`. This makes the dequantization match:
+  #   (qvar - zp) * scale == (q_clipped - zero_mse_quant) * scale
+  qvar = q_clipped - zero_mse_quant  # Ranges from [-7, 7] for 4 bits
+
+  # ml_drift expects a zero point for dequantization. To match the effective
+  # dequantization of mse_quant, we set the zero point to 0 (or None).
+  zp = None
+
+  container_dtype = (
+      jnp.int8 if bits <= 8 else jnp.int16 if bits <= 16 else jnp.int32
+  )
+  qvar = qvar.astype(container_dtype)
+
+  return qvar, scale, zp
+
+
+def reduce_precision(
+    t: JTensor,
+    contract_dims: Optional[Sequence[int]],
+    need_gradient: bool = False,
+    bits: int = 8,
+    optimization_on_bound: bool = False,
+    p_value: float = 1.0,
+    percentile: float = 1.0,
+    use_symmetric: bool = True,
+    use_fp: bool = False,
+    add_scale_eps: bool = False,
+    per_channel: bool = False,
+    random_rounding: bool = False,
+    key: Optional[jax.Array] = None,
+) -> Tuple[JTensor, JTensor, Optional[JTensor]]:
+  """Reduce the precision of a tensor.
+
+  Generic for all tensors.
+
+  Args:
+    t: Input tensor.
+    contract_dims: Specifies contracting dimensions of the input tensor.
+    need_gradient: If gradient is needed out of this function.
+    bits: Target number of bits.
+    optimization_on_bound: If MAE bound optimizer is used.
+    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
+    percentile: Percentile Factor to apply on the min/max range. Setting this to
+      other than 1.0 disables optimization_on_bound.
+    use_symmetric: If the input tensor is quantized symmetrically.
+    use_fp: Use floating point.
+    add_scale_eps: Add eps value or replace zero value by 1 to avoid division by
+      zero.
+    per_channel: use per-channel clipping optimization.
+    random_rounding: round with uniform random.
+    key: rng key for rounding.
+
+  Returns:
+    A tuple of quantized tensor, quantization scale
+      and quantization zero point (optional).
+  """
+  min_value, max_value = get_min_max(bits, use_fp=use_fp)
+
+  if use_symmetric:
+    bound = jnp.max(jnp.abs(t), axis=contract_dims, keepdims=True)
+    scale_bound = max_value
+  else:
+    t_max = jnp.max(t, axis=contract_dims, keepdims=True)
+    t_min = jnp.min(t, axis=contract_dims, keepdims=True)
+    bound = t_max - t_min
+    scale_bound = max_value - min_value
+
+  if percentile < 1.0:
+    bound = jnp.multiply(bound, percentile)
+  elif optimization_on_bound:
+    bound = get_best_bound(
+        t, bound, min_value, max_value, p_value, per_channel=per_channel
+    )
+
+  scale = bound / scale_bound
+
+  if add_scale_eps:
+    # Add epsilon to avoid divide-by-zero.
+    scale = scale + jnp.finfo(t.dtype).eps
+  else:
+    scale = jnp.where(scale == 0.0, 1.0, scale)
+
+  if use_symmetric:
+    zp = None
+    t = jnp.divide(t, scale)
+  else:
+    zp = min_value - t_min / scale
+    t = jnp.divide(t, scale) + zp
+    zp = jnp.multiply(scale, zp)
+
+  if use_fp:
+    # No need to round.
+    t = jnp.clip(t, min_value, max_value).astype(jnp.float8_e4m3fn)
+    # TODO: refactor to remove this logic.
+    t = jax.lax.bitcast_convert_type(t, new_dtype=jnp.int8)
+  else:
+    if need_gradient:
+      t = pass_through(t, jnp.round)
+      t = jnp.clip(t, min_value, max_value)
+    else:
+      if random_rounding:
+        t = t + jax.random.uniform(
+            key=key, shape=t.shape, minval=-0.5, maxval=0.5
+        )
+      t = jnp.round(t)
+      container_dtype = (
+          jnp.int8 if bits <= 8 else jnp.int16 if bits <= 16 else jnp.int32
+      )
+      t = jnp.clip(t, min_value, max_value).astype(container_dtype)
+
+  return t, scale, zp
+
+
+def quantize_tensor(
+    var: np.ndarray,
+    axis: List[int],
+    factor: float = 1.0,
+    sym: bool = True,
+    number_bits: int = 8,
+    use_fp: bool = False,
+    add_scale_eps: bool = False,
+    optimization_on_bound: bool = False,
+    p_value: float = 1.0,
+    per_channel: bool = False,
+    block_size: int = 0,
+    use_mse_quant: bool = False,
+) -> Union[
+    Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray]
+]:
+  """Quantize a tensor.
+
+  Args:
+    var: The variable to be quantized.
+    axis: The axis along which variable will be quantized.
+    factor: The clipping factor.
+    sym: Symmetric or asymmetric quantize the variable.
+    number_bits: Number of bits for quantized value.
+    use_fp: do fp with number of bits (i.e. fp8)
+    add_scale_eps: add epsilon to scale to avoid division by zero, else it will
+      replace zero scale by 1.
+    optimization_on_bound: If p-mean bound optimizer is used.
+    p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
+    per_channel: use per-channel clipping optimization.
+    block_size:  block size for sub-channel quantization. Defaults to 0, which
+      means off.
+    use_mse_quant: Whether to use MSE quantization.
+
+  Returns:
+    Quantized tensors, along with scales and zero point.
+  """
+  # TODO: support jnp.float8_e5m2
+  assert number_bits == 8 or number_bits == 4
+  jnp_var = jnp.asarray(var)
+  # When using sub-channel, the contracting dim is split into a sub-channel
+  # dim followed by the block dim. Therefore the contracting dim
+  # (quantize_axis) should increment by one, and the corresponding pack_dim
+  # should also increment by one.
+  if block_size > 0:
+    # We hard-code in a special case for [0, 1]. It has not been thoroughly
+    # tested or verified yet, so we log a warning.
+    if len(axis) == 2 and axis == [0, 1]:
+      print('Using experimental blockwise int4 axis=[0, 1] compression')
+      jnp_var = jnp.reshape(jnp_var, [-1, jnp_var.shape[-1]])
+      axis = [0]
+    shape = list(jnp_var.shape)
+    assert len(axis) == 1, 'Only support 1D sub-channel quantization'
+    sub_channels, rem = divmod(shape[axis[0]], block_size)
+    assert rem == 0
+    shape.insert(axis[0], sub_channels)
+    axis[0] += 1
+    shape[axis[0]] = block_size
+    jnp_var = jnp.reshape(jnp_var, shape)
+
+  if use_mse_quant:
+    qvar, scale, zp = mse_reduce_precision(
+        jnp_var,
+        contract_dims=axis,
+        need_gradient=False,
+        bits=number_bits,
+        use_symmetric=sym,
+        use_fp=False,
+    )
+  else:
+    qvar, scale, zp = reduce_precision(
+        jnp_var,
+        contract_dims=axis,
+        need_gradient=False,
+        bits=number_bits,
+        optimization_on_bound=optimization_on_bound,
+        percentile=factor,
+        use_symmetric=sym,
+        use_fp=use_fp,
+        add_scale_eps=add_scale_eps,
+        p_value=p_value,
+        per_channel=per_channel,
+    )
+  if sym:
+    return np.array(qvar), np.array(jnp.squeeze(scale, axis=axis))  # pytype: disable=wrong-arg-types  # jnp-type
+  else:
+    return (
+        np.array(qvar),
+        # CAVEAT: the following squeezes should squeeze along the quantization
+        # axis only.
+        np.array(jnp.squeeze(scale)),
+        np.array(jnp.squeeze(zp)),
+    )
+
+
+def pack_4bit(
+    x: np.ndarray, pack_dim: int, packed_dtype: jnp.dtype = jnp.int32
+) -> np.ndarray:
+  """Pack int8 or uint8 tensor where its values are actually int4 or uint4, to int32 or int8 nibble format along pack_dim.
+
+  Args:
+    x: Original int8 or uint8 tensor to pack.
+    pack_dim: Dimension to pack along. x.shape[pack_dim] must be divisible by 8,
+      when packed_dtype is int32 and divisible by 2 when target_type is int8.
+      Also pack_dim must be < x.ndim - 1.
+    packed_dtype: Target type to pack to, int32 or int8.
+
+  Returns:
+    int32 or int8 packed tensor where the pack_dim size is dividend by 8
+    from the original tensor x.
+  """
+  x = jnp.asarray(x)
+  if packed_dtype == jnp.int8 and x.dtype == jnp.uint8:
+    # It doesn't make sense to pack uint8 numbers into int4 as we'll
+    # the range overlap between uint8 and int4 is [0..7].
+    raise ValueError(
+        'only int8 input dtype is supported when packing into int8. '
+        f'Given {x.dtype}'
+    )
+
+  if x.dtype != jnp.int8 and x.dtype != jnp.uint8:
+    raise ValueError(
+        f'input dtype must be either int8 or uint8. Given {x.dtype}'
+    )
+  if pack_dim >= x.ndim - 1:
+    raise ValueError(
+        f'pack_dim must be < input ndim - 1. input shape {x.shape} and pack_dim'
+        f' {pack_dim}'
+    )
+  if packed_dtype != jnp.int32 and packed_dtype != jnp.int8:
+    raise ValueError(
+        f'packed_dtype must be either int32 or int8. Given {packed_dtype}'
+    )
+  if packed_dtype == jnp.int32 and x.shape[pack_dim] % 8 != 0:
+    raise ValueError(
+        'input shape[pack_dim] must be divisible by 8 when target_type '
+        f'is int32. Given shape {x.shape}'
+    )
+  if packed_dtype == jnp.int8 and x.shape[pack_dim] % 2 != 0:
+    raise ValueError(
+        'input shape[pack_dim] must be divisible by 2 when target_type '
+        f'is int8. Given shape {x.shape}'
+    )
+
+  int4s_per_packed_type = 8 if packed_dtype == jnp.int32 else 2
+
+  rep_shape = list(x.shape)
+  rep_shape.insert(pack_dim + 1, int4s_per_packed_type)
+  rep_shape[pack_dim] //= int4s_per_packed_type
+
+  shifts = lax.broadcasted_iota(packed_dtype, rep_shape, pack_dim + 1)
+  shifts <<= 2
+
+  # Promote x to packed_dtype
+  x = x & jnp.array(0x0F, packed_dtype)
+  x = lax.reshape(x, rep_shape)
+  x = x << shifts
+  x = lax.reduce(x, jnp.array(0x0, packed_dtype), lax.add, [pack_dim + 1])
+  return np.asarray(x)
+
+
+def pack_2bit_into_int8(x: np.ndarray, pack_dim: int) -> np.ndarray:
+  """Pack (u)int8 tensor with int2 or uint2 values into int8 along pack_dim.
+
+  Args:
+    x: Original int8 or uint8 tensor to pack.
+    pack_dim: Dimension to pack along. x.shape[pack_dim] must be divisible by 4.
+      Also pack_dim must be < x.ndim - 1.
+
+  Returns:
+    int8 packed tensor where the pack_dim size is divided by 4 from the original
+    tensor x.
+  """
+  x = jnp.asarray(x)
+  packed_dtype = jnp.int8
+  bits = 2
+  bits_per_packed_type = 8 // bits  # 4 int2 values per int8
+
+  if x.dtype != jnp.int8 and x.dtype != jnp.uint8:
+    raise ValueError(
+        f'input dtype must be either int8 or uint8. Given {x.dtype}'
+    )
+  if pack_dim >= x.ndim - 1:
+    raise ValueError(
+        f'pack_dim must be < input ndim - 1. input shape {x.shape} and pack_dim'
+        f' {pack_dim}'
+    )
+  if x.shape[pack_dim] % bits_per_packed_type != 0:
+    raise ValueError(
+        f'input shape[pack_dim] must be divisible by {bits_per_packed_type}'
+        f' when packing {bits}-bit values into {packed_dtype}. Given shape'
+        f' {x.shape}'
+    )
+
+  rep_shape = list(x.shape)
+  rep_shape.insert(pack_dim + 1, bits_per_packed_type)
+  rep_shape[pack_dim] //= bits_per_packed_type
+
+  # Create shifts: [0, 1, 2, 3]
+  shifts = lax.broadcasted_iota(packed_dtype, rep_shape, pack_dim + 1)
+  # Adjust shifts for 2-bit packing: [0, 2, 4, 6]
+  shifts <<= bits - 1  # Equivalent to shifts * 2
+
+  # Promote x to packed_dtype and mask to keep only the lowest 2 bits (0x03)
+  x = x & jnp.array(0x03, packed_dtype)
+  x = lax.reshape(x, rep_shape)
+  x = x << shifts
+  x = lax.reduce(x, jnp.array(0x0, packed_dtype), lax.add, [pack_dim + 1])
+  return np.asarray(x)
+
+
+def update_to_uint4(
+    qx: np.ndarray, scale: np.ndarray, zp: Optional[np.ndarray] = None
+):
+  """Updates the quantized weights from int4 to uint4.
+
+  This is a conversion function designed for XNNPack as it expects the 4-bit
+  quantized weight to be represented differently from the original Pax setting.
+  Specifically, the differences are:
+    1) The dynamic range of weight values: int4 (Pax) vs. uint4 (XNNPack).
+    2) The dynamic range of zero-point: float (Pax) vs. uint4 (XNNPack).
+    3) The number of zero-point: per-channel (Pax) vs. per-tensor (XNNPack).
+
+  Args:
+    qx: np.array of shape [..., channel], which is the quantized weight values
+      from Pax in the shape of. The values are in the dynamic range of int4 but
+      are hosted as int8 type. Note that if the first dimension is 3, it means
+      the qkv matrices are concatenated together and should be treated
+      differently.
+    scale: np.array of shape [1(3), channel] as np.float type, which are the
+      scaling factors for dequantization per channel.
+    zp: (optional) np.array of shape [1 (or 3), channel] as np.float type, which
+      are the zero points for dequantization per channel.
+
+  Returns:
+    A tuple (qx, scale, zp):
+      qx: The updated np.array of shape [..., channel] as np.int8 type with
+        updated dynamic range as uint4 (with 8 as the default zero points).
+      scale: Same as the input scale.
+      zp: (optional) np.array of shape [1 (or 3)] as np.int8 type with the
+        updated zero point values in the dynamic range as uint4.
+  """
+  if qx.dtype != np.int8 or ('float' not in str(scale.dtype)):
+    raise ValueError(
+        'Unexpected dtype qx:' + str(qx.dtype) + ' scale:' + str(scale.dtype)
+    )
+
+  scale = scale.astype(np.float32)
+
+  def get_new_zp(old_zp):
+    new_zp = old_zp / (scale + np.finfo(np.float32).eps)
+    per_tensor_zp = np.mean(new_zp)
+    per_tensor_zp = per_tensor_zp.astype(np.int8) + _UINT4_ZP
+    return per_tensor_zp
+
+  if zp is not None:
+    if qx.shape[0] == 3:
+      per_tensor_zp = np.stack([get_new_zp(szp) for szp in zp], axis=0)
+    else:
+      per_tensor_zp = get_new_zp(zp)
+  else:
+    per_tensor_zp = (
+        _UINT4_ZP * np.ones(shape=(3)) if qx.shape[0] == 3 else _UINT4_ZP
+    )
+
+  qx = qx + _UINT4_ZP
+  return qx, scale, np.array(per_tensor_zp, dtype=np.int32)
+
+
+def max_length_axis(shape, axis):
+  """Get the axis with the maximum length."""
+  if axis is None:
+    axis = range(len(shape))
+  elif isinstance(axis, int):
+    axis = [axis]
+  axis = tuple(a % len(shape) for a in axis)
+  axis = sorted(axis, reverse=True)  # Prefers larger dim if there is a tie.
+  max_len_ind = max(enumerate([shape[a] for a in axis]), key=lambda x: x[1])[0]
+  return axis[max_len_ind]
+
+
+def dequantize_blockwise(values, scales, block_size):
+  """Dequantizes a given blockwise-compressed tensor."""
+  contracted_dims = []
+  if len(values.shape) != len(scales.shape):
+    raise ValueError('Values and scales are not the same size')
+
+  # Find the contracted dimensions.
+  for i in range(len(values.shape)):
+    if values.shape[i] != scales.shape[i]:
+      contracted_dims.append(i)
+
+  # Make sure decompressing results in an exact shape match.
+  test_shape = list(scales.shape)
+  for i in contracted_dims:
+    test_shape[i] *= block_size
+
+  if tuple(test_shape) != values.shape:
+    raise ValueError('Expanded dimensions did not match')
+
+  # We need to repeat the scale tensor along the target dimension to match the
+  # shape of the quantized value tensor.
+  target_dim = max_length_axis(values.shape, contracted_dims)
+  scales = jnp.repeat(scales, block_size, axis=target_dim)
+  return values.astype(jnp.int32).astype(jnp.float32) * scales

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/quantization_util_test.py

@@ -0,0 +1,292 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for quantization_util."""
+
+from absl.testing import absltest
+import jax
+from jax import numpy as jnp
+import numpy as np
+
+from mediapipe.tasks.python.genai.converter import quantization_util
+
+
+_dtype = lambda x: getattr(x, 'dtype', None) or np.asarray(x).dtype
+
+
+class TestCase(absltest.TestCase):
+
+  def assertAllClose(
+      self, x, y, check_dtypes=True, rtol=1e-5, atol=1e-5, **kwargs
+  ):
+    """Wrapper for np.testing.assert_allclose()."""
+    x = np.asarray(x)
+    y = np.asarray(y)
+    if check_dtypes:
+      self.assertDtypesMatch(x, y)
+    x = x.astype(np.float32) if x.dtype == jnp.bfloat16 else x
+    y = y.astype(np.float32) if y.dtype == jnp.bfloat16 else y
+    np.testing.assert_allclose(x, y, rtol=rtol, atol=atol, **kwargs)
+
+  def assertDtypesMatch(self, x, y):
+    self.assertEqual(
+        jax.dtypes.canonicalize_dtype(_dtype(x)),
+        jax.dtypes.canonicalize_dtype(_dtype(y)),
+    )
+
+
+class Quantize8BTest(TestCase):
+
+  def test_quantize_symmetric(self):
+    inputs = np.array([[1.2, 3.1, 5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
+    qx, scale = quantization_util.quantize_tensor(inputs, axis=[1])
+
+    self.assertAllClose(
+        qx, np.array([[28, 72, 127, 67], [6, -48, 105, 127]], dtype=np.int8)
+    )
+    self.assertAllClose(
+        scale, np.array([0.04330709, 0.03149606], dtype=np.float32)
+    )
+
+  def test_quantize_symmetric_with_dimension_size_one_unquantized(self):
+    # inputs shape: (2, 1, 4), quantization axis 2.
+    inputs = np.array([[[1.2, 3.1, 5.5, 2.9]], [[0.2, -1.5, 3.3, 4.0]]])
+    qx, scale = quantization_util.quantize_tensor(inputs, axis=[2])
+
+    self.assertAllClose(
+        qx, np.array([[[28, 72, 127, 67]], [[6, -48, 105, 127]]], dtype=np.int8)
+    )
+    # expected scale shape: (2, 1)
+    self.assertAllClose(
+        scale, np.array([[0.04330709], [0.03149606]], dtype=np.float32)
+    )
+
+  def test_quantize_asymmetric(self):
+    inputs = np.array([[1.2, 3.1, 5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
+    qx, scale, zp = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=False
+    )
+
+    self.assertAllClose(
+        qx,
+        np.array([[-128, -15, 127, -27], [-49, -128, 95, 127]], dtype=np.int8),
+    )
+    self.assertAllClose(scale, np.array([0.016863, 0.021569], dtype=np.float32))
+    self.assertAllClose(zp, np.array([-3.358431, -1.260784], dtype=np.float32))
+
+
+class Quantize8BFPTest(TestCase):
+
+  def test_quantize_symmetric(self):
+    inputs = np.array([[1.0, 2.0, 5.5, 2.9], [0.02, -0.01, 3.3, 4.0]])
+    qx, scale = quantization_util.quantize_tensor(inputs, axis=[1], use_fp=True)
+
+    self.assertAllClose(
+        qx,
+        np.array([[106, 114, 126, 119], [65, -71, 124, 126]], dtype=np.int8),
+    )
+    self.assertAllClose(
+        scale, np.array([0.01227679, 0.00892857], dtype=np.float32)
+    )
+
+  def test_quantize_symmetric_with_dimension_size_one_unquantized(self):
+    # inputs shape: (2, 1, 4), quantization axis 2.
+    inputs = np.array([[[1.0, 2.0, 5.5, 2.9]], [[0.02, -0.01, 3.3, 4.0]]])
+    qx, scale = quantization_util.quantize_tensor(inputs, axis=[2], use_fp=True)
+
+    self.assertAllClose(
+        qx,
+        np.array(
+            [[[106, 114, 126, 119]], [[65, -71, 124, 126]]], dtype=np.int8
+        ),
+    )
+    # expected scale shape: (2, 1)
+    self.assertAllClose(
+        scale, np.array([[0.01227679], [0.00892857]], dtype=np.float32)
+    )
+
+  def test_quantize_asymmetric(self):
+    inputs = np.array([[-1.0, -2.0, -2.01, 2.9], [0.02, -0.15, 3.3, 4.0]])
+    qx, scale, zp = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=False, use_fp=True
+    )
+
+    self.assertAllClose(
+        qx,
+        np.array([[-8, -2, -2, 126], [-3, -2, 121, 126]], dtype=np.int8),
+    )
+    self.assertAllClose(
+        scale, np.array([0.00547991, 0.0046317], dtype=np.float32)
+    )
+    self.assertAllClose(
+        zp, np.array([-0.4449999, -1.9250002], dtype=np.float32)
+    )
+
+  def test_quantize_add_scale_eps(self):
+    inputs = np.array([[0.0, 0.0, 0.0, 0.0], [-4.0, -4.0, -4.0, -4.0]])
+    _, scale, _ = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=False, use_fp=True, add_scale_eps=True
+    )
+    self.assertAllClose(
+        scale, np.array([np.finfo(np.float32).eps, np.finfo(np.float32).eps])
+    )
+    _, scale, _ = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=False, use_fp=True, add_scale_eps=False
+    )
+    self.assertAllClose(scale, np.array([1.0, 1.0]))
+
+
+class Quantize4BTest(TestCase):
+
+  def test_quantize_symmetric(self):
+    inputs = np.array([[1.2, 3.1, 5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
+    qx, scale = quantization_util.quantize_tensor(
+        inputs, axis=[1], number_bits=4
+    )
+    self.assertAllClose(
+        qx, np.array([[2, 4, 7, 4], [0, -3, 6, 7]], dtype=np.int8)
+    )
+    self.assertAllClose(
+        scale, np.array([0.78571427, 0.5714286], dtype=np.float32)
+    )
+
+  def test_quantize_symmetric_with_dimension_size_one_unquantized(self):
+    # inputs shape: (2, 1, 4), quantization axis 2.
+    inputs = np.array([[[1.2, 3.1, 5.5, 2.9]], [[0.2, -1.5, 3.3, 4.0]]])
+    qx, scale = quantization_util.quantize_tensor(
+        inputs, axis=[2], number_bits=4
+    )
+
+    self.assertAllClose(
+        qx, np.array([[[2, 4, 7, 4]], [[0, -3, 6, 7]]], dtype=np.int8)
+    )
+    # expected scale shape: (2, 1)
+    self.assertAllClose(
+        scale, np.array([[0.78571427], [0.5714286]], dtype=np.float32)
+    )
+
+  def test_quantize_asymmetric(self):
+    inputs = np.array([[1.2, 3.1, 5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
+    qx, scale, zp = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=False, number_bits=4
+    )
+
+    self.assertAllClose(
+        qx,
+        np.array([[-8, -1, 7, -2], [-3, -8, 5, 7]], dtype=np.int8),
+    )
+    self.assertAllClose(
+        scale, np.array([0.2866667, 0.36666667], dtype=np.float32)
+    )
+    self.assertAllClose(
+        zp, np.array([-3.4933336, -1.4333334], dtype=np.float32)
+    )
+
+
+class QuantizationUtilTest(TestCase):
+
+  def test_update_to_uint4_sym(self):
+    inputs = np.array([[1.2, 3.1, -5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
+    qx, scale = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=True, number_bits=4
+    )
+    dequant_from_int4 = qx * np.expand_dims(scale, -1)
+    qx_n, scale_n, zp_n = quantization_util.update_to_uint4(qx, scale)
+    self.assertEmpty(zp_n.shape)  # A scalar numpy array.
+    dequant_from_uint4 = np.expand_dims(scale_n, -1) * (qx_n - zp_n)
+    np.testing.assert_allclose(dequant_from_int4, dequant_from_uint4)
+
+  def test_update_to_uint4_sym_combined(self):
+    inputs = np.array(
+        [[-1.2, 3.5, -6.2, 1.7], [1.2, 3.1, -5.5, 2.9], [0.2, -1.5, 3.3, 4.0]]
+    )
+    qx, scale = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=True, number_bits=4
+    )
+    dequant_from_int4 = qx * np.expand_dims(scale, -1)
+    qx_n, scale_n, zp_n = quantization_util.update_to_uint4(qx, scale)
+    self.assertEqual(zp_n.shape[0], 3)
+    dequant_from_uint4 = np.expand_dims(scale_n, -1) * (
+        qx_n - np.expand_dims(zp_n, -1)
+    )
+    np.testing.assert_allclose(dequant_from_int4, dequant_from_uint4)
+
+  def test_update_to_uint4_asym(self):
+    inputs = np.array([[1.0, 8.0, -3.0, 2.0], [-3.0, 2.0, 1.0, 8.0]])
+    qx, scale, zp = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=False, number_bits=4
+    )
+    qx_n, scale_n, zp_n = quantization_util.update_to_uint4(qx, scale, zp)
+    expected_dequant = np.array([
+        [0.0, 7.333333, -3.666667, 1.466667],
+        [-3.666667, 1.466667, 0.0, 7.333333],
+    ])
+    dequant_from_uint4 = np.expand_dims(scale_n, -1) * (qx_n - zp_n)
+    np.testing.assert_allclose(dequant_from_uint4, expected_dequant, rtol=1e-05)
+
+  def test_update_to_uint4_asym_combined(self):
+    inputs = np.array(
+        [[1.0, 8.0, -3.0, 2.0], [-3.0, 2.0, 1.0, 8.0], [2.0, 1.0, 8.0, -3.0]]
+    )
+    qx, scale, zp = quantization_util.quantize_tensor(
+        inputs, axis=[1], sym=False, number_bits=4
+    )
+    qx_n, scale_n, zp_n = quantization_util.update_to_uint4(qx, scale, zp)
+    self.assertEqual(zp_n.shape[0], 3)
+    expected_dequant = np.array([
+        [0.0, 7.333333, -3.666667, 1.466667],
+        [-3.666667, 1.466667, 0.0, 7.333333],
+        [1.466667, 0.0, 7.333333, -3.666667],
+    ])
+    dequant_from_uint4 = np.expand_dims(scale_n, -1) * (
+        qx_n - np.expand_dims(zp_n, -1)
+    )
+    np.testing.assert_allclose(dequant_from_uint4, expected_dequant, rtol=1e-05)
+
+
+class QuantizeMseTest(TestCase):
+
+  def test_mse_reduce_precision_4bit(self):
+    inputs = np.array([[-1.0, 1.0, 0.0], [1.0, 2.0, 3.0]], dtype=np.float32)
+    qx, scale = quantization_util.quantize_tensor(
+        inputs, axis=[1], number_bits=4, use_mse_quant=True
+    )
+
+    # For 4 bits, maxq = 15, zero_mse_quant = (15 - 1) / 2 = 7.
+    # The scale multiplier is 0.37755.
+
+    # For inputs row 1 ([-1.0, 1.0, 0.0]):
+    # rms = sqrt(mean((-1)^2 + 1^2 + 0^2)) = sqrt(2/3) approx 0.8164966
+    # scale = 0.37755 * 0.8164966 = 0.3079979
+    # t_shifted(-1.0) = -1.0 / 0.3079979 + 7 = -3.2461 + 7 = 3.7539
+    #   -> round 4 -> qvar = -3
+    # t_shifted(1.0) = 1.0 / 0.3079979 + 7 = 3.2461 + 7 = 10.2461
+    #   -> round 10 -> qvar = 3
+    # t_shifted(0.0) = 0.0 / 0.3079979 + 7 = 7.0 -> round 7 -> qvar = 0
+
+    # For inputs row 2 ([1.0, 2.0, 3.0]):
+    # rms = sqrt(mean(1^2 + 2^2 + 3^2)) = sqrt(14/3) approx 2.1602469
+    # scale = 0.37755 * 2.1602469 = 0.8159988
+    # t_shifted(1.0) = 1.0 / 0.8159988 + 7 = 1.2255 + 7 = 8.2255
+    #   -> round 8 -> qvar = 1
+    # t_shifted(2.0) = 2.0 / 0.8159988 + 7 = 2.4509 + 7 = 9.4509
+    #   -> round 9 -> qvar = 2
+    # t_shifted(3.0) = 3.0 / 0.8159988 + 7 = 3.6764 + 7 = 10.6764
+    #   -> round 11 -> qvar = 4
+    self.assertAllClose(qx, np.array([[-3, 3, 0], [1, 2, 4]], dtype=np.int8))
+    self.assertAllClose(scale, np.array([0.308268, 0.815601], dtype=np.float32))
+
+
+if __name__ == '__main__':
+  absltest.main()

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/safetensors_converter.py

@@ -0,0 +1,635 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CkptLoader implementation for loading the Safetensors."""
+
+import array
+from typing import Iterator
+import enum
+import glob
+import json
+import os
+from typing import List, Optional
+
+import numpy as np
+
+from mediapipe.tasks.python.genai.converter import converter_base
+from mediapipe.tasks.python.genai.converter import external_dependencies
+
+torch = external_dependencies.torch
+
+DTYPE_MAP = {
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+    "F32": torch.float32,
+}
+
+
+class _SafetensorsShardReader:
+  """Reads a single safetensors shard."""
+
+  _HEAD_BYTES = 8
+
+  def __init__(self, shard_path: str):
+    self._shard_path = shard_path
+    if not os.path.exists(self._shard_path):
+      raise ValueError(f"{self._shard_path} does not exists.")
+    with open(self._shard_path, "rb") as f:
+      head_bytes = f.read(self._HEAD_BYTES)
+      metadata_bytes_num = np.frombuffer(head_bytes, dtype=np.uint64)[0]
+      metadata_bytes = f.read(metadata_bytes_num)
+      self.layers_info = json.loads(metadata_bytes)
+      self.metadata_bytes_num = metadata_bytes_num
+
+  def read_tensor_as_numpy(self, tensor_name) -> np.ndarray:
+    """Reads a tensor from the model file as a numpy array with np.float32 type."""
+    tensor_info = self.layers_info[tensor_name]
+    with open(self._shard_path, "rb") as f:
+      shape = tensor_info["shape"]
+      dtype = tensor_info["dtype"]
+      if dtype not in DTYPE_MAP:
+        raise ValueError(f"{dtype} is not supported.")
+      data_offsets = tensor_info["data_offsets"]
+      f.seek(int(self._HEAD_BYTES + self.metadata_bytes_num + data_offsets[0]))
+      tensor_bytes = f.read(data_offsets[1] - data_offsets[0])
+      raw_tensor = torch.frombuffer(
+          array.array("b", tensor_bytes), dtype=DTYPE_MAP[dtype]
+      ).reshape(shape)
+      return raw_tensor.float().t().contiguous().numpy()
+
+  def get_tensor_names(self) -> List[str]:
+    names = list(self.layers_info.keys())
+    if "__metadata__" in names:
+      names.remove("__metadata__")
+    return names
+
+
+class _SafetensorsReader:
+  """Reads all the safetensors shards."""
+
+  def __init__(self, ckpt_path: str):
+    shards = []
+    if os.path.isdir(ckpt_path):
+      # Read all safetensors files within checkpoint
+      for shard_path in glob.glob(os.path.join(ckpt_path, "*.safetensors")):
+        shards.append(_SafetensorsShardReader(shard_path))
+    else:
+      # Assume the ckpt_path is a file or a file pattern to match.
+      for shard_path in glob.glob(ckpt_path):
+        shards.append(_SafetensorsShardReader(shard_path))
+    assert shards is not None
+
+    self._ckpt_path = ckpt_path
+    self._tensors_map = {}
+    for shard in shards:
+      tensor_names = shard.get_tensor_names()
+      for tensor_name in tensor_names:
+        if tensor_name in self._tensors_map:
+          raise ValueError(f"Duplicate tensor name: {tensor_name}")
+        self._tensors_map[tensor_name] = shard
+
+  def get_tensor_names(self) -> List[str]:
+    return list(self._tensors_map.keys())
+
+  def read_tensor_as_numpy(self, tensor_name: str) -> np.ndarray:
+    return self._tensors_map[tensor_name].read_tensor_as_numpy(tensor_name)
+
+
+class LayerType(enum.Enum):
+  """Enum for layer type."""
+
+  NONE = 0
+  ATTENTION = 1  # Layer is part of the attention module.
+  FEEDFORWARD = 2  # Layer is part of the feedforward module in the Transformer.
+  EMBEDDING = 3  # Layer is the embedding lookup or final projection layer.
+  LAYER_NORM = (
+      4  # Layer is layer normalization before and after attention layer.
+  )
+  LORA = 5  # Layer is LoRA weights augmented on the base model layers.
+  SKIP = 6  # Log a warning, and remove this not-yet-supported layer entirely.
+
+  @classmethod
+  def get_layer_type(cls, layer_name: str):
+    """Gets the layer type of the given layer name."""
+    ffn_layers = [
+        "mlp",
+    ]
+    attn_layers = [
+        "self_attn",
+    ]
+    emb_layers = [
+        "embed_tokens",
+        "embedder",
+        "lm_head",
+    ]
+    layer_norms = [
+        "input_layernorm",
+        "post_attention_layernorm",
+        "final_layernorm",
+        "model.norm.weight",
+        "pre_feedforward_layernorm",
+        "post_feedforward_layernorm",
+    ]
+    lora_layers = ["lora"]
+    skip_layers = [
+        "multi_modal_projector",
+        "vision_tower",
+    ]
+    if any(sub_name in layer_name for sub_name in skip_layers):
+      return LayerType.SKIP
+    if any(sub_name in layer_name for sub_name in lora_layers):
+      return LayerType.LORA
+    if any(sub_name in layer_name for sub_name in attn_layers):
+      return LayerType.ATTENTION
+    if any(sub_name in layer_name for sub_name in ffn_layers):
+      return LayerType.FEEDFORWARD
+    if any(sub_name in layer_name for sub_name in emb_layers):
+      return LayerType.EMBEDDING
+    if any(sub_name in layer_name for sub_name in layer_norms):
+      return LayerType.LAYER_NORM
+    else:
+      return LayerType.NONE
+
+
+class StablelmMapper(converter_base.LayerActionMapperBase):
+  """LayerActionMapper for handling the StableLM model."""
+
+  def __init__(
+      self,
+      is_symmetric: bool,
+      attention_quant_bits: int,
+      feedforward_quant_bits: int,
+      embedding_quant_bits: int,
+      backend: str,
+      reader: _SafetensorsReader,
+  ):
+    super().__init__(
+        is_symmetric=is_symmetric,
+        attention_quant_bits=attention_quant_bits,
+        feedforward_quant_bits=feedforward_quant_bits,
+        embedding_quant_bits=embedding_quant_bits,
+        backend=backend,
+    )
+    self._reader = reader
+
+  def map_to_actions(
+      self, layer_name: str
+  ) -> Optional[List[converter_base.QuantizationAction]]:
+    """Map the given layer name to actions."""
+    tensor_value = self._reader.read_tensor_as_numpy(layer_name)
+    quantize_axis = None
+    quantize_bits = None
+    layer_type = LayerType.get_layer_type(layer_name)
+
+    if layer_type != LayerType.LAYER_NORM and layer_name.endswith(".weight"):
+      quantize_axis = [0]
+      if layer_type == LayerType.FEEDFORWARD:
+        quantize_bits = self._feedforward_quant_bits
+      elif layer_type == LayerType.ATTENTION:
+        quantize_bits = self._attention_quant_bits
+        if self._backend == "cpu" and ".o_proj." in layer_name:
+          tensor_value = np.transpose(tensor_value)
+          quantize_axis = [1]
+      elif layer_type == LayerType.EMBEDDING:
+        quantize_bits = self._embedding_quant_bits
+        if self._backend == "cpu" and ".embed_tokens." in layer_name:
+          tensor_value = np.transpose(tensor_value)
+          quantize_axis = [1]
+    target_name = self.update_target_name(layer_name)
+
+    actions = [
+        converter_base.QuantizationAction(
+            tensor_name=layer_name,
+            tensor_value=tensor_value,
+            target_name=target_name,
+            quantize_axis=quantize_axis,
+            quantize_bits=quantize_bits,
+            pack_dim=0,
+        )
+    ]
+    return actions
+
+  def update_target_name(self, target_name: str) -> str:
+    """Updates the target name to match the tensor name convention."""
+    target_name = target_name.replace(
+        "model.layers.", "params.lm.transformer.x_layers_"
+    )
+    target_name = target_name.replace("mlp.up_proj", "ff_layer.ffn_layer1")
+    target_name = target_name.replace("mlp.down_proj", "ff_layer.ffn_layer2")
+    target_name = target_name.replace(
+        "mlp.gate_proj", "ff_layer.ffn_layer1_gate"
+    )
+    target_name = target_name.replace("input_layernorm", "pre_layer_norm")
+    target_name = target_name.replace(
+        "pre_layer_norm.weight", "pre_layer_norm.scale"
+    )
+    if self._backend == "cpu":
+      target_name = target_name.replace(
+          "post_attention_layernorm", "ff_layer.pre_layer_norm"
+      )
+      target_name = target_name.replace(
+          "ff_layer.pre_layer_norm.weight", "ff_layer.pre_layer_norm.scale"
+      )
+    else:
+      target_name = target_name.replace(
+          "post_attention_layernorm", "post_layer_norm"
+      )
+      target_name = target_name.replace(
+          "post_layer_norm.weight", "post_layer_norm.scale"
+      )
+    target_name = target_name.replace("self_attn.q_proj", "self_attention.q")
+    target_name = target_name.replace("self_attn.k_proj", "self_attention.k")
+    target_name = target_name.replace("self_attn.v_proj", "self_attention.v")
+    target_name = target_name.replace("self_attn.o_proj", "self_attention.post")
+    target_name = target_name.replace(
+        "model.embed_tokens", "params.lm.token_embedding"
+    )
+    target_name = target_name.replace("model.norm", "params.lm.final_ln")
+    target_name = target_name.replace("final_ln.weight", "final_ln.scale")
+    target_name = target_name.replace("lm_head", "params.lm.softmax.logits_ffn")
+    target_name = target_name.replace(".weight", ".w")
+
+    return target_name
+
+
+class PhiMapper(converter_base.LayerActionMapperBase):
+  """LayerActionMapper for handling the Phi model."""
+
+  def __init__(
+      self,
+      is_symmetric: bool,
+      attention_quant_bits: int,
+      feedforward_quant_bits: int,
+      embedding_quant_bits: int,
+      backend: str,
+      reader: _SafetensorsReader,
+  ):
+    super().__init__(
+        is_symmetric=is_symmetric,
+        attention_quant_bits=attention_quant_bits,
+        feedforward_quant_bits=feedforward_quant_bits,
+        embedding_quant_bits=embedding_quant_bits,
+        backend=backend,
+    )
+    self._reader = reader
+
+  def map_to_actions(
+      self, layer_name: str
+  ) -> Optional[List[converter_base.QuantizationAction]]:
+    """Map the given layer name to actions."""
+    tensor_value = self._reader.read_tensor_as_numpy(layer_name)
+    quantize_axis = None
+    quantize_bits = None
+    layer_type = LayerType.get_layer_type(layer_name)
+
+    if (
+        layer_type != LayerType.LAYER_NORM
+        and layer_name.endswith(".weight")
+        and layer_type != LayerType.LORA
+    ):
+      quantize_axis = [0]
+      if layer_type == LayerType.FEEDFORWARD:
+        quantize_bits = self._feedforward_quant_bits
+      elif layer_type == LayerType.ATTENTION:
+        quantize_bits = self._attention_quant_bits
+        if self._backend == "cpu" and ".dense." in layer_name:
+          tensor_value = np.transpose(tensor_value)
+          quantize_axis = [1]
+      elif layer_type == LayerType.EMBEDDING:
+        quantize_bits = self._embedding_quant_bits
+        if self._backend == "cpu" and ".embed_tokens." in layer_name:
+          tensor_value = np.transpose(tensor_value)
+          quantize_axis = [1]
+    target_name = self.update_target_name(layer_name)
+
+    actions = [
+        converter_base.QuantizationAction(
+            tensor_name=layer_name,
+            tensor_value=tensor_value,
+            target_name=target_name,
+            quantize_axis=quantize_axis,
+            quantize_bits=quantize_bits,
+            pack_dim=0,
+        )
+    ]
+    return actions
+
+  def update_target_name(self, target_name: str) -> str:
+    """Updates the target name to match the tensor name convention."""
+    target_name = target_name.replace("base_model.model.", "")
+    target_name = target_name.replace(
+        "model.layers.", "params.lm.transformer.x_layers_"
+    )
+
+    layer_type = LayerType.get_layer_type(target_name)
+    if layer_type == LayerType.FEEDFORWARD:
+      target_name = target_name.replace(".weight", ".linear.w")
+      target_name = target_name.replace(".bias", ".bias.b")
+      target_name = target_name.replace("mlp.fc1", "ff_layer.ffn_layer1")
+      target_name = target_name.replace("mlp.fc2", "ff_layer.ffn_layer2")
+
+    elif layer_type == LayerType.ATTENTION:
+      target_name = target_name.replace(".weight", ".linear.w")
+      target_name = target_name.replace(".bias", ".bias.b")
+      target_name = target_name.replace("self_attn.q_proj", "self_attention.q")
+      target_name = target_name.replace("self_attn.k_proj", "self_attention.k")
+      target_name = target_name.replace("self_attn.v_proj", "self_attention.v")
+      target_name = target_name.replace(
+          "self_attn.dense", "self_attention.post"
+      )
+    elif layer_type == LayerType.EMBEDDING:
+      target_name = target_name.replace(
+          "model.embed_tokens", "params.lm.token_embedding"
+      )
+      target_name = target_name.replace(
+          "lm_head", "params.lm.softmax.logits_ffn"
+      )
+      target_name = target_name.replace(
+          "logits_ffn.weight", "logits_ffn.linear.w"
+      )
+      target_name = target_name.replace("logits_ffn.bias", "logits_ffn.bias.b")
+    elif layer_type == LayerType.LAYER_NORM:
+      target_name = target_name.replace("input_layernorm", "pre_layer_norm")
+      target_name = target_name.replace(
+          "pre_layer_norm.weight", "pre_layer_norm.scale"
+      )
+      target_name = target_name.replace(
+          "model.final_layernorm", "params.lm.final_ln"
+      )
+      target_name = target_name.replace("final_ln.weight", "final_ln.scale")
+    target_name = target_name.replace(".weight", ".w")
+
+    # For LoRA weights
+    if "post" in target_name:
+      target_name = target_name.replace("lora_A.linear.w", "w_prime_right")
+      target_name = target_name.replace("lora_B.linear.w", "w_prime_left")
+    else:
+      target_name = target_name.replace("lora_A.linear.w", "w_prime_left")
+      target_name = target_name.replace("lora_B.linear.w", "w_prime_right")
+
+    return target_name
+
+
+class GemmaMapper(converter_base.LayerActionMapperBase):
+  """LayerActionMapper for handling the Gemma model."""
+
+  def __init__(
+      self,
+      is_symmetric: bool,
+      attention_quant_bits: int,
+      feedforward_quant_bits: int,
+      embedding_quant_bits: int,
+      backend: str,
+      reader: _SafetensorsReader,
+      is_v2: bool,
+  ):
+    super().__init__(
+        is_symmetric=is_symmetric,
+        attention_quant_bits=attention_quant_bits,
+        feedforward_quant_bits=feedforward_quant_bits,
+        embedding_quant_bits=embedding_quant_bits,
+        backend=backend,
+    )
+    self._reader = reader
+    self._is_v2 = is_v2
+
+  def map_to_actions(
+      self, layer_name: str
+  ) -> Optional[List[converter_base.QuantizationAction]]:
+    """Map the given layer name to actions."""
+
+    layer_type = LayerType.get_layer_type(layer_name)
+    if layer_type == LayerType.SKIP:
+      print("Unsupported Gemma layer: %s, skipping." % layer_name)
+      return
+
+    tensor_value = self._reader.read_tensor_as_numpy(layer_name)
+    quantize_axis = None
+    quantize_bits = None
+
+    if (
+        layer_type != LayerType.LAYER_NORM
+        and layer_name.endswith(".weight")
+        and layer_type != LayerType.LORA
+    ):
+      quantize_axis = [0]
+      if layer_type == LayerType.FEEDFORWARD:
+        quantize_bits = self._feedforward_quant_bits
+      elif layer_type == LayerType.ATTENTION:
+        quantize_bits = self._attention_quant_bits
+        if "o_proj" in layer_name:
+          tensor_value = np.transpose(tensor_value)
+          quantize_axis = [1]
+        if "q_norm" in layer_name:
+          # TODO: Look into quantization for q_norm and k_norm
+          quantize_axis = None
+          quantize_bits = None
+        if "k_norm" in layer_name:
+          quantize_axis = None
+          quantize_bits = None
+      elif layer_type == LayerType.EMBEDDING:
+        quantize_bits = self._embedding_quant_bits
+
+    target_name = self.update_target_name(layer_name)
+
+    actions = [
+        converter_base.QuantizationAction(
+            tensor_name=layer_name,
+            tensor_value=tensor_value,
+            target_name=target_name,
+            quantize_axis=quantize_axis,
+            quantize_bits=quantize_bits,
+            pack_dim=0,
+        )
+    ]
+    return actions
+
+  def update_target_name(self, target_name: str) -> str:
+    """Updates the target name to match the tensor name convention."""
+
+    # For removing multimodality stack from Gemma3-4B
+    target_name = target_name.replace("language_model.", "")
+
+    target_name = target_name.replace("base_model.model.", "")
+    target_name = target_name.replace(
+        "model.layers.", "params.lm.transformer.x_layers_"
+    )
+    target_name = target_name.replace("mlp.up_proj", "ff_layer.ffn_layer1")
+    target_name = target_name.replace("mlp.down_proj", "ff_layer.ffn_layer2")
+    target_name = target_name.replace(
+        "mlp.gate_proj", "ff_layer.ffn_layer1_gate"
+    )
+    target_name = target_name.replace("input_layernorm", "pre_layer_norm")
+    target_name = target_name.replace(
+        "pre_layer_norm.weight", "pre_layer_norm.scale"
+    )
+
+    # Gemma and Gemma2/Gemma3 differ slightly in their use of the
+    # "post_attention_layernorm" tensor name.
+    if self._is_v2:
+      target_name = target_name.replace(
+          "post_attention_layernorm", "post_layer_norm"
+      )
+    else:
+      target_name = target_name.replace(
+          "post_attention_layernorm", "ff_layer.pre_layer_norm"
+      )
+
+    target_name = target_name.replace(
+        "pre_feedforward_layernorm", "ff_layer.pre_layer_norm"
+    )
+    target_name = target_name.replace(
+        "post_feedforward_layernorm", "ff_layer.post_layer_norm"
+    )
+    target_name = target_name.replace(
+        "ff_layer.pre_layer_norm.weight", "ff_layer.pre_layer_norm.scale"
+    )
+    target_name = target_name.replace(
+        "ff_layer.post_layer_norm.weight", "ff_layer.post_layer_norm.scale"
+    )
+    target_name = target_name.replace(
+        "post_layer_norm.weight", "post_layer_norm.scale"
+    )
+    target_name = target_name.replace("self_attn.q_proj", "self_attention.q")
+    target_name = target_name.replace("self_attn.k_proj", "self_attention.k")
+    target_name = target_name.replace("self_attn.v_proj", "self_attention.v")
+    target_name = target_name.replace("self_attn.o_proj", "self_attention.post")
+
+    # Gemma3 QK norms
+    target_name = target_name.replace(
+        "self_attn.q_norm", "self_attention.q_norm"
+    )
+    target_name = target_name.replace(
+        "self_attention.q_norm.weight", "self_attention.q_norm.scale"
+    )
+    target_name = target_name.replace(
+        "self_attn.k_norm", "self_attention.k_norm"
+    )
+    target_name = target_name.replace(
+        "self_attention.k_norm.weight", "self_attention.k_norm.scale"
+    )
+
+    target_name = target_name.replace(
+        "model.embed_tokens", "params.lm.softmax.logits_ffn"
+    )
+    # Alternative Gemma3 embedding layer name
+    target_name = target_name.replace(
+        "embedder", "params.lm.softmax.logits_ffn"
+    )
+
+    target_name = target_name.replace("model.norm", "params.lm.final_ln")
+    target_name = target_name.replace("final_ln.weight", "final_ln.scale")
+    target_name = target_name.replace(".weight", ".w")
+
+    # For LoRA weights
+    if "post" in target_name:
+      target_name = target_name.replace("lora_A.w", "w_prime_right")
+      target_name = target_name.replace("lora_B.w", "w_prime_left")
+    else:
+      target_name = target_name.replace("lora_A.w", "w_prime_left")
+      target_name = target_name.replace("lora_B.w", "w_prime_right")
+
+    return target_name
+
+
+class SafetensorsCkptLoader(converter_base.CkptLoaderBase):
+  """CkptLoader implementation for loading the Safetensors."""
+
+  def __init__(
+      self,
+      ckpt_path: str,
+      is_symmetric: bool,
+      attention_quant_bits: int,
+      feedforward_quant_bits: int,
+      embedding_quant_bits: int,
+      special_model: str,
+      backend: str,
+  ):
+    """Initializes the loader.
+
+    Args:
+      ckpt_path: The filepath to the safetensors file.
+      is_symmetric: Whether to apply symmetric or asymmetric quantization.
+      attention_quant_bits: An integer that specify the target quantization bits
+        (support 8 or 4) for the attention layers.
+      feedforward_quant_bits: An integer that specify the target quantization
+        bits (support 8 or 4) for the feedforward layers in each Transformer
+        blocks.
+      embedding_quant_bits: An integer that specify the target quantization bits
+        (support 8 or 4) for the embedding (and the final projection) layers.
+      special_model: A string that indicates which input model is and whether
+        any special treatment is needed.
+      backend: A string indicating the backend used when converting this model.
+        Valid options are "cpu" and "gpu".
+    """
+    super().__init__(
+        ckpt_path,
+        is_symmetric,
+        attention_quant_bits,
+        feedforward_quant_bits,
+        embedding_quant_bits,
+    )
+
+    self._special_model = special_model
+    self._reader = _SafetensorsReader(ckpt_path)
+    if special_model in ["STABLELM_4E1T_3B"]:
+      self.mapper = StablelmMapper(
+          is_symmetric,
+          attention_quant_bits,
+          feedforward_quant_bits,
+          embedding_quant_bits,
+          backend,
+          self._reader,
+      )
+    elif special_model in ["PHI_2"]:
+      self.mapper = PhiMapper(
+          is_symmetric,
+          attention_quant_bits,
+          feedforward_quant_bits,
+          embedding_quant_bits,
+          backend,
+          self._reader,
+      )
+    elif special_model in [
+        "GEMMA_2B",
+        "GEMMA_7B",
+        "GEMMA2_2B",
+        "GEMMA3_1B",
+        "GEMMA3_4B",
+        "GEMMA3_12B",
+        "GEMMA3_27B",
+        "GEMMA3_300M",
+        "FUNCTIONGEMMA3_300M",
+    ]:
+      self.mapper = GemmaMapper(
+          is_symmetric,
+          attention_quant_bits,
+          feedforward_quant_bits,
+          embedding_quant_bits,
+          backend,
+          self._reader,
+          False if special_model in ["GEMMA_2B", "GEMMA_7B"] else True,
+      )
+    else:
+      raise ValueError(f"Unknown special model: {special_model}")
+
+  def load_to_actions(
+      self,
+  ) -> Iterator[List[converter_base.QuantizationAction]]:
+    tensor_names = self._reader.get_tensor_names()
+    for tensor_name in tensor_names:
+      tensor_actions = self.mapper.map_to_actions(tensor_name)
+      if tensor_actions is None:
+        continue
+      yield tensor_actions

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/safetensors_converter_test.py

@@ -0,0 +1,83 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for safetensors_converter."""
+
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from mediapipe.tasks.python.genai.converter import safetensors_converter
+from mediapipe.tasks.python.test import test_utils
+
+_TEST_DATA_DIR = 'mediapipe/tasks/testdata/text'
+_SAFETENSORS_FILE = test_utils.get_test_data_path(
+    os.path.join(_TEST_DATA_DIR, 'stablelm_3b_4e1t_test_weight.safetensors')
+)
+
+
+class SafetensorsConverterTest(parameterized.TestCase):
+  VARIABLE_NAMES = [
+      'model.embed_tokens.weight',
+      'model.layers.0.input_layernorm.bias',
+      'model.layers.0.input_layernorm.weight',
+      'model.layers.0.mlp.down_proj.weight',
+      'model.layers.0.mlp.gate_proj.weight',
+      'model.layers.0.mlp.up_proj.weight',
+      'model.layers.0.post_attention_layernorm.bias',
+      'model.layers.0.post_attention_layernorm.weight',
+      'model.layers.0.self_attn.k_proj.weight',
+      'model.layers.0.self_attn.o_proj.weight',
+      'model.layers.0.self_attn.q_proj.weight',
+      'model.layers.0.self_attn.v_proj.weight',
+      'model.norm.bias',
+      'model.norm.weight',
+      'lm_head.weight',
+  ]
+
+  def test_init(self):
+    loader = safetensors_converter.SafetensorsCkptLoader(
+        ckpt_path=_SAFETENSORS_FILE,
+        is_symmetric=True,
+        attention_quant_bits=8,
+        feedforward_quant_bits=8,
+        embedding_quant_bits=8,
+        special_model='STABLELM_4E1T_3B',
+        backend='gpu',
+    )
+    self.assertEqual(loader._ckpt_path, _SAFETENSORS_FILE)
+    self.assertEqual(loader._is_symmetric, True)
+    self.assertEqual(loader._attention_quant_bits, 8)
+    self.assertEqual(loader._feedforward_quant_bits, 8)
+
+  @parameterized.product(
+      quant_bits=(4, 8),
+  )
+  def test_load_to_actions(self, quant_bits):
+    loader = safetensors_converter.SafetensorsCkptLoader(
+        ckpt_path=_SAFETENSORS_FILE,
+        is_symmetric=True,
+        attention_quant_bits=8,
+        feedforward_quant_bits=quant_bits,
+        embedding_quant_bits=8,
+        special_model='STABLELM_4E1T_3B',
+        backend='gpu',
+    )
+    actions = loader.load_to_actions()
+    self.assertLen(list(actions), 15)
+
+
+if __name__ == '__main__':
+  absltest.main()

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/weight_bins_writer.py

@@ -0,0 +1,127 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""ModelWriter for writing a set of weights as binary files."""
+
+import contextlib
+import os
+from typing import Dict, Tuple
+
+import numpy as np
+
+from mediapipe.tasks.python.genai.converter import converter_base
+from mediapipe.tasks.python.genai.converter import external_dependencies
+from mediapipe.tasks.python.genai.converter import quantization_util
+
+jnp = external_dependencies.jnp
+
+
+@contextlib.contextmanager
+def filemanager(filename: str, mode: str):
+  try:
+    with open(filename, mode) as f:
+      yield f
+  finally:
+    pass
+
+
+def removeprefix(s, prefix):
+  """Removes the prefix from a string."""
+  if s.startswith(prefix):
+    return s[len(prefix) :]
+  return s
+
+
+class WeightBinsWriter(converter_base.ModelWriterBase):
+  """A ModelWriter for writing a set of weights as binary files."""
+
+  def get_weight_info(self, var_name: str, weight: np.ndarray) -> str:
+    """Gets the string that describes the weights."""
+    dtype_str = str(weight.dtype)
+    shape_str = '_'.join(map(str, weight.shape))
+    return f'mdl_vars.{var_name}.{dtype_str}.{shape_str}\n'
+
+  def write_variables(
+      self,
+      variables: Dict[str, Tuple[np.ndarray, bool]],
+      use_fake_values: bool = False,
+  ):
+    """Writes variable to the binary files. One for each layer.
+
+    Args:
+      variables: A dictionary that maps from the target variable names to the
+        quantized tensor values along with a boolean that indicates whether to
+        pack the values (only applicable for the 4-bit quantized tensors).
+      use_fake_values: Whether to use fake values for the weights. If set to
+        True, the weights will be filled with zeros.
+    """
+    weights_info = []
+    for var_name, value in variables.items():
+      output = value[0]
+      if use_fake_values:
+        output.fill(0)
+      if output.dtype == jnp.int2:
+        print('saving int2 weight: ', var_name)
+        output = output.astype(jnp.int8)
+        output = np.expand_dims(np.ravel(output), axis=-1)
+        output = quantization_util.pack_2bit_into_int8(output, 0)
+      if value[1]:
+        # Squeeze the tensor to make sure it is a 1D array for packing.
+        output = np.expand_dims(np.ravel(output), axis=-1)
+        # Extra pack needed for 4 bit. We always pack the weights along the
+        # first dimension since the tensor has already been squeezed.
+        output = quantization_util.pack_4bit(output, 0, jnp.int8)
+      if 'combined_qkv' in var_name:
+        var_name = removeprefix(var_name, 'mld_vars.')
+        var_name_q = var_name.replace('combined_qkv', 'q')
+        var_name_k = var_name.replace('combined_qkv', 'k')
+        var_name_v = var_name.replace('combined_qkv', 'v')
+        if output.shape[0] == 3:
+          weight_q, weight_k, weight_v = output
+          assert weight_q.shape == weight_k.shape == weight_v.shape
+        else:  # LoRA right weight is shared across q, k, v
+          weight_q = weight_k = weight_v = output
+          weights_info.append(self.get_weight_info(var_name_q, weight_q))
+        path_q = os.path.join(self._output_dir, var_name_q)
+        with filemanager(path_q, 'wb') as f:
+          f.write(weight_q.tobytes())
+          weights_info.append(self.get_weight_info(var_name_k, weight_k))
+        path_k = os.path.join(self._output_dir, var_name_k)
+        with filemanager(path_k, 'wb') as f:
+          f.write(weight_k.tobytes())
+        path_v = os.path.join(self._output_dir, var_name_v)
+        with filemanager(path_v, 'wb') as f:
+          f.write(weight_v.tobytes())
+          weights_info.append(self.get_weight_info(var_name_v, weight_v))
+      else:
+        if 'key' in var_name:
+          var_name = var_name.replace('key', 'k')
+        if 'query' in var_name:
+          var_name = var_name.replace('query', 'q')
+        if 'value' in var_name:
+          var_name = var_name.replace('value', 'v')
+        path = os.path.join(
+            self._output_dir, removeprefix(var_name, 'mdl_vars.')
+        )
+        with filemanager(path, 'wb') as f:
+          f.write(output.tobytes())
+          weights_info.append(self.get_weight_info(var_name, output))
+
+      # Sort weights_info
+      weights_info.sort()
+      with filemanager(
+          os.path.join(self._output_dir, 'layer_info.txt'), 'a'
+      ) as finfo:
+        for line in weights_info:
+          finfo.write(line + '\n')

backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/weight_bins_writer_test.py

@@ -0,0 +1,95 @@
+# Copyright 2024 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for pax_converter."""
+
+import os
+
+from absl import flags
+from absl.testing import absltest
+from absl.testing import parameterized
+import numpy as np
+
+from mediapipe.tasks.python.genai.converter import weight_bins_writer
+
+
+class WeightBinsWriterTest(parameterized.TestCase):
+
+  def test_get_weight_info(self):
+    output_dir = os.path.join(flags.FLAGS.test_tmpdir, 'output_dir')
+    writer = weight_bins_writer.WeightBinsWriter(
+        output_dir=output_dir, backend='cpu'
+    )
+    var_name = 'params.lm.softmax.logits_ffn.linear.w'
+    weight_info = writer.get_weight_info(
+        var_name, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
+    )
+    self.assertEqual(
+        weight_info,
+        'mdl_vars.params.lm.softmax.logits_ffn.linear.w.float32.2_3\n',
+    )
+
+  def test_load_to_actions(self):
+    output_dir = os.path.join(flags.FLAGS.test_tmpdir, 'output_dir')
+    writer = weight_bins_writer.WeightBinsWriter(
+        output_dir=output_dir, backend='cpu'
+    )
+    variables = {
+        'mdl_vars.params.lm.softmax.logits_ffn.linear.w': (
+            np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
+            False,
+        ),
+    }
+    writer.write_variables(variables)
+    file_size = os.path.getsize(
+        os.path.join(output_dir, 'params.lm.softmax.logits_ffn.linear.w')
+    )
+    self.assertEqual(file_size, 6 * 4)
+
+  @parameterized.named_parameters(
+      (
+          'real_weights',
+          np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
+          False,
+      ),
+      (
+          'fake_weights',
+          np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
+          True,
+      ),
+  )
+  def test_write_variables(self, var_values, use_fake_values):
+    output_dir = os.path.join(flags.FLAGS.test_tmpdir, 'output_dir')
+    writer = weight_bins_writer.WeightBinsWriter(
+        output_dir=output_dir, backend='gpu'
+    )
+    variables = {
+        'mdl_vars.params.lm.softmax.logits_ffn.linear.w': (
+            var_values,
+            False,
+        ),
+    }
+    writer.write_variables(variables, use_fake_values=use_fake_values)
+    with open(
+        os.path.join(output_dir, 'params.lm.softmax.logits_ffn.linear.w'), 'rb'
+    ) as f:
+      data = np.frombuffer(f.read(), dtype=np.float32).reshape(var_values.shape)
+      expected_values = var_values if not use_fake_values else np.zeros_like(
+          var_values
+      )
+      self.assertTrue(np.all(data == expected_values))
+
+
+if __name__ == '__main__':
+  absltest.main()