mediapipe-nightly 0.10.21.post20250114__cp310-cp310-manylinux_2_28_x86_64.whl

mediapipe/tasks/python/genai/converter/weight_bins_writer.py

@@ -0,0 +1,120 @@
+# 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
+
+from jax import numpy as jnp
+import numpy as np
+
+from mediapipe.tasks.python.genai.converter import converter_base
+from mediapipe.tasks.python.genai.converter import quantization_util
+
+
+@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 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')

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()

mediapipe/tasks/python/metadata/__init__.py

@@ -0,0 +1,13 @@
+# 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.