mediapipe-nightly 0.10.10.post20240216__cp38-cp38-macosx_11_0_universal2.whl → 0.10.10.post20240220__cp38-cp38-macosx_11_0_universal2.whl

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

@@ -0,0 +1,111 @@
+# 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 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]]):
+    """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).
+    """
+    weights_info = []
+    for var_name, value in variables.items():
+      output = value[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)
+      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'), 'w'
+      ) as finfo:
+        for line in weights_info:
+          finfo.write(line + '\n')

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

@@ -0,0 +1,62 @@
+# 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)
+
+
+if __name__ == '__main__':
+  absltest.main()

mediapipe/version.txt

@@ -1 +1 @@
-0.10.10-20240216
+0.10.10-20240220

{mediapipe_nightly-0.10.10.post20240216.dist-info → mediapipe_nightly-0.10.10.post20240220.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mediapipe-nightly
-Version: 0.10.10.post20240216
+Version: 0.10.10.post20240220
 Summary: MediaPipe is the simplest way for researchers and developers to build world-class ML solutions and applications for mobile, edge, cloud and the web.
 Home-page: https://github.com/google/mediapipe
 Author: The MediaPipe Authors

{mediapipe_nightly-0.10.10.post20240216.dist-info → mediapipe_nightly-0.10.10.post20240220.dist-info}/RECORD

@@ -1,11 +1,11 @@
-mediapipe_nightly-0.10.10.post20240216.dist-info/RECORD,,
-mediapipe_nightly-0.10.10.post20240216.dist-info/LICENSE,sha256=hwfu8FM5h-_FsVXWR2HutuIHk_ULm9Gmja0c9HGdDtg,12331
-mediapipe_nightly-0.10.10.post20240216.dist-info/WHEEL,sha256=ne9oAesYhmqqzZoZWVnfbfuFiLIE1hGPUf33054jNWs,109
-mediapipe_nightly-0.10.10.post20240216.dist-info/top_level.txt,sha256=LG-epD1oIiiHFRqLp--7jacjB3dbx2RfMcLYjCIhmxU,175
-mediapipe_nightly-0.10.10.post20240216.dist-info/METADATA,sha256=UOivjWZI-gIVxzgH2pfjxL15cS2JmVVJHX3pZkWGe9Y,9672
-mediapipe/__init__.py,sha256=qqZ9eB_asECNYDc7jx1XIZ-5heE3E3wuuTnMqgP6Zf8,816
+mediapipe_nightly-0.10.10.post20240220.dist-info/RECORD,,
+mediapipe_nightly-0.10.10.post20240220.dist-info/LICENSE,sha256=hwfu8FM5h-_FsVXWR2HutuIHk_ULm9Gmja0c9HGdDtg,12331
+mediapipe_nightly-0.10.10.post20240220.dist-info/WHEEL,sha256=ne9oAesYhmqqzZoZWVnfbfuFiLIE1hGPUf33054jNWs,109
+mediapipe_nightly-0.10.10.post20240220.dist-info/top_level.txt,sha256=LG-epD1oIiiHFRqLp--7jacjB3dbx2RfMcLYjCIhmxU,175
+mediapipe_nightly-0.10.10.post20240220.dist-info/METADATA,sha256=EhYSeSZKjHI0a_l13O29t73CzTLaRhFQix-W_s_lVX8,9672
+mediapipe/__init__.py,sha256=rdbukqBkiz5MY5H-FrACK0jAtVcmijeoMWNeVR_bhrU,816
 mediapipe/tasks/__init__.py,sha256=sVJS2p8J2PNVl8DLRPVY6KLpHenP_z3QVPRU0x_iL5g,571
-mediapipe/tasks/python/__init__.py,sha256=i-0yWpv_VgAcZeBFQZiLlgIb0FuFb1fHdT--0EqLCAU,838
+mediapipe/tasks/python/__init__.py,sha256=wIM_WOWboOVI1MeehN8fkN_DjoA0MEBVw5mShAd8AS4,858
 mediapipe/tasks/python/benchmark/__init__.py,sha256=epEucluzX0HinwBZoS7Tgb19j_qgfTuBf-vBkqemch8,587
 mediapipe/tasks/python/benchmark/benchmark_utils.py,sha256=5qbqGxxYlJJQLJzbW0cwMCcGl4c8ZfKL3rNmm7xMVAE,2241
 mediapipe/tasks/python/benchmark/vision/benchmark.py,sha256=gifRumSkesmXVU51GHRct-UYf8S9Dn2isD28Aw7BClQ,3491
@@ -85,6 +85,19 @@ mediapipe/tasks/python/text/text_embedder.py,sha256=JB-jQyVgcbfH8tOuJjV2Qyk2ewyU
 mediapipe/tasks/python/text/text_classifier.py,sha256=AJbYep6iL8vkf6JKRrGArr9sNd5ugLEliOK1tNr3DsE,7917
 mediapipe/tasks/python/text/core/base_text_task_api.py,sha256=OHt7j_0n5c3HBdOrCb_BGWCdKWMKvDULp6tKA5mDZAc,1822
 mediapipe/tasks/python/text/core/__init__.py,sha256=ZKC2XRtShVe6k6u6LxDt1pG7DQIn5nZnjurs6Pcvm6A,593
+mediapipe/tasks/python/genai/__init__.py,sha256=7rri6fT6wNurla8O2c5yKiLs9_3qIY0vKkyVAUDe-18,620
+mediapipe/tasks/python/genai/converter/pytorch_converter_test.py,sha256=hJOjozkWpzjedhkfWoF-vinSMs7TqCPO-W9VpleQo2U,3008
+mediapipe/tasks/python/genai/converter/safetensors_converter_test.py,sha256=bdBHCQOTqAyHpC19di_H0hV3zraBOXXfV_6NbFT0ouw,2796
+mediapipe/tasks/python/genai/converter/quantization_util.py,sha256=B6i13GqRRIwMabEJWO8rFHPMBjIgdOhFpHiwMD4GzRc,17196
+mediapipe/tasks/python/genai/converter/safetensors_converter.py,sha256=IbnvTH_n8hnVd9nChOT8oGx2Ssw4iCFCnlWB1708ujc,18710
+mediapipe/tasks/python/genai/converter/weight_bins_writer_test.py,sha256=6qgNYXODNOsbveZ0ighEW4JBdawil9mPcC16MZ0mdm8,1994
+mediapipe/tasks/python/genai/converter/converter_base.py,sha256=nRIhPYUg2pkUZcmwMIDVM4lskRRUwoOYep9MY78b2xE,6519
+mediapipe/tasks/python/genai/converter/pytorch_converter.py,sha256=RiuCr2aXbZVn3Ssa3OcSx_RQK56rfZRCpvG7qYIji4w,10675
+mediapipe/tasks/python/genai/converter/__init__.py,sha256=jfUkinDJR5BVldnbJMbo5vIr2Xc5Z4TTnaCJTNoAUvg,893
+mediapipe/tasks/python/genai/converter/converter_factory.py,sha256=2K16PZBQym0WhXM2HOdBMHMugykohoD4OTaOIo-UKko,2928
+mediapipe/tasks/python/genai/converter/llm_converter.py,sha256=_S8fj-WAA6OMQFj_QZPChlfkJ4gp_DE2t8i4pwJAHjI,7782
+mediapipe/tasks/python/genai/converter/quantization_util_test.py,sha256=ICujhTFeREGuHGmNk1PlBpf1AUThFvv-Wl5UuZ-xWAk,9060
+mediapipe/tasks/python/genai/converter/weight_bins_writer.py,sha256=Ar8XBs1GhVSFObqN29K754iVZ_WSxe4k0Rty1a4vla0,4308
 mediapipe/tasks/python/metadata/metadata.py,sha256=EECQnM-Af0angD60jaBBOuNMgt7HExH6SqVtVMFNHGc,33763
 mediapipe/tasks/python/metadata/metadata_displayer_cli.py,sha256=tLhF0B1mXG0igFTA9nPh8t1efRpRw2hQ00XpTPYdk_o,1202
 mediapipe/tasks/python/metadata/__init__.py,sha256=YGHXQMz1ZGPcNgSXggu03b0USZKE8d9Xqvn6NDUl898,586
@@ -296,7 +309,7 @@ mediapipe/python/solution_base.py,sha256=48PszV6dCssQUkyPBaqhqVehVjPtskawtGUc54k
 mediapipe/python/timestamp_test.py,sha256=oWKTZMsV586jH57OBV30rihcymETyGC29VbYURNLJQQ,2528
 mediapipe/python/image_frame_test.py,sha256=ZSjdE-an2t8i6MiA4_Xri91VMH5_CCx45fjhWUQptMY,8602
 mediapipe/python/__init__.py,sha256=xTuq4e55ofMUGUpV81ZbVzYHmjMmzt9dBOfS4ugFN1c,1428
-mediapipe/python/_framework_bindings.cpython-38-darwin.so,sha256=PsnSj2OKbw8zknVNhwab8JZsPLbLnTZczpg2y-jQVWU,78020819
+mediapipe/python/_framework_bindings.cpython-38-darwin.so,sha256=Eg1Puj-umoAIaSTLEQSbfasgJCtoy6a2ZFkhDnQkLs0,78069827
 mediapipe/python/solution_base_test.py,sha256=1u5Lo4aEUrMKj8Ha_34XMyKnI-3A1AvpaX3MCI0b2MM,15632
 mediapipe/python/packet_creator.py,sha256=34MBIMwykbZSLV-gdVYTzK8R07yzWV5yg1HuyRbS4d0,11414
 mediapipe/python/packet_getter.py,sha256=QkBxKCjXrOC6j2dJ5zcVNGaPB6zjurKztQpW5kOYLj4,4205