Merge pull request #385 from googlesamples/add-llm-bundling-example

Add llm bundling notebook.

examples/llm_inference/bundling/llm_bundling.ipynb

@@ -0,0 +1 @@
+{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"authorship_tag":"ABX9TyOM8l8EraBR86XcbcyMFOI7"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"code","execution_count":null,"metadata":{"id":"pGj9jjk0pKxs"},"outputs":[],"source":["#@title License information { display-mode: \"form\" }\n","#@markdown Copyright 2024 The MediaPipe Authors.\n","#@markdown Licensed under the Apache License, Version 2.0 (the \"License\");\n","#@markdown\n","#@markdown you may not use this file except in compliance with the License.\n","#@markdown You may obtain a copy of the License at\n","#@markdown\n","#@markdown https://www.apache.org/licenses/LICENSE-2.0\n","#@markdown\n","#@markdown Unless required by applicable law or agreed to in writing, software\n","#@markdown distributed under the License is distributed on an \"AS IS\" BASIS,\n","#@markdown WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n","#@markdown See the License for the specific language governing permissions and\n","#@markdown limitations under the License."]},{"cell_type":"code","source":["#@title Setup { display-mode: \"form\" }\n","import ipywidgets as widgets\n","from IPython.display import display\n","from google.colab import files\n","install_out = widgets.Output()\n","display(install_out)\n","with install_out:\n","  !pip install mediapipe\n","  from mediapipe.tasks.python.genai import bundler\n","\n","install_out.clear_output()\n","with install_out:\n","  print(\"Setup done.\")"],"metadata":{"id":"owxiD-Pxp26x"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Why do we need task bundles?\n","\n","Executing a text generation pipeline in its entirety requires more than merely employing the core transformer model. It requires preparing the input text to align with the model's required format, running the model autoregressively, and sampling during each iteration. Consequently, once the user has converted their model into `tflite` format (comprising the model's graph and parameters), it becomes essential to augment it with additional metadata to ensure successful end-to-end execution of the model.\n","\n","Task bundler offers a practical approach to creating such bundles. The example below illustrates how a task bundle can be generated for a converted Gemma model."],"metadata":{"id":"5Z6c-cpvqd25"}},{"cell_type":"code","source":["tflite_model=\"PATH/gemma.tflite\" # @param {type:\"string\"}\n","tokenizer_model=\"PATH/tokenizer.model\" # @param {type:\"string\"}\n","start_token=\"<bos>\" # @param {type:\"string\"}\n","stop_token=\"<eos>\" # @param {type:\"string\"}\n","output_filename=\"PATH/gemma.task\" # @param {type:\"string\"}\n","enable_bytes_to_unicode_mapping=False # @param [\"False\", \"True\"] {type:\"raw\"}\n","\n","config = bundler.BundleConfig(\n","    tflite_model=tflite_model,\n","    tokenizer_model=tokenizer_model,\n","    start_token=start_token,\n","    stop_tokens=[stop_token],\n","    output_filename=output_filename,\n","    enable_bytes_to_unicode_mapping=enable_bytes_to_unicode_mapping,\n",")\n","bundler.create_bundle(config)"],"metadata":{"id":"Cs4KJzSzqb7y"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["**Notes:**\n","\n","* The current task pipeline only supports SentencePiece tokenizer models.\n","* Certain models (e.g., phi-2) use bytes to unicode mapping. Use `enable_bytes_to_unicode_mapping` flag accordingly. Such information, including `start_token` and `stop_tokens` are often provided along with model artifacts.\n","* The generated output bundle file must end with `.task`. If not, `create_bundle` automatically adds the extension. Do not remove or change that."],"metadata":{"id":"_Unaye8xsfio"}},{"cell_type":"code","source":["#@title Download the Bundle { display-mode: \"form\" }\n","#@markdown Run this cell to download the generated `.task` file.\n","files.download(output_filename)"],"metadata":{"id":"FnFkpfr1bwtx"},"execution_count":null,"outputs":[]}]}