\n"
+ ]
},
{
- "cell_type": "markdown",
- "metadata": {
- "id": "-61KZz7EHiIS"
- },
- "source": [
- "Use the `sampler` to check if your model can perform translation. The `total_generation_steps` argument in [`gemma.sampler.Sampler`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/sampler.py#L88) is the number of steps performed when generating a response. To ensure the input matches the training format, use the prefix `Translate this into French:\\n` with a newline character at the end. This signals the model to begin translation.\n",
- "\n",
- "**Note:** Due to hardware restrictions, the number of training parameters used in the gemma Transformer may not be sufficient to produce \"stable\" results in this demo.\n",
- "\n",
- "**Note:** If you run out of memory, click on **Runtime** > **Disconnect and delete runtime**, and then **Runtime** > **Run all**.\n"
- ]
- },
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Example 0:\n",
+ "input_tokens: [[ 2 49688 736 1280 6987 235292 108 10924 665 12302\n",
+ " 235341 108 2 4397 63011 1437 38696 1241 1 0]\n",
+ " [ 2 49688 736 1280 6987 235292 108 13835 1517 235265\n",
+ " 108 2 69875 540 19713 235265 1 0 0 0]\n",
+ " [ 2 49688 736 1280 6987 235292 108 6956 1586 235297\n",
+ " 235265 108 2 78368 1586 235297 235265 1 0 0]]\n",
+ "target_mask: [[False False False False False False False False False False False False\n",
+ " True True True True True True True False]\n",
+ " [False False False False False False False False False False False True\n",
+ " True True True True True False False False]\n",
+ " [False False False False False False False False False False False False\n",
+ " True True True True True True False False]]\n",
+ "\n",
+ "Example 1:\n",
+ "input_tokens: [[ 2 49688 736 1280 6987 235292 108 18874 235341 108\n",
+ " 2 115905 6425 1241 1 0 0 0 0 0]\n",
+ " [ 2 49688 736 1280 6987 235292 108 7574 3356 235341\n",
+ " 108 2 7997 20707 1241 1 0 0 0 0]\n",
+ " [ 2 49688 736 1280 6987 235292 108 8703 665 235265\n",
+ " 108 2 235338 235303 90006 20133 235265 1 0 0]]\n",
+ "target_mask: [[False False False False False False False False False False True True\n",
+ " True True True False False False False False]\n",
+ " [False False False False False False False False False False False True\n",
+ " True True True True False False False False]\n",
+ " [False False False False False False False False False False False True\n",
+ " True True True True True True False False]]\n",
+ "\n"
+ ]
+ }
+ ],
+ "source": [
+ "tokenizer = GemmaTokenizer(vocab)\n",
+ "\n",
+ "dataset_builder = MTNTDatasetBuilder(tokenizer, max_seq_len=20)\n",
+ "ds = dataset_builder.get_train_dataset(3, 1)\n",
+ "ds = ds.take(2)\n",
+ "ds = ds.as_numpy_iterator()\n",
+ "\n",
+ "for idx, example in enumerate(ds):\n",
+ " print(f'Example {idx}:')\n",
+ " for key, val in example.items():\n",
+ " print(f'{key}: {val}')\n",
+ " print()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "7IY8Muu1zRF4"
+ },
+ "source": [
+ "## Configure the model\n",
+ "\n",
+ "Before you begin fine-tuning the Gemma model, you need to configure it.\n",
+ "\n",
+ "First, load and format the Gemma model checkpoint with the [`gemma.params.load_and_format_params`](https://github.com/google-deepmind/gemma/blob/c6bd156c246530e1620a7c62de98542a377e3934/gemma/params.py#L27) method:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "metadata": {
+ "id": "by6eWKtqzxRf"
+ },
+ "outputs": [],
+ "source": [
+ "params = params_lib.load_and_format_params(CKPT_PATH)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "BtJhJkkZzsy1"
+ },
+ "source": [
+ "To automatically load the correct configuration from the Gemma model checkpoint, use [`gemma.transformer.TransformerConfig`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L65). The `cache_size` argument is the number of time steps in the Gemma `Transformer` cache. Afterwards, instantiate the Gemma model as `model_2b` with [`gemma.transformer.Transformer`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L136) (which inherits from [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html)).\n",
+ "\n",
+ "**Note:** The vocabulary size is smaller than the number of input embeddings because of unused tokens in the current Gemma release."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 15,
+ "metadata": {
+ "id": "_jjlFAkazzit"
+ },
+ "outputs": [],
+ "source": [
+ "config_2b = transformer_lib.TransformerConfig.from_params(\n",
+ " params,\n",
+ " cache_size=30\n",
+ ")\n",
+ "\n",
+ "model_2b = transformer_lib.Transformer(config=config_2b)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "t7UL2Af536x_"
+ },
+ "source": [
+ "## Fine-tune the model\n",
+ "\n",
+ "In this section, you will:\n",
+ "\n",
+ "- Use the `gemma.transformer.Transformer` class to create the forward pass and loss function.\n",
+ "- Build the position and attention mask vectors for tokens\n",
+ "- Build a training step function with Flax.\n",
+ "- Build the validation step without the backwards pass.\n",
+ "- Create the training loop.\n",
+ "- Fine-tune the Gemma model."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aJhtJumH7H8_"
+ },
+ "source": [
+ "Define the forward pass and the loss function using the [`gemma.transformer.Transformer`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L136) class. The Gemma `Transformer` inherits from [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html), and offers two essential methods:\n",
+ "\n",
+ "- `init`: Initializes the model's parameters.\n",
+ "- `apply`: Executes the model's `__call__` function using a given set of parameters.\n",
+ "\n",
+ " Since you are working with pre-trained Gemma weights, you don't need to use the `init` function."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 16,
+ "metadata": {
+ "id": "iEcV0XEEEcoZ"
+ },
+ "outputs": [],
+ "source": [
+ "def forward_and_loss_fn(params,\n",
+ " *,\n",
+ " model: transformer_lib.Transformer,\n",
+ " input_tokens: jax.Array, # Shape [B, L]\n",
+ " input_mask: jax.Array, # Shape [B, L]\n",
+ " positions: jax.Array, # Shape [B, L]\n",
+ " attention_mask: jax.Array, # [B, L, L]\n",
+ " ) -> jax.Array:\n",
+ " \"\"\"The forward pass and the loss function.\n",
+ "\n",
+ " Args:\n",
+ " params: Model's input parameters.\n",
+ " model: The Gemma transformer model to call.\n",
+ " input_tokens: Input tokens sequence, shape [B, L].\n",
+ " input_mask: Tokens to ignore when computing the loss, shape [B, L].\n",
+ " positions: Relative position of each token, shape [B, L].\n",
+ " attention_mask: Input attention mask, shape [B, L].\n",
+ "\n",
+ " Returns:\n",
+ " The softmax cross-entropy loss for the next-token prediction task.\n",
+ " \"\"\"\n",
+ "\n",
+ " # The forward pass on the input data.\n",
+ " # No attention cache is needed here.\n",
+ " logits, _ = model.apply(\n",
+ " params,\n",
+ " input_tokens,\n",
+ " positions,\n",
+ " None, # Attention cache is None.\n",
+ " attention_mask,\n",
+ " )\n",
+ "\n",
+ " # Exclude the last step as it does not appear in the targets.\n",
+ " logits = logits[0, :-1]\n",
+ "\n",
+ " # Similarly, the first token cannot be predicted.\n",
+ " target_tokens = input_tokens[0, 1:]\n",
+ " target_mask = input_mask[0, 1:]\n",
+ "\n",
+ " # Convert the target labels to one-hot encoded vectors.\n",
+ " one_hot = jax.nn.one_hot(target_tokens, logits.shape[-1])\n",
+ "\n",
+ " # Don't update on unwanted tokens.\n",
+ " one_hot = one_hot * target_mask.astype(one_hot.dtype)[...,None]\n",
+ "\n",
+ " # Define the normalization factor.\n",
+ " norm_factor = 1 / (jnp.sum(target_mask) + 1e-8)\n",
+ "\n",
+ " # Return the negative log likelihood (NLL) loss.\n",
+ " return -jnp.sum(jax.nn.log_softmax(logits) * one_hot) * norm_factor"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "WxbxsKcd7Ot7"
+ },
+ "source": [
+ "The [`gemma.transformer.Transformer`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L136) class requires an `attention_mask` and a `positions` vector alongside each input. You can generate these by creating a custom function that uses [`Transformer.build_positions_from_mask`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L48) and [`Transformer.make_causal_attn_mask`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L29):\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 17,
+ "metadata": {
+ "id": "cbWfdHf0EcoZ"
+ },
+ "outputs": [],
+ "source": [
+ "def get_attention_mask_and_positions(example: jax.Array,\n",
+ " pad_id : int,\n",
+ " )-> tuple[jax.Array, jax.Array]:\n",
+ " \"\"\"Builds the position and attention mask vectors from the given tokens.\"\"\"\n",
+ " pad_mask = example != pad_id\n",
+ " current_token_position = transformer_lib.build_positions_from_mask(pad_mask)\n",
+ " attention_mask = transformer_lib.make_causal_attn_mask(pad_mask)\n",
+ " return current_token_position, attention_mask"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "uRkeF6ed8tOI"
+ },
+ "source": [
+ "Build the `train_step` function that performs the backward pass and updates the model's parameters accordingly, where:\n",
+ "\n",
+ "- [`jax.value_and_grad`](https://jax.readthedocs.io/en/latest/_autosummary/jax.value_and_grad.html) is for evaluating the loss function and gradients during the forward and backward passes.\n",
+ "- [`optax.apply_updates`](https://optax.readthedocs.io/en/latest/api/apply_updates.html#optax.apply_updates) is for updating the parameters."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 18,
+ "metadata": {
+ "id": "cPSfp7ZUEcoZ"
+ },
+ "outputs": [],
+ "source": [
+ "def train_step(model: transformer_lib.Transformer,\n",
+ " params,\n",
+ " optimizer: optax.GradientTransformation,\n",
+ " opt_state: optax.OptState,\n",
+ " pad_id: int,\n",
+ " example: TrainingInput):\n",
+ " \"\"\"Train step.\n",
+ "\n",
+ " Args:\n",
+ " model: The Gemma transformer model.\n",
+ " params: The model's input parameters.\n",
+ " optimizer: The Optax optimizer to use.\n",
+ " opt_state: The input optimizer's state.\n",
+ " pad_id: ID of the pad token.\n",
+ " example: Input batch.\n",
+ "\n",
+ " Returns:\n",
+ " The training loss, the updated parameters, and the updated optimizer state.\n",
+ " \"\"\"\n",
+ "\n",
+ " # Build the position and attention mask vectors.\n",
+ " positions, attention_mask = get_attention_mask_and_positions(example.input_tokens, pad_id)\n",
+ "\n",
+ " # The forward and backward passes.\n",
+ " train_loss, grads = jax.value_and_grad(forward_and_loss_fn)(params,\n",
+ " model=model,\n",
+ " input_tokens=example.input_tokens,\n",
+ " input_mask=example.target_mask,\n",
+ " positions=positions,\n",
+ " attention_mask=attention_mask)\n",
+ " # Update the parameters.\n",
+ " updates, opt_state = optimizer.update(grads, opt_state)\n",
+ " params = optax.apply_updates(params, updates)\n",
+ "\n",
+ " return train_loss, params, opt_state"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "8ZKSa-jJ809n"
+ },
+ "source": [
+ "Build the `validation_step` function without the backward pass:\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 19,
+ "metadata": {
+ "id": "yU4oR92YEcoa"
+ },
+ "outputs": [],
+ "source": [
+ "def validation_step(model: transformer_lib.Transformer,\n",
+ " params,\n",
+ " pad_id: int,\n",
+ " example: TrainingInput,\n",
+ " ):\n",
+ " positions, attention_mask = get_attention_mask_and_positions(example.input_tokens, pad_id)\n",
+ " val_loss = forward_and_loss_fn(params,\n",
+ " model=model,\n",
+ " input_tokens=example.input_tokens,\n",
+ " input_mask=example.target_mask,\n",
+ " positions=positions,\n",
+ " attention_mask=attention_mask)\n",
+ " return val_loss"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "bNqVhj7v87f4"
+ },
+ "source": [
+ "Define the training loop using [`optax.sgd`](https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.sgd) for the SGD optimizer:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 20,
+ "metadata": {
+ "id": "xT4bAqNLEcoa"
+ },
+ "outputs": [],
+ "source": [
+ "@chex.dataclass(frozen=True)\n",
+ "class TrainingConfig:\n",
+ " learning_rate: float\n",
+ " num_epochs: int\n",
+ " eval_every_n: int\n",
+ " batch_size: int\n",
+ " max_steps: int | None = None\n",
+ "\n",
+ "def train_loop(\n",
+ " model: transformer_lib.Transformer,\n",
+ " params,\n",
+ " dataset_builder: MTNTDatasetBuilder,\n",
+ " training_cfg: TrainingConfig):\n",
+ "\n",
+ " # Apply `jax.jit` on the training step, making the whole loop much more efficient.\n",
+ " compiled_train_step = jax.jit(train_step, static_argnames=['model', 'optimizer'])\n",
+ "\n",
+ " # Apply `jax.jit` on the validation step.\n",
+ " compiled_validation_step = jax.jit(validation_step, static_argnames=['model'])\n",
+ "\n",
+ " # To save memory, use the SGD optimizer instead of the usual Adam optimizer.\n",
+ " # Note that for this specific example, SGD is more than enough.\n",
+ " optimizer = optax.sgd(training_cfg.learning_rate)\n",
+ " opt_state = optimizer.init(params)\n",
+ "\n",
+ " # Build the training dataset.\n",
+ " train_ds = dataset_builder.get_train_dataset(batch_size=training_cfg.batch_size,\n",
+ " num_epochs=training_cfg.num_epochs)\n",
+ " train_ds = train_ds.as_numpy_iterator()\n",
+ "\n",
+ " # Build the validation dataset, with a limited number of samples for this demo.\n",
+ " validation_ds = dataset_builder.get_validation_dataset(batch_size=training_cfg.batch_size)\n",
+ " validation_ds = validation_ds.take(50)\n",
+ "\n",
+ " n_steps = 0\n",
+ " avg_loss=0\n",
+ "\n",
+ " # A first round of the validation loss.\n",
+ " n_steps_eval = 0\n",
+ " eval_loss = 0\n",
+ " val_iterator = validation_ds.as_numpy_iterator()\n",
+ " for val_example in val_iterator:\n",
+ " eval_loss += compiled_validation_step(model,\n",
+ " params,\n",
+ " dataset_builder._tokenizer.pad_id,\n",
+ " val_example)\n",
+ " n_steps_eval += 1\n",
+ " print(f\"Start, validation loss: {eval_loss/n_steps_eval}\")\n",
+ "\n",
+ " for train_example in train_ds:\n",
+ " train_loss, params, opt_state = compiled_train_step(model=model,\n",
+ " params=params,\n",
+ " optimizer=optimizer,\n",
+ " opt_state=opt_state,\n",
+ " pad_id=dataset_builder._tokenizer.pad_id,\n",
+ " example=train_example)\n",
+ " n_steps += 1\n",
+ " avg_loss += train_loss\n",
+ " if n_steps % training_cfg.eval_every_n == 0:\n",
+ " eval_loss = 0\n",
+ "\n",
+ " n_steps_eval = 0\n",
+ " val_iterator = validation_ds.as_numpy_iterator()\n",
+ " for val_example in val_iterator:\n",
+ " eval_loss += compiled_validation_step(model,\n",
+ " params,\n",
+ " dataset_builder._tokenizer.pad_id,\n",
+ " val_example)\n",
+ " n_steps_eval +=1\n",
+ " avg_loss /= training_cfg.eval_every_n\n",
+ " eval_loss /= n_steps_eval\n",
+ " print(f\"STEP {n_steps} training loss: {avg_loss} - eval loss: {eval_loss}\")\n",
+ " avg_loss=0\n",
+ " if training_cfg.max_steps is not None and n_steps > training_cfg.max_steps:\n",
+ " break\n",
+ " return params"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ecv6lp5MCzFc"
+ },
+ "source": [
+ "Begin fine-tuning the Gemma model on a limited number of steps (`SEQ_SIZE`) to make sure this fits in the memory:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 21,
+ "metadata": {
+ "id": "7SL2VAmVEcoa"
+ },
+ "outputs": [
{
- "cell_type": "code",
- "execution_count": 23,
- "metadata": {
- "id": "S5F3fk22Ecod"
- },
- "outputs": [
- {
- "data": {
- "text/plain": [
- "[\"C'est Bonjour, mon nom est Morgane.C'est Bonjour, mon nom est Morgane.\"]"
- ]
- },
- "execution_count": 23,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "sampler(\n",
- " [\"Translate this into French:\\nHello, my name is Morgane.\\n\"],\n",
- " total_generation_steps=100,\n",
- " ).text"
- ]
- },
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Start, validation loss: 10.647212982177734\n",
+ "STEP 20 training loss: 3.3015992641448975 - eval loss: 2.686880111694336\n",
+ "STEP 40 training loss: 5.375057220458984 - eval loss: 2.6751961708068848\n",
+ "STEP 60 training loss: 2.6599338054656982 - eval loss: 2.663877010345459\n",
+ "STEP 80 training loss: 4.822389125823975 - eval loss: 2.3333375453948975\n",
+ "STEP 100 training loss: 2.0131142139434814 - eval loss: 2.360811948776245\n"
+ ]
+ }
+ ],
+ "source": [
+ "SEQ_SIZE = 25\n",
+ "tokenizer = GemmaTokenizer(vocab)\n",
+ "dataset_builder= MTNTDatasetBuilder(tokenizer, SEQ_SIZE)\n",
+ "training_cfg = TrainingConfig(learning_rate=1e-4,\n",
+ " num_epochs=1,\n",
+ " eval_every_n=20,\n",
+ " batch_size=1,\n",
+ " max_steps=100)\n",
+ "\n",
+ "params = train_loop(model=model_2b,\n",
+ " params={'params': params['transformer']},\n",
+ " dataset_builder=dataset_builder,\n",
+ " training_cfg=training_cfg)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "EtfVo3pDDAZV"
+ },
+ "source": [
+ "Both the training loss and the validation loss should have gone down with each step count.\n",
+ "\n",
+ "Create a `sampler` with [`gemma.sampler.Sampler`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/sampler.py#L88). It uses the Gemma model checkpoint and the tokenizer."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 22,
+ "metadata": {
+ "id": "dQ1oCF10Ecod"
+ },
+ "outputs": [],
+ "source": [
+ "sampler = sampler_lib.Sampler(\n",
+ " transformer=model_2b,\n",
+ " vocab=vocab,\n",
+ " params=params['params'],\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "-61KZz7EHiIS"
+ },
+ "source": [
+ "Use the `sampler` to check if your model can perform translation. The `total_generation_steps` argument in [`gemma.sampler.Sampler`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/sampler.py#L88) is the number of steps performed when generating a response. To ensure the input matches the training format, use the prefix `Translate this into French:\\n` with a newline character at the end. This signals the model to begin translation.\n",
+ "\n",
+ "**Note:** Due to hardware restrictions, the number of training parameters used in the gemma Transformer may not be sufficient to produce \"stable\" results in this demo.\n",
+ "\n",
+ "**Note:** If you run out of memory, click on **Runtime** > **Disconnect and delete runtime**, and then **Runtime** > **Run all**.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 23,
+ "metadata": {
+ "id": "S5F3fk22Ecod"
+ },
+ "outputs": [
{
- "cell_type": "markdown",
- "metadata": {
- "id": "Jao0Qk-ZIqyD"
- },
- "source": [
- "## Learn more\n",
- "\n",
- "- You can learn more about the Google DeepMind [`gemma` library on GitHub](https://github.com/google-deepmind/gemma), which contains docstrings of methods you used in this tutorial, such as [`gemma.params`](https://github.com/google-deepmind/gemma/blob/main/gemma/params.py),\n",
- "[`gemma.transformer`](https://github.com/google-deepmind/gemma/blob/main/gemma/transformer.py), and\n",
- "[`gemma.sampler`](https://github.com/google-deepmind/gemma/blob/main/gemma/sampler.py).\n",
- "- The following libraries have their own documentation sites: [core JAX](https://jax.readthedocs.io), [Flax](https://flax.readthedocs.io), [Chex](https://chex.readthedocs.io/en/latest/), [Optax](https://optax.readthedocs.io/en/latest/), and [Orbax](https://orbax.readthedocs.io/).\n",
- "- For `sentencepiece` tokenizer/detokenizer documentation, check out [Google's `sentencepiece` GitHub repo](https://github.com/google/sentencepiece).\n",
- "- For `kagglehub` documentation, check out `README.md` on [Kaggle's `kagglehub` GitHub repo](https://github.com/Kaggle/kagglehub).\n",
- "- Learn how to [use Gemma models with Google Cloud Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/open-models/use-gemma).\n",
- "- If you are using Google Cloud TPUs (v3-8 and newer), make sure to also update to the latest `jax[tpu]` package (`!pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html`), restart the runtime, and check that `jax` and `jaxlib` versions match (`!pip list | grep jax`). This can prevent the `RuntimeError` that can arise because of the `jaxlib` and `jax` version mismatch. For more JAX installation instructions, refer to the [JAX docs](https://jax.readthedocs.io/en/latest/tutorials/installation.html#install-google-tpu)."
+ "data": {
+ "text/plain": [
+ "[\"C'est Bonjour, mon nom est Morgane.C'est Bonjour, mon nom est Morgane.\"]"
]
+ },
+ "execution_count": 23,
+ "metadata": {},
+ "output_type": "execute_result"
}
- ],
- "metadata": {
- "accelerator": "GPU",
- "colab": {
- "name": "jax_finetune.ipynb",
- "toc_visible": true
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- }
+ ],
+ "source": [
+ "sampler(\n",
+ " [\"Translate this into French:\\nHello, my name is Morgane.\\n\"],\n",
+ " total_generation_steps=100,\n",
+ " ).text"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Jao0Qk-ZIqyD"
+ },
+ "source": [
+ "## Learn more\n",
+ "\n",
+ "- You can learn more about the Google DeepMind [`gemma` library on GitHub](https://github.com/google-deepmind/gemma), which contains docstrings of modules you used in this tutorial, such as [`gemma.params`](https://github.com/google-deepmind/gemma/blob/main/gemma/params.py),\n",
+ "[`gemma.transformer`](https://github.com/google-deepmind/gemma/blob/main/gemma/transformer.py), and\n",
+ "[`gemma.sampler`](https://github.com/google-deepmind/gemma/blob/main/gemma/sampler.py).\n",
+ "- The following libraries have their own documentation sites: [core JAX](https://jax.readthedocs.io), [Flax](https://flax.readthedocs.io), [Chex](https://chex.readthedocs.io/en/latest/), [Optax](https://optax.readthedocs.io/en/latest/), and [Orbax](https://orbax.readthedocs.io/).\n",
+ "- For `sentencepiece` tokenizer/detokenizer documentation, check out [Google's `sentencepiece` GitHub repo](https://github.com/google/sentencepiece).\n",
+ "- For `kagglehub` documentation, check out `README.md` on [Kaggle's `kagglehub` GitHub repo](https://github.com/Kaggle/kagglehub).\n",
+ "- Learn how to [use Gemma models with Google Cloud Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/open-models/use-gemma).\n",
+ "- If you are using Google Cloud TPUs (v3-8 and newer), make sure to also update to the latest `jax[tpu]` package (`!pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html`), restart the runtime, and check that `jax` and `jaxlib` versions match (`!pip list | grep jax`). This can prevent the `RuntimeError` that can arise because of the `jaxlib` and `jax` version mismatch. For more JAX installation instructions, refer to the [JAX docs](https://jax.readthedocs.io/en/latest/tutorials/installation.html#install-google-tpu)."
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "name": "jax_finetune.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
},
- "nbformat": 4,
- "nbformat_minor": 0
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.11.7"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
}
diff --git a/site/en/gemma/docs/jax_inference.ipynb b/site/en/gemma/docs/jax_inference.ipynb
index a7d5b82f3..899d4f58b 100644
--- a/site/en/gemma/docs/jax_inference.ipynb
+++ b/site/en/gemma/docs/jax_inference.ipynb
@@ -1,500 +1,513 @@
{
- "cells": [
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "Tce3stUlHN0L"
- },
- "source": [
- "##### Copyright 2024 Google LLC."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "tuOe1ymfHZPu"
- },
- "outputs": [],
- "source": [
- "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
- "# you may not use this file except in compliance with the License.\n",
- "# You may obtain a copy of the License at\n",
- "#\n",
- "# https://www.apache.org/licenses/LICENSE-2.0\n",
- "#\n",
- "# Unless required by applicable law or agreed to in writing, software\n",
- "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
- "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
- "# See the License for the specific language governing permissions and\n",
- "# limitations under the License."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "FUOiKRSF7jc1"
- },
- "source": [
- "# Inference with Gemma using JAX and Flax"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "60KmTK7o6ppd"
- },
- "source": [
- ""
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "Tdlq6K0znh3O"
- },
- "source": [
- "## Overview\n",
- "\n",
- "Gemma is a family of lightweight, state-of-the-art open large language models, based on the Google DeepMind Gemini research and technology. This tutorial demonstrates how to perform basic sampling/inference with the Gemma 2B Instruct model using [Google DeepMind's `gemma` library](https://github.com/google-deepmind/gemma) that was written with [JAX](https://jax.readthedocs.io) (a high-performance numerical computing library), [Flax](https://flax.readthedocs.io) (the JAX-based neural network library), [Orbax](https://orbax.readthedocs.io/) (a JAX-based library for training utilities like checkpointing), and [SentencePiece](https://github.com/google/sentencepiece) (a tokenizer/detokenizer library). Although Flax is not used directly in this notebook, Flax was used to create Gemma.\n",
- "\n",
- "This notebook can run on Google Colab with free T4 GPU (go to **Edit** > **Notebook settings** > Under **Hardware accelerator** select **T4 GPU**)."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "aKvTsIkL98BG"
- },
- "source": [
- "## Setup"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "WCgCkmQSPxkE"
- },
- "source": [
- "### 1. Set up Kaggle access for Gemma\n",
- "\n",
- "To complete this tutorial, you first need to follow the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup), which show you how to do the following:\n",
- "\n",
- "* Get access to Gemma on [kaggle.com](https://www.kaggle.com/models/google/gemma/).\n",
- "* Select a Colab runtime with sufficient resources to run the Gemma model.\n",
- "* Generate and configure a Kaggle username and API key.\n",
- "\n",
- "After you've completed the Gemma setup, move on to the next section, where you'll set environment variables for your Colab environment.\n",
- "\n",
- "### Set environment variables\n",
- "\n",
- "Set environment variables for `KAGGLE_USERNAME` and `KAGGLE_KEY`. When prompted with the \"Grant access?\" messages, agree to provide secret access."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 1,
- "metadata": {
- "id": "lKoW-nhE-gNO"
- },
- "outputs": [],
- "source": [
- "import os\n",
- "from google.colab import userdata # `userdata` is a Colab API.\n",
- "\n",
- "os.environ[\"KAGGLE_USERNAME\"] = userdata.get('KAGGLE_USERNAME')\n",
- "os.environ[\"KAGGLE_KEY\"] = userdata.get('KAGGLE_KEY')"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "AO7a1Q4Yyc9Z"
- },
- "source": [
- "### 2. Install the `gemma` library\n",
- "\n",
- "This notebook focuses on using a free Colab GPU. To enable hardware acceleration, click on **Edit** > **Notebook settings** > Select **T4 GPU** > **Save**.\n",
- "\n",
- "Next, you need to install the Google DeepMind `gemma` library from [`github.com/google-deepmind/gemma`](https://github.com/google-deepmind/gemma). If you get an error about \"pip's dependency resolver\", you can usually ignore it.\n",
- "\n",
- "**Note:** By installing `gemma`, you will also install [`flax`](https://flax.readthedocs.io), core [`jax`](https://jax.readthedocs.io), [`optax`](https://optax.readthedocs.io/en/latest/) (the JAX-based gradient processing and optimization library), [`orbax`](https://orbax.readthedocs.io/), and [`sentencepiece`](https://github.com/google/sentencepiece)."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 2,
- "metadata": {
- "id": "WWEzVJR4Fx9g"
- },
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- " Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
- " Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
- " Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
- "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m133.7/133.7 kB\u001b[0m \u001b[31m3.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
- "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m244.4/244.4 kB\u001b[0m \u001b[31m13.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
- "\u001b[?25h Building wheel for gemma (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
- "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
- "tensorflow-metadata 1.14.0 requires absl-py<2.0.0,>=0.9, but you have absl-py 2.1.0 which is incompatible.\u001b[0m\u001b[31m\n",
- "\u001b[0m"
- ]
- }
- ],
- "source": [
- "!pip install -q git+https://github.com/google-deepmind/gemma.git"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "VKLjBAe1m3Ck"
- },
- "source": [
- "## Load and prepare the Gemma model\n",
- "\n",
- "1. Load the Gemma model with [`kagglehub.model_download`](https://github.com/Kaggle/kagglehub/blob/bddefc718182282882b72f814d407d89e5d178c4/src/kagglehub/models.py#L12), which takes three arguments:\n",
- "\n",
- "- `handle`: The model handle from Kaggle\n",
- "- `path`: (Optional string) The local path\n",
- "- `force_download`: (Optional boolean) Forces to re-download the model\n",
- "\n",
- "**Note:** Be mindful that the `gemma-2b-it` model is around 3.7Gb in size."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 3,
- "metadata": {
- "id": "_W3FUd9lt8VT"
- },
- "outputs": [],
- "source": [
- "GEMMA_VARIANT = '2b-it' # @param ['2b', '2b-it'] {type:\"string\"}"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 4,
- "metadata": {
- "id": "kFCmWEKdMA_Y"
- },
- "outputs": [
- {
- "name": "stderr",
- "output_type": "stream",
- "text": [
- "Downloading from https://www.kaggle.com/api/v1/models/google/gemma/flax/2b-it/2/download...\n",
- "100%|██████████| 3.67G/3.67G [00:35<00:00, 110MB/s]\n",
- "Extracting model files...\n"
- ]
- }
- ],
- "source": [
- "import kagglehub\n",
- "\n",
- "GEMMA_PATH = kagglehub.model_download(f'google/gemma/flax/{GEMMA_VARIANT}')"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 5,
- "metadata": {
- "id": "nYmYTMk8aELi"
- },
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "GEMMA_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2\n"
- ]
- }
- ],
- "source": [
- "print('GEMMA_PATH:', GEMMA_PATH)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ytNi47xSlw71"
- },
- "source": [
- "**Note:** The path from the output above is where the model weights and tokenizer are saved locally, you will need them for later."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "92BcvYdemXbd"
- },
- "source": [
- "2. Check the location of the model weights and the tokenizer, then set the path variables. The tokenizer directory will be in the main directory where you downloaded the model, while the model weights will be in a sub-directory. For example:\n",
- "\n",
- "- The `tokenizer.model` file will be in `/LOCAL/PATH/TO/gemma/flax/2b-it/2`).\n",
- "- The model checkpoint will be in `/LOCAL/PATH/TO/gemma/flax/2b-it/2/2b-it`)."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 6,
- "metadata": {
- "id": "QY6OnASOpZbW"
- },
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "CKPT_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/2b-it\n",
- "TOKENIZER_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/tokenizer.model\n"
- ]
- }
- ],
- "source": [
- "CKPT_PATH = os.path.join(GEMMA_PATH, GEMMA_VARIANT)\n",
- "TOKENIZER_PATH = os.path.join(GEMMA_PATH, 'tokenizer.model')\n",
- "print('CKPT_PATH:', CKPT_PATH)\n",
- "print('TOKENIZER_PATH:', TOKENIZER_PATH)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "jc0ZzYIW0TSN"
- },
- "source": [
- "## Perform sampling/inference"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "aEe3p8geqekV"
- },
- "source": [
- "1. Load and format the Gemma model checkpoint with the [`gemma.params`](https://github.com/google-deepmind/gemma/blob/main/gemma/params.py) method:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 7,
- "metadata": {
- "id": "Mnr52JQVqKRw"
- },
- "outputs": [],
- "source": [
- "from gemma import params as params_lib\n",
- "\n",
- "params = params_lib.load_and_format_params(CKPT_PATH)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "-Xpnb2igrGjk"
- },
- "source": [
- "2. Load the Gemma tokenizer, constructed using [`sentencepiece.SentencePieceProcessor`](https://github.com/google/sentencepiece/blob/4d6a1f41069c4636c51a5590f7578a0dbed83450/python/src/sentencepiece/__init__.py#L423):"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 8,
- "metadata": {
- "id": "-T0ZHff5rNSy"
- },
- "outputs": [
- {
- "data": {
- "text/plain": [
- "True"
- ]
- },
- "execution_count": 8,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "import sentencepiece as spm\n",
- "\n",
- "vocab = spm.SentencePieceProcessor()\n",
- "vocab.Load(TOKENIZER_PATH)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "IkAf4fkNrY-3"
- },
- "source": [
- "3. To automatically load the correct configuration from the Gemma model checkpoint, use [`gemma.transformer.TransformerConfig`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L65). The `cache_size` argument is the number of time steps in the Gemma `transformer` cache. Afterwards, instantiate the Gemma model as `transformer` with [`gemma.transformer.Transformer`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L136) (which inherits from [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html)).\n",
- "\n",
- "**Note:** The vocabulary size is smaller than the number of input embeddings because of unused tokens in the current Gemma release."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 9,
- "metadata": {
- "id": "4PNWxDhvrRXJ"
- },
- "outputs": [],
- "source": [
- "from gemma import transformer as transformer_lib\n",
- "\n",
- "transformer_config = transformer_lib.TransformerConfig.from_params(\n",
- " params=params,\n",
- " cache_size=1024\n",
- ")\n",
- "\n",
- "transformer = transformer_lib.Transformer(transformer_config)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "vs0vgmXVroBq"
- },
- "source": [
- "3. Create a `sampler` with [`gemma.sampler.Sampler`](https://github.com/google-deepmind/gemma/blob/c6bd156c246530e1620a7c62de98542a377e3934/gemma/sampler.py#L88) on top of the Gemma model checkpoint/weights and the tokenizer:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 10,
- "metadata": {
- "id": "4GX4pFP6rtyN"
- },
- "outputs": [],
- "source": [
- "from gemma import sampler as sampler_lib\n",
- "\n",
- "sampler = sampler_lib.Sampler(\n",
- " transformer=transformer,\n",
- " vocab=vocab,\n",
- " params=params['transformer'],\n",
- ")"
- ]
- },
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Tce3stUlHN0L"
+ },
+ "source": [
+ "##### Copyright 2024 Google LLC."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "cellView": "form",
+ "id": "tuOe1ymfHZPu"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# you may not use this file except in compliance with the License.\n",
+ "# You may obtain a copy of the License at\n",
+ "#\n",
+ "# https://www.apache.org/licenses/LICENSE-2.0\n",
+ "#\n",
+ "# Unless required by applicable law or agreed to in writing, software\n",
+ "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+ "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+ "# See the License for the specific language governing permissions and\n",
+ "# limitations under the License."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "FUOiKRSF7jc1"
+ },
+ "source": [
+ "# Inference with Gemma using JAX and Flax"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "60KmTK7o6ppd"
+ },
+ "source": [
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Tdlq6K0znh3O"
+ },
+ "source": [
+ "## Overview\n",
+ "\n",
+ "Gemma is a family of lightweight, state-of-the-art open large language models, based on the Google DeepMind Gemini research and technology. This tutorial demonstrates how to perform basic sampling/inference with the Gemma 2B Instruct model using [Google DeepMind's `gemma` library](https://github.com/google-deepmind/gemma) that was written with [JAX](https://jax.readthedocs.io) (a high-performance numerical computing library), [Flax](https://flax.readthedocs.io) (the JAX-based neural network library), [Orbax](https://orbax.readthedocs.io/) (a JAX-based library for training utilities like checkpointing), and [SentencePiece](https://github.com/google/sentencepiece) (a tokenizer/detokenizer library). Although Flax is not used directly in this notebook, Flax was used to create Gemma.\n",
+ "\n",
+ "This notebook can run on Google Colab with free T4 GPU (go to **Edit** > **Notebook settings** > Under **Hardware accelerator** select **T4 GPU**)."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aKvTsIkL98BG"
+ },
+ "source": [
+ "## Setup"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "WCgCkmQSPxkE"
+ },
+ "source": [
+ "### 1. Set up Kaggle access for Gemma\n",
+ "\n",
+ "To complete this tutorial, you first need to follow the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup), which show you how to do the following:\n",
+ "\n",
+ "* Get access to Gemma on [kaggle.com](https://www.kaggle.com/models/google/gemma/).\n",
+ "* Select a Colab runtime with sufficient resources to run the Gemma model.\n",
+ "* Generate and configure a Kaggle username and API key.\n",
+ "\n",
+ "After you've completed the Gemma setup, move on to the next section, where you'll set environment variables for your Colab environment.\n",
+ "\n",
+ "### 2. Set environment variables\n",
+ "\n",
+ "Set environment variables for `KAGGLE_USERNAME` and `KAGGLE_KEY`. When prompted with the \"Grant access?\" messages, agree to provide secret access."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {
+ "id": "lKoW-nhE-gNO"
+ },
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "from google.colab import userdata # `userdata` is a Colab API.\n",
+ "\n",
+ "os.environ[\"KAGGLE_USERNAME\"] = userdata.get('KAGGLE_USERNAME')\n",
+ "os.environ[\"KAGGLE_KEY\"] = userdata.get('KAGGLE_KEY')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "AO7a1Q4Yyc9Z"
+ },
+ "source": [
+ "### 3. Install the `gemma` library\n",
+ "\n",
+ "This notebook focuses on using a free Colab GPU. To enable hardware acceleration, click on **Edit** > **Notebook settings** > Select **T4 GPU** > **Save**.\n",
+ "\n",
+ "Next, you need to install the Google DeepMind `gemma` library from [`github.com/google-deepmind/gemma`](https://github.com/google-deepmind/gemma). If you get an error about \"pip's dependency resolver\", you can usually ignore it.\n",
+ "\n",
+ "**Note:** By installing `gemma`, you will also install [`flax`](https://flax.readthedocs.io), core [`jax`](https://jax.readthedocs.io), [`optax`](https://optax.readthedocs.io/en/latest/) (the JAX-based gradient processing and optimization library), [`orbax`](https://orbax.readthedocs.io/), and [`sentencepiece`](https://github.com/google/sentencepiece)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {
+ "id": "WWEzVJR4Fx9g"
+ },
+ "outputs": [
{
- "cell_type": "markdown",
- "metadata": {
- "id": "V9yU99Xxr59w"
- },
- "source": [
- "4. Write a prompt in `input_batch` and perform inference. You can tweak `total_generation_steps` (the number of steps performed when generating a response — this example uses `100` to preserve host memory).\n",
- "\n",
- "**Note:** If you run out of memory, click on **Runtime** > **Disconnect and delete runtime**, and then **Runtime** > **Run all**."
- ]
- },
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ " Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
+ " Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
+ " Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
+ "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m133.7/133.7 kB\u001b[0m \u001b[31m3.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+ "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m244.4/244.4 kB\u001b[0m \u001b[31m13.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+ "\u001b[?25h Building wheel for gemma (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
+ "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
+ "tensorflow-metadata 1.14.0 requires absl-py<2.0.0,>=0.9, but you have absl-py 2.1.0 which is incompatible.\u001b[0m\u001b[31m\n",
+ "\u001b[0m"
+ ]
+ }
+ ],
+ "source": [
+ "!pip install -q git+https://github.com/google-deepmind/gemma.git"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "VKLjBAe1m3Ck"
+ },
+ "source": [
+ "## Load and prepare the Gemma model\n",
+ "\n",
+ "1. Load the Gemma model with [`kagglehub.model_download`](https://github.com/Kaggle/kagglehub/blob/bddefc718182282882b72f814d407d89e5d178c4/src/kagglehub/models.py#L12), which takes three arguments:\n",
+ "\n",
+ "- `handle`: The model handle from Kaggle\n",
+ "- `path`: (Optional string) The local path\n",
+ "- `force_download`: (Optional boolean) Forces to re-download the model\n",
+ "\n",
+ "**Note:** Be mindful that the `gemma-2b-it` model is around 3.7Gb in size."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {
+ "id": "_W3FUd9lt8VT"
+ },
+ "outputs": [],
+ "source": [
+ "GEMMA_VARIANT = '2b-it' # @param ['2b', '2b-it'] {type:\"string\"}"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {
+ "id": "kFCmWEKdMA_Y"
+ },
+ "outputs": [
{
- "cell_type": "code",
- "execution_count": 11,
- "metadata": {
- "id": "Gj9jRFI5Hrv2"
- },
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Prompt:\n",
- "\n",
- "# What is the meaning of life?\n",
- "Output:\n",
- "\n",
- "\n",
- "The question of what the meaning of life is one that has occupied the minds of philosophers, theologians, and individuals for centuries. There is no single, universally accepted answer, but there are many different perspectives on this complex and multifaceted question.\n",
- "\n",
- "**Some common perspectives on the meaning of life include:**\n",
- "\n",
- "* **Biological perspective:** From a biological standpoint, the meaning of life is to survive and reproduce.\n",
- "* **Existential perspective:** Existentialists believe that life is not inherently meaningful and that\n"
- ]
- }
- ],
- "source": [
- "prompt = [\n",
- " \"\\n# What is the meaning of life?\",\n",
- "]\n",
- "\n",
- "reply = sampler(input_strings=prompt,\n",
- " total_generation_steps=100,\n",
- " )\n",
- "\n",
- "for input_string, out_string in zip(prompt, reply.text):\n",
- " print(f\"Prompt:\\n{input_string}\\nOutput:\\n{out_string}\")"
- ]
- },
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "Downloading from https://www.kaggle.com/api/v1/models/google/gemma/flax/2b-it/2/download...\n",
+ "100%|██████████| 3.67G/3.67G [00:35<00:00, 110MB/s]\n",
+ "Extracting model files...\n"
+ ]
+ }
+ ],
+ "source": [
+ "import kagglehub\n",
+ "\n",
+ "GEMMA_PATH = kagglehub.model_download(f'google/gemma/flax/{GEMMA_VARIANT}')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "metadata": {
+ "id": "nYmYTMk8aELi"
+ },
+ "outputs": [
{
- "cell_type": "markdown",
- "metadata": {
- "id": "njxRJy3qsBWw"
- },
- "source": [
- "5. (Option) Run this cell to free up memory if you have completed the notebook and want to try another prompt. Afterwards, you can instantiate the `sampler` again in step 3 and customize and run the prompt in step 4."
- ]
- },
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "GEMMA_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2\n"
+ ]
+ }
+ ],
+ "source": [
+ "print('GEMMA_PATH:', GEMMA_PATH)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ytNi47xSlw71"
+ },
+ "source": [
+ "**Note:** The path from the output above is where the model weights and tokenizer are saved locally, you will need them for later."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "92BcvYdemXbd"
+ },
+ "source": [
+ "2. Check the location of the model weights and the tokenizer, then set the path variables. The tokenizer directory will be in the main directory where you downloaded the model, while the model weights will be in a sub-directory. For example:\n",
+ "\n",
+ "- The `tokenizer.model` file will be in `/LOCAL/PATH/TO/gemma/flax/2b-it/2`).\n",
+ "- The model checkpoint will be in `/LOCAL/PATH/TO/gemma/flax/2b-it/2/2b-it`)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "metadata": {
+ "id": "QY6OnASOpZbW"
+ },
+ "outputs": [
{
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "qxX6qfFdNGHy"
- },
- "outputs": [],
- "source": [
- "del sampler"
- ]
- },
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "CKPT_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/2b-it\n",
+ "TOKENIZER_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/tokenizer.model\n"
+ ]
+ }
+ ],
+ "source": [
+ "CKPT_PATH = os.path.join(GEMMA_PATH, GEMMA_VARIANT)\n",
+ "TOKENIZER_PATH = os.path.join(GEMMA_PATH, 'tokenizer.model')\n",
+ "print('CKPT_PATH:', CKPT_PATH)\n",
+ "print('TOKENIZER_PATH:', TOKENIZER_PATH)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "jc0ZzYIW0TSN"
+ },
+ "source": [
+ "## Perform sampling/inference"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aEe3p8geqekV"
+ },
+ "source": [
+ "1. Load and format the Gemma model checkpoint with the [`gemma.params.load_and_format_params`](https://github.com/google-deepmind/gemma/blob/c6bd156c246530e1620a7c62de98542a377e3934/gemma/params.py#L27) method:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "metadata": {
+ "id": "Mnr52JQVqKRw"
+ },
+ "outputs": [],
+ "source": [
+ "from gemma import params as params_lib\n",
+ "\n",
+ "params = params_lib.load_and_format_params(CKPT_PATH)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "-Xpnb2igrGjk"
+ },
+ "source": [
+ "2. Load the Gemma tokenizer, constructed using [`sentencepiece.SentencePieceProcessor`](https://github.com/google/sentencepiece/blob/4d6a1f41069c4636c51a5590f7578a0dbed83450/python/src/sentencepiece/__init__.py#L423):"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "metadata": {
+ "id": "-T0ZHff5rNSy"
+ },
+ "outputs": [
{
- "cell_type": "markdown",
- "metadata": {
- "id": "bzKsCGIN0yX5"
- },
- "source": [
- "## Learn more\n",
- "\n",
- "- You can learn more about the Google DeepMind [`gemma` library on GitHub](https://github.com/google-deepmind/gemma), which contains docstrings of methods you used in this tutorial, such as [`gemma.params`](https://github.com/google-deepmind/gemma/blob/main/gemma/params.py),\n",
- "[`gemma.transformer`](https://github.com/google-deepmind/gemma/blob/main/gemma/transformer.py), and\n",
- "[`gemma.sampler`](https://github.com/google-deepmind/gemma/blob/main/gemma/sampler.py).\n",
- "- The following libraries have their own documentation sites: [core JAX](https://jax.readthedocs.io), [Flax](https://flax.readthedocs.io), and [Orbax](https://orbax.readthedocs.io/).\n",
- "- For `sentencepiece` tokenizer/detokenizer documentation, check out [Google's `sentencepiece` GitHub repo](https://github.com/google/sentencepiece).\n",
- "- For `kagglehub` documentation, check out `README.md` on [Kaggle's `kagglehub` GitHub repo](https://github.com/Kaggle/kagglehub).\n",
- "- Learn how to [use Gemma models with Google Cloud Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/open-models/use-gemma)."
+ "data": {
+ "text/plain": [
+ "True"
]
+ },
+ "execution_count": 8,
+ "metadata": {},
+ "output_type": "execute_result"
}
- ],
- "metadata": {
- "accelerator": "GPU",
- "colab": {
- "name": "jax_inference.ipynb",
- "toc_visible": true
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
+ ],
+ "source": [
+ "import sentencepiece as spm\n",
+ "\n",
+ "vocab = spm.SentencePieceProcessor()\n",
+ "vocab.Load(TOKENIZER_PATH)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "IkAf4fkNrY-3"
+ },
+ "source": [
+ "3. To automatically load the correct configuration from the Gemma model checkpoint, use [`gemma.transformer.TransformerConfig`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L65). The `cache_size` argument is the number of time steps in the Gemma `Transformer` cache. Afterwards, instantiate the Gemma model as `transformer` with [`gemma.transformer.Transformer`](https://github.com/google-deepmind/gemma/blob/56e501ce147af4ea5c23cc0ddf5a9c4a6b7bd0d0/gemma/transformer.py#L136) (which inherits from [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html)).\n",
+ "\n",
+ "**Note:** The vocabulary size is smaller than the number of input embeddings because of unused tokens in the current Gemma release."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "metadata": {
+ "id": "4PNWxDhvrRXJ"
+ },
+ "outputs": [],
+ "source": [
+ "from gemma import transformer as transformer_lib\n",
+ "\n",
+ "transformer_config = transformer_lib.TransformerConfig.from_params(\n",
+ " params=params,\n",
+ " cache_size=1024\n",
+ ")\n",
+ "\n",
+ "transformer = transformer_lib.Transformer(transformer_config)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "vs0vgmXVroBq"
+ },
+ "source": [
+ "3. Create a `sampler` with [`gemma.sampler.Sampler`](https://github.com/google-deepmind/gemma/blob/c6bd156c246530e1620a7c62de98542a377e3934/gemma/sampler.py#L88) on top of the Gemma model checkpoint/weights and the tokenizer:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "metadata": {
+ "id": "4GX4pFP6rtyN"
+ },
+ "outputs": [],
+ "source": [
+ "from gemma import sampler as sampler_lib\n",
+ "\n",
+ "sampler = sampler_lib.Sampler(\n",
+ " transformer=transformer,\n",
+ " vocab=vocab,\n",
+ " params=params['transformer'],\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "V9yU99Xxr59w"
+ },
+ "source": [
+ "4. Write a prompt in `input_batch` and perform inference. You can tweak `total_generation_steps` (the number of steps performed when generating a response — this example uses `100` to preserve host memory).\n",
+ "\n",
+ "**Note:** If you run out of memory, click on **Runtime** > **Disconnect and delete runtime**, and then **Runtime** > **Run all**."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 11,
+ "metadata": {
+ "id": "Gj9jRFI5Hrv2"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Prompt:\n",
+ "\n",
+ "# What is the meaning of life?\n",
+ "Output:\n",
+ "\n",
+ "\n",
+ "The question of what the meaning of life is one that has occupied the minds of philosophers, theologians, and individuals for centuries. There is no single, universally accepted answer, but there are many different perspectives on this complex and multifaceted question.\n",
+ "\n",
+ "**Some common perspectives on the meaning of life include:**\n",
+ "\n",
+ "* **Biological perspective:** From a biological standpoint, the meaning of life is to survive and reproduce.\n",
+ "* **Existential perspective:** Existentialists believe that life is not inherently meaningful and that\n"
+ ]
}
+ ],
+ "source": [
+ "prompt = [\n",
+ " \"\\n# What is the meaning of life?\",\n",
+ "]\n",
+ "\n",
+ "reply = sampler(input_strings=prompt,\n",
+ " total_generation_steps=100,\n",
+ " )\n",
+ "\n",
+ "for input_string, out_string in zip(prompt, reply.text):\n",
+ " print(f\"Prompt:\\n{input_string}\\nOutput:\\n{out_string}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "njxRJy3qsBWw"
+ },
+ "source": [
+ "5. (Optional) Run this cell to free up memory if you have completed the notebook and want to try another prompt. Afterwards, you can instantiate the `sampler` again in step 3 and customize and run the prompt in step 4."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "qxX6qfFdNGHy"
+ },
+ "outputs": [],
+ "source": [
+ "del sampler"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "bzKsCGIN0yX5"
+ },
+ "source": [
+ "## Learn more\n",
+ "\n",
+ "- You can learn more about the Google DeepMind [`gemma` library on GitHub](https://github.com/google-deepmind/gemma), which contains docstrings of modules you used in this tutorial, such as [`gemma.params`](https://github.com/google-deepmind/gemma/blob/main/gemma/params.py),\n",
+ "[`gemma.transformer`](https://github.com/google-deepmind/gemma/blob/main/gemma/transformer.py), and\n",
+ "[`gemma.sampler`](https://github.com/google-deepmind/gemma/blob/main/gemma/sampler.py).\n",
+ "- The following libraries have their own documentation sites: [core JAX](https://jax.readthedocs.io), [Flax](https://flax.readthedocs.io), and [Orbax](https://orbax.readthedocs.io/).\n",
+ "- For `sentencepiece` tokenizer/detokenizer documentation, check out [Google's `sentencepiece` GitHub repo](https://github.com/google/sentencepiece).\n",
+ "- For `kagglehub` documentation, check out `README.md` on [Kaggle's `kagglehub` GitHub repo](https://github.com/Kaggle/kagglehub).\n",
+ "- Learn how to [use Gemma models with Google Cloud Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/open-models/use-gemma)."
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "name": "jax_inference.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
},
- "nbformat": 4,
- "nbformat_minor": 0
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.11.7"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
}
![](\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\")
View on ai.google.dev\n",
- " ![](\"https://www.tensorflow.org/images/colab_logo_32px.png\"){kind=logo}
Run in Google Colab\n",
- " ![](\"https://ai.google.dev/images/cloud-icon.svg\"){kind=icon}
Open in Vertex AI\n",
- " ![](\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\")
View source on GitHub\n",
- "