diff --git "a/Finetune_PaliGemma_for_image_description.ipynb" "b/Finetune_PaliGemma_for_image_description.ipynb"
new file mode 100644--- /dev/null
+++ "b/Finetune_PaliGemma_for_image_description.ipynb"
@@ -0,0 +1,7033 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "c3hpiPPEqmf6"
+ },
+ "source": [
+ "##### Copyright 2024 Google LLC."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "cellView": "form",
+ "id": "bVm-2hW9z9HR"
+ },
+ "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": "u71STQRgnQ3a"
+ },
+ "source": [
+ "# Fine-tune PaliGemma for Image Description with Custom Dataset"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "wR53lePHuiP-"
+ },
+ "source": [
+ "This notebook guides you through the process of fine-tuning [PaliGemma](https://ai.google.dev/gemma/docs/paligemma), a powerful vision-language model, for bird description using [JAX](https://jax.readthedocs.io/en/latest/installation.html). We will leverage a curated subset of a bird species dataset and enrich it with descriptive text for each bird. The resulting dataset, comprising 3,692 image-description pairs, will be used to fine-tune PaliGemma, enabling it to generate accurate and detailed descriptions of bird images.\n",
+ "\n",
+ "
\n",
+ " \n"
+ ],
+ "text/plain": [
+ " class id filepaths labels \\\n",
+ "0 0.0 ABBOTTS BABBLER/001.jpg ABBOTTS BABBLER \n",
+ "1 0.0 ABBOTTS BABBLER/007.jpg ABBOTTS BABBLER \n",
+ "2 0.0 ABBOTTS BABBLER/008.jpg ABBOTTS BABBLER \n",
+ "3 0.0 ABBOTTS BABBLER/009.jpg ABBOTTS BABBLER \n",
+ "4 0.0 ABBOTTS BABBLER/002.jpg ABBOTTS BABBLER \n",
+ "... ... ... ... \n",
+ "3917 524.0 BLACK BREASTED PUFFBIRD/3.jpg BLACK BREASTED PUFFBIRD \n",
+ "3918 524.0 BLACK BREASTED PUFFBIRD/4.jpg BLACK BREASTED PUFFBIRD \n",
+ "3919 524.0 BLACK BREASTED PUFFBIRD/1.jpg BLACK BREASTED PUFFBIRD \n",
+ "3920 524.0 BLACK BREASTED PUFFBIRD/2.jpg BLACK BREASTED PUFFBIRD \n",
+ "3921 524.0 BLACK BREASTED PUFFBIRD/5.jpg BLACK BREASTED PUFFBIRD \n",
+ "\n",
+ " data set scientific name \\\n",
+ "0 train MALACOCINCLA ABBOTTI \n",
+ "1 train MALACOCINCLA ABBOTTI \n",
+ "2 train MALACOCINCLA ABBOTTI \n",
+ "3 train MALACOCINCLA ABBOTTI \n",
+ "4 train MALACOCINCLA ABBOTTI \n",
+ "... ... ... \n",
+ "3917 valid NOTHARCHUS PECTORALIS \n",
+ "3918 valid NOTHARCHUS PECTORALIS \n",
+ "3919 valid NOTHARCHUS PECTORALIS \n",
+ "3920 valid NOTHARCHUS PECTORALIS \n",
+ "3921 valid NOTHARCHUS PECTORALIS \n",
+ "\n",
+ " bird_description \n",
+ "0 Abbott's Babbler: Look for this small insectiv... \n",
+ "1 Abbott's Babbler: Look for this small insectiv... \n",
+ "2 Abbott's Babbler: Look for this small insectiv... \n",
+ "3 Abbott's Babbler: Look for this small insectiv... \n",
+ "4 Abbott's Babbler: Look for this small insectiv... \n",
+ "... ... \n",
+ "3917 Black-breasted Puffbird: Observe the medium-si... \n",
+ "3918 Black-breasted Puffbird: Observe the medium-si... \n",
+ "3919 Black-breasted Puffbird: Observe the medium-si... \n",
+ "3920 Black-breasted Puffbird: Observe the medium-si... \n",
+ "3921 Black-breasted Puffbird: Observe the medium-si... \n",
+ "\n",
+ "[3922 rows x 6 columns]"
+ ]
+ },
+ "execution_count": 9,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "# Load the descriptions DataFrame\n",
+ "birds_discription_df = pd.read_csv(\"birds_description.csv\")\n",
+ "birds_discription_df"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "CZg6Rj0Njqjp"
+ },
+ "source": [
+ "### Display number of unique species, number of training and validation samples"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "uA-l53mTjqjp"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Number of unique bird species: 23\n",
+ "Number of training samples: 3692\n",
+ "Number of validation samples: 115\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(\"Number of unique bird species:\", len(birds_discription_df['labels'].unique()))\n",
+ "print(\"Number of training samples:\", (birds_discription_df['data set'] == \"train\").sum())\n",
+ "print(\"Number of validation samples:\", (birds_discription_df['data set'] == \"valid\").sum())"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "1Qkqjs2Yjqjp"
+ },
+ "source": [
+ "## Convert DataFrame to JSON Lines for Finetuning\n",
+ "\n",
+ "Since PaliGemma expects data in JSON Lines format for finetuning, we convert the DataFrame containing bird information into separate JSON Lines files for training and validation data.\n",
+ "\n",
+ "For fine-tuning, we only need two columns from the DataFrame:\n",
+ "\n",
+ "* `\"filepaths\"`: This column contains the paths to the bird images.\n",
+ "* `\"bird_description\"`: This column contains descriptions for each bird species."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "jzFnL_wujqjp"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "DataFrame converted to JSON Lines format and saved to train.jsonl and valid.jsonl\n"
+ ]
+ }
+ ],
+ "source": [
+ "def df_to_jsonl(df):\n",
+ " \"\"\"Converts a pandas DataFrame to separate JSON Lines files for train and validation data.\n",
+ "\n",
+ " Args:\n",
+ " df: The pandas DataFrame to convert.\n",
+ "\n",
+ " Returns:\n",
+ " None. Writes the JSON Lines to separate 'train.jsonl' and 'valid.jsonl' files.\n",
+ " \"\"\"\n",
+ "\n",
+ " train_file = open('train.jsonl', 'w')\n",
+ " valid_file = open('valid.jsonl', 'w')\n",
+ "\n",
+ " try:\n",
+ " for index, row in df.iterrows():\n",
+ " if row['data set'] == 'train':\n",
+ " data = {\n",
+ " 'prefix': '',\n",
+ " 'suffix': row['bird_description'],\n",
+ " 'image': row['filepaths']\n",
+ " }\n",
+ " json.dump(data, train_file)\n",
+ " train_file.write('\\n')\n",
+ " elif row['data set'] == 'valid':\n",
+ " data = {\n",
+ " 'prefix': '',\n",
+ " 'suffix': row['bird_description'],\n",
+ " 'image': row['filepaths']\n",
+ " }\n",
+ " json.dump(data, valid_file)\n",
+ " valid_file.write('\\n')\n",
+ "\n",
+ " finally:\n",
+ " train_file.close()\n",
+ " valid_file.close()\n",
+ "\n",
+ " print(\"DataFrame converted to JSON Lines format and saved to train.jsonl and valid.jsonl\")\n",
+ "\n",
+ "# Convert the birds description DataFrame to JSON Lines format\n",
+ "df_to_jsonl(birds_discription_df)\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aMyiB9bUjqjp"
+ },
+ "source": [
+ "## Display first record from JSON lines file"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "xnud6HhMjqjp"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "First Record:\n",
+ "{\n",
+ " \"prefix\": \"\",\n",
+ " \"suffix\": \"Abbott's Babbler: Look for this small insectivorous bird with distinctive streaked brown plumage and pale buff underparts, making its home in the diverse landscapes of South Asia.\",\n",
+ " \"image\": \"ABBOTTS BABBLER/001.jpg\"\n",
+ "}\n"
+ ]
+ }
+ ],
+ "source": [
+ "def display_first_record(filename):\n",
+ " \"\"\"Opens a JSON Lines file and displays only the first record.\n",
+ "\n",
+ " Args:\n",
+ " filename: The path to the JSON Lines file.\n",
+ " \"\"\"\n",
+ " try:\n",
+ " with open(filename, 'r') as f:\n",
+ " first_line = f.readline().strip()\n",
+ " if first_line: # Check if there's data in the file\n",
+ " data = json.loads(first_line)\n",
+ " print(f\"First Record:\\n{json.dumps(data, indent=2)}\")\n",
+ " else:\n",
+ " print(\"File is empty or corrupt.\")\n",
+ " except FileNotFoundError:\n",
+ " print(f\"Error: File '{filename}' not found.\")\n",
+ "\n",
+ "# Display the first record in 'train.jsonl'\n",
+ "display_first_record('train.jsonl')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "rv7w-cGuLj5o"
+ },
+ "source": [
+ "# Download and configure PaliGemma model\n",
+ "This section retrieves the pre-trained model and tokenizer weights."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "v5e1RGCDjqjq"
+ },
+ "outputs": [],
+ "source": [
+ "# Define model\n",
+ "model_config = ml_collections.FrozenConfigDict({\n",
+ " \"llm\": {\"vocab_size\": 257_152},\n",
+ " \"img\": {\"variant\": \"So400m/14\", \"pool_type\": \"none\", \"scan\": True, \"dtype_mm\": \"float16\"}\n",
+ "})\n",
+ "model = paligemma.Model(**model_config)\n",
+ "tokenizer = sentencepiece.SentencePieceProcessor(TOKENIZER_PATH)\n",
+ "\n",
+ "# Load params - this can take up to 1 minute in a notebook.\n",
+ "params = paligemma.load(None, MODEL_PATH, model_config)\n",
+ "\n",
+ "# Define `decode` function to sample outputs from the model.\n",
+ "decode_fn = predict_fns.get_all(model)['decode']\n",
+ "decode = functools.partial(decode_fn, devices=jax.devices(), eos_token=tokenizer.eos_id())"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "uidBwmb8LwZ5"
+ },
+ "source": [
+ "### Move model parameters into GPU/TPU memory"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "RWOdf_fw2SAO"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ " == Model params == \n",
+ "img/Transformer/encoder_norm/bias (1152,) float16\n",
+ "img/Transformer/encoder_norm/scale (1152,) float16\n",
+ "img/Transformer/encoderblock/LayerNorm_0/bias (27, 1152) float16\n",
+ "img/Transformer/encoderblock/LayerNorm_0/scale (27, 1152) float16\n",
+ "img/Transformer/encoderblock/LayerNorm_1/bias (27, 1152) float16\n",
+ "img/Transformer/encoderblock/LayerNorm_1/scale (27, 1152) float16\n",
+ "img/Transformer/encoderblock/MlpBlock_0/Dense_0/bias (27, 4304) float16\n",
+ "img/Transformer/encoderblock/MlpBlock_0/Dense_0/kernel (27, 1152, 4304) float16\n",
+ "img/Transformer/encoderblock/MlpBlock_0/Dense_1/bias (27, 1152) float16\n",
+ "img/Transformer/encoderblock/MlpBlock_0/Dense_1/kernel (27, 4304, 1152) float16\n",
+ "img/Transformer/encoderblock/MultiHeadDotProductAttention_0/key/bias (27, 16, 72) float16\n",
+ "img/Transformer/encoderblock/MultiHeadDotProductAttention_0/key/kernel (27, 1152, 16, 72) float16\n",
+ "img/Transformer/encoderblock/MultiHeadDotProductAttention_0/out/bias (27, 1152) float16\n",
+ "img/Transformer/encoderblock/MultiHeadDotProductAttention_0/out/kernel (27, 16, 72, 1152) float16\n",
+ "img/Transformer/encoderblock/MultiHeadDotProductAttention_0/query/bias (27, 16, 72) float16\n",
+ "img/Transformer/encoderblock/MultiHeadDotProductAttention_0/query/kernel (27, 1152, 16, 72) float16\n",
+ "img/Transformer/encoderblock/MultiHeadDotProductAttention_0/value/bias (27, 16, 72) float16\n",
+ "img/Transformer/encoderblock/MultiHeadDotProductAttention_0/value/kernel (27, 1152, 16, 72) float16\n",
+ "img/embedding/bias (1152,) float16\n",
+ "img/embedding/kernel (14, 14, 3, 1152) float16\n",
+ "img/head/bias (2048,) float16\n",
+ "img/head/kernel (1152, 2048) float16\n",
+ "img/pos_embedding (1, 256, 1152) float16\n",
+ "llm/embedder/input_embedding (257152, 2048) float16\n",
+ "llm/final_norm/scale (2048,) float16\n",
+ "llm/layers/attn/attn_vec_einsum/w (18, 8, 256, 2048) float32\n",
+ "llm/layers/attn/kv_einsum/w (18, 2, 1, 2048, 256) float32\n",
+ "llm/layers/attn/q_einsum/w (18, 8, 2048, 256) float32\n",
+ "llm/layers/mlp/gating_einsum (18, 2, 2048, 16384) float16\n",
+ "llm/layers/mlp/linear (18, 16384, 2048) float16\n",
+ "llm/layers/pre_attention_norm/scale (18, 2048) float16\n",
+ "llm/layers/pre_ffw_norm/scale (18, 2048) float16\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Create a pytree mask of the trainable params.\n",
+ "def is_trainable_param(name, param): # pylint: disable=unused-argument\n",
+ " if name.startswith(\"llm/layers/attn/\"): return True\n",
+ " if name.startswith(\"llm/\"): return False\n",
+ " if name.startswith(\"img/\"): return False\n",
+ " raise ValueError(f\"Unexpected param name {name}\")\n",
+ "trainable_mask = big_vision.utils.tree_map_with_names(is_trainable_param, params)\n",
+ "\n",
+ "# If more than one device is available (e.g. multiple GPUs) the parameters can\n",
+ "# be sharded across them to reduce HBM usage per device.\n",
+ "mesh = jax.sharding.Mesh(jax.devices(), (\"data\"))\n",
+ "\n",
+ "data_sharding = jax.sharding.NamedSharding(\n",
+ " mesh, jax.sharding.PartitionSpec(\"data\"))\n",
+ "\n",
+ "params_sharding = big_vision.sharding.infer_sharding(\n",
+ " params, strategy=[('.*', 'fsdp(axis=\"data\")')], mesh=mesh)\n",
+ "\n",
+ "# Yes: Some donated buffers are not usable.\n",
+ "warnings.filterwarnings(\n",
+ " \"ignore\", message=\"Some donated buffers were not usable\")\n",
+ "\n",
+ "@functools.partial(jax.jit, donate_argnums=(0,), static_argnums=(1,))\n",
+ "def maybe_cast_to_f32(params, trainable):\n",
+ " return jax.tree.map(lambda p, m: p.astype(jnp.float32) if m else p,\n",
+ " params, trainable)\n",
+ "\n",
+ "# Loading all params in simultaneous - albeit much faster and more succinct -\n",
+ "# requires more RAM than the notebook runtimes have by default.\n",
+ "# Instead we do it param by param.\n",
+ "params, treedef = jax.tree.flatten(params)\n",
+ "sharding_leaves = jax.tree.leaves(params_sharding)\n",
+ "trainable_leaves = jax.tree.leaves(trainable_mask)\n",
+ "for idx, (sharding, trainable) in enumerate(zip(sharding_leaves, trainable_leaves)):\n",
+ " params[idx] = big_vision.utils.reshard(params[idx], sharding)\n",
+ " params[idx] = maybe_cast_to_f32(params[idx], trainable)\n",
+ " params[idx].block_until_ready()\n",
+ "params = jax.tree.unflatten(treedef, params)\n",
+ "\n",
+ "# Print params to show what the model is made of.\n",
+ "def parameter_overview(params):\n",
+ " for path, arr in big_vision.utils.tree_flatten_with_names(params)[0]:\n",
+ " print(f\"{path:80s} {str(arr.shape):22s} {arr.dtype}\")\n",
+ "\n",
+ "print(\" == Model params == \")\n",
+ "parameter_overview(params)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "iD_9XXQkn1Mv"
+ },
+ "source": [
+ "# Prepare to tune the model"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "8SRW0NuU4UcW"
+ },
+ "outputs": [],
+ "source": [
+ "def preprocess_image(image, size=224):\n",
+ " # Model has been trained to handle images of different aspects ratios\n",
+ " # resized to 224x224 in the range [-1, 1]. Bilinear and antialias resize\n",
+ " # options are helpful to improve quality in some tasks.\n",
+ " image = np.asarray(image)\n",
+ " if image.ndim == 2: # Convert image without last channel into greyscale.\n",
+ " image = np.stack((image,)*3, axis=-1)\n",
+ " image = image[..., :3] # Remove alpha layer.\n",
+ " assert image.shape[-1] == 3\n",
+ "\n",
+ " image = tf.constant(image)\n",
+ " image = tf.image.resize(image, (size, size), method='bilinear', antialias=True)\n",
+ " return image.numpy() / 127.5 - 1.0 # [0, 255]->[-1,1]\n",
+ "\n",
+ "def preprocess_tokens(prefix, suffix=None, seqlen=None):\n",
+ " # Model has been trained to handle tokenized text composed of a prefix with\n",
+ " # full attention and a suffix with causal attention.\n",
+ " separator = \"\\n\"\n",
+ " tokens = tokenizer.encode(prefix, add_bos=True) + tokenizer.encode(separator)\n",
+ " mask_ar = [0] * len(tokens) # 0 to use full attention for prefix.\n",
+ " mask_loss = [0] * len(tokens) # 0 to not use prefix tokens in the loss.\n",
+ "\n",
+ " if suffix:\n",
+ " suffix = tokenizer.encode(suffix, add_eos=True)\n",
+ " tokens += suffix\n",
+ " mask_ar += [1] * len(suffix) # 1 to use causal attention for suffix.\n",
+ " mask_loss += [1] * len(suffix) # 1 to use suffix tokens in the loss.\n",
+ "\n",
+ " mask_input = [1] * len(tokens) # 1 if it's a token, 0 if padding.\n",
+ " if seqlen:\n",
+ " padding = [0] * max(0, seqlen - len(tokens))\n",
+ " tokens = tokens[:seqlen] + padding\n",
+ " mask_ar = mask_ar[:seqlen] + padding\n",
+ " mask_loss = mask_loss[:seqlen] + padding\n",
+ " mask_input = mask_input[:seqlen] + padding\n",
+ "\n",
+ " return jax.tree.map(np.array, (tokens, mask_ar, mask_loss, mask_input))\n",
+ "\n",
+ "def postprocess_tokens(tokens):\n",
+ " tokens = tokens.tolist() # np.array to list[int]\n",
+ " try: # Remove tokens at and after EOS if any.\n",
+ " eos_pos = tokens.index(tokenizer.eos_id())\n",
+ " tokens = tokens[:eos_pos]\n",
+ " except ValueError:\n",
+ " pass\n",
+ " return tokenizer.decode(tokens)\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ovgWBgdHJZq3"
+ },
+ "source": [
+ "### Create the training and validation iterators"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "whzWOojGOtzi"
+ },
+ "outputs": [],
+ "source": [
+ "SEQLEN = 128\n",
+ "\n",
+ "TRAIN_DATA_DIR = 'train/'\n",
+ "VALID_DATA_DIR = 'valid/'\n",
+ "\n",
+ "# Load training data\n",
+ "train_dataset = big_vision.datasets.jsonl.DataSource(\n",
+ " os.path.join(\"train.jsonl\"),\n",
+ " fopen_keys={\"image\": TRAIN_DATA_DIR})\n",
+ "\n",
+ "# Load validation data\n",
+ "val_dataset = big_vision.datasets.jsonl.DataSource(\n",
+ " os.path.join(\"valid.jsonl\"),\n",
+ " fopen_keys={\"image\": VALID_DATA_DIR})\n",
+ "\n",
+ "def train_data_iterator():\n",
+ " \"\"\"Never ending iterator over training examples.\"\"\"\n",
+ " # Shuffle examples and repeat so one can train for many epochs.\n",
+ " dataset = train_dataset.get_tfdata().shuffle(1_000).repeat()\n",
+ " for example in dataset.as_numpy_iterator():\n",
+ " image = Image.open(io.BytesIO(example[\"image\"]))\n",
+ " image = preprocess_image(image)\n",
+ "\n",
+ " # Define prefix for tokenization\n",
+ " prefix = \"describe en\"\n",
+ " suffix = example[\"suffix\"].decode().lower()\n",
+ " tokens, mask_ar, mask_loss, _ = preprocess_tokens(prefix, suffix, SEQLEN)\n",
+ "\n",
+ " yield {\n",
+ " \"image\": np.asarray(image),\n",
+ " \"text\": np.asarray(tokens),\n",
+ " \"mask_ar\": np.asarray(mask_ar),\n",
+ " \"mask_loss\": np.asarray(mask_loss),\n",
+ " }\n",
+ "\n",
+ "\n",
+ "def validation_data_iterator():\n",
+ " \"\"\"Single iterator over validation examples.\"\"\"\n",
+ " for example in val_dataset.get_tfdata(ordered=True).as_numpy_iterator():\n",
+ " image = Image.open(io.BytesIO(example[\"image\"]))\n",
+ " image = preprocess_image(image)\n",
+ "\n",
+ " # Define prefix for tokenization\n",
+ " prefix = \"describe en\" # Could also be a different prefix per example 'describe en'\n",
+ " tokens, mask_ar, _, mask_input = preprocess_tokens(prefix, seqlen=SEQLEN)\n",
+ "\n",
+ " yield {\n",
+ " \"image\": np.asarray(image),\n",
+ " \"text\": np.asarray(tokens),\n",
+ " \"mask_ar\": np.asarray(mask_ar),\n",
+ " \"mask_input\": np.asarray(mask_input),\n",
+ " }\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "84olaM5dCiAl"
+ },
+ "source": [
+ "### View sample training examples"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "BzJfb5t0nsLq"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Training Examples\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "
\n",
+ " \n",
+ "
yellow-billed chough: observe this medium-sized crow easily recognizable by its sleek black feathers and contrasting long yellow bill, a common sight in the mountainous regions of europe and asia, noted for its acrobatic flight and yellow bill.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " \n",
+ "
american bittern: notice this medium-sized heron, a master of camouflage, blending seamlessly into the marshes of north and central america, distinguished by its cryptic plumage and elongated neck.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " \n",
+ "
american pipit: look for this small, ground-dwelling songbird with streaked brown plumage and long tail, filling the air with melodic tunes across the grasslands of north america, europe, and asia.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " \n",
+ "
alexandrine parakeet: notice this medium-sized parrot bursting with a rainbow of colors, a vibrant resident of the forests and woodlands of south and southeast asia, distinguished by its large size and distinctive red beak.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " \n",
+ "
crowned crane: observe this epitome of elegance with a golden crown, calling the wetlands and grasslands of sub-saharan africa home, distinguished by its tall stature, long legs, and regal posture.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " \n",
+ "
african pygmy goose: look for this small freshwater goose with dark brown feathers and contrasting white markings, gracing the lakes and rivers of sub-saharan africa, noted for its petite size and distinctive facial markings.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " \n",
+ "
american bittern: notice this medium-sized heron, a master of camouflage, blending seamlessly into the marshes of north and central america, distinguished by its cryptic plumage and elongated neck.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " \n",
+ "
african pygmy goose: look for this small freshwater goose with dark brown feathers and contrasting white markings, gracing the lakes and rivers of sub-saharan africa, noted for its petite size and distinctive facial markings.
\n",
+ "
\n",
+ " "
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "def render_inline(image, resize=(128, 128)):\n",
+ " \"\"\"Converts an image array into inline HTML.\n",
+ "\n",
+ " Args:\n",
+ " image (numpy.ndarray): The image array to convert.\n",
+ " resize (tuple): Optional. Size to resize the image to before conversion. Default is (128, 128).\n",
+ "\n",
+ " Returns:\n",
+ " str: HTML representation of the image encoded in base64.\n",
+ " \"\"\"\n",
+ " image = Image.fromarray(image)\n",
+ " image.resize(resize)\n",
+ " with io.BytesIO() as buffer:\n",
+ " image.save(buffer, format='jpeg')\n",
+ " image_b64 = str(base64.b64encode(buffer.getvalue()), \"utf-8\")\n",
+ " return f\"data:image/jpeg;base64,{image_b64}\"\n",
+ "\n",
+ "def render_example(image, description):\n",
+ " \"\"\"Generates HTML for displaying an image with its description.\n",
+ "\n",
+ " Args:\n",
+ " image (numpy.ndarray): The image array to display.\n",
+ " description (str): The description of the image.\n",
+ "\n",
+ " Returns:\n",
+ " str: HTML representation of the image with description.\n",
+ " \"\"\"\n",
+ " image = ((image + 1)/2 * 255).astype(np.uint8) # [-1,1] -> [0, 255]\n",
+ " return f\"\"\"\n",
+ "
Validation: abbott's booby: keep an eye out for this majestic seabird soaring over the tropical indian and pacific oceans, especially around christmas island, noted for its large size and long wingspan while in flight.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: albatross: observe this large seabird known for its brooding behavior, providing a visual spectacle with its wingspan and long flight feathers, and for nesting on remote islands in the southern ocean.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: albatross: observe this large seafaring bird known for its impressive wingspan and effortless gliding flight, a majestic sight in the southern hemisphere, noted for its long wings and buoyant flight.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: african black-crowned night-jar: observe this large nocturnal bird of prey in the savannas and woodlands of sub-saharan africa, noted for its distinctive black crown and yellow bill.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: yellow-billed chough: observe this medium-sized crow easily recognizable by its sleek black feathers and contrasting long yellow bill, a common sight in the mountainous regions of europe and asia, noted for its acrobatic flight and yellow bill.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: american redstart: observe this small insectivorous bird with red markings on its wings and tail, found in the deciduous forests of north america.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: american redstart: observe this small insectivorous warbler with contrasting black feathers and vibrant orange patches on its wings, a resident of north america, noted for its active foraging behavior and constant tail flicking.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: american redstart: observe this small insectivorous bird with red underparts and black feathers, found in the woodlands of north and central america, noted for its slender body and long tail.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: american redstart: observe this small insectivorous warbler with contrasting black feathers and vibrant orange patches on its wings, a resident of north america, noted for its active foraging behavior and constant tail flicking.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: american flamingo: observe this large wading bird, a standout in any tropical landscape, known for its vibrant pink plumage and long legs, a true icon of the caribbean.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: american flamingo: look for this large wading bird recognizable by its vibrant pink to red plumage, long legs, and distinctively curved bill, a striking inhabitant of the caribbean, galΓ‘pagos, and coastal regions of south america, noted for its long neck and legs.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: cuban toreador: a close-knit family of medium-sized birds with elongated snouts and robust wings, inhabiting the forests of central and south america, noted for their impressive casque of feathers.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: african emerald cuckoo: look for this small parasitic bird with shimmering emerald feathers, a unique inhabitant of the forests of sub-saharan africa, characterized by its slender body and long tail.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: african black hornbill: observe this medium-sized african bird with an impressive red bill and long beak, a distinctive sight endemic to the forests of southern and east africa, known for its impressive size and robust build.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: abyssinian ground hornbill: notice this large terrestrial bird commanding attention with its distinctive helmet-like structure on its head, a resident of the savannas and woodlands of sub-saharan africa, characterized by its powerful bill and long eyelashes.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: american goldfinch: observe this small, colorful finch familiarizing himself with backyard feeders, noted for its distinctive yellow plumage and black cap.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: american goldfinch: notice this small, colorful finch with bright yellow feathers and contrasting black wings, bringing joy to the landscapes of north america, noted for its conical bill and black cap.
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
Prediction: long-billed dowitch: look for this large, long-billed shorebird in the grasslands and savannas of southern africa, characterized by its distinctive black and white patterning and long, slender bill.
\n",
+ "
\n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "
Validation: alexandrine parakeet: notice this medium-sized parrot bursting with a rainbow of colors, a vibrant resident of the forests and woodlands of south and southeast asia, distinguished by its large size and distinctive red beak.
\n",
+ "
\n",
+ " "
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "def data_iterator(data_type=None):\n",
+ " \"\"\"Iterates over examples for validation or prediction.\n",
+ "\n",
+ " Args:\n",
+ " data_type (str): Type of data to iterate ('prediction' or 'validation').\n",
+ "\n",
+ " Yields:\n",
+ " dict: Dictionary containing image and text data based on data_type.\n",
+ " For 'prediction': {'image': np.array, 'text': np.array, 'mask_ar': np.array, 'mask_input': np.array}\n",
+ " For 'validation': {'image': np.array, 'text': np.array, 'mask_ar': np.array, 'mask_loss': np.array}\n",
+ " \"\"\"\n",
+ " for example in val_dataset.get_tfdata(ordered=True).as_numpy_iterator():\n",
+ " image = Image.open(io.BytesIO(example[\"image\"]))\n",
+ " image = preprocess_image(image)\n",
+ "\n",
+ " prefix = \"describe en\"\n",
+ " if data_type == \"prediction\":\n",
+ " tokens, mask_ar, _, mask_input = preprocess_tokens(prefix, seqlen=SEQLEN)\n",
+ "\n",
+ " yield {\n",
+ " \"image\": np.asarray(image),\n",
+ " \"text\": np.asarray(tokens),\n",
+ " \"mask_ar\": np.asarray(mask_ar),\n",
+ " \"mask_input\": np.asarray(mask_input),\n",
+ " }\n",
+ " elif data_type == \"validation\":\n",
+ " suffix = example[\"suffix\"].decode().lower()\n",
+ " tokens, mask_ar, mask_loss, _ = preprocess_tokens(prefix, suffix, SEQLEN)\n",
+ "\n",
+ " yield {\n",
+ " \"image\": np.asarray(image),\n",
+ " \"text\": np.asarray(tokens),\n",
+ " \"mask_ar\": np.asarray(mask_ar),\n",
+ " \"mask_loss\": np.asarray(mask_loss),\n",
+ " }\n",
+ "\n",
+ "def render_example(image, description):\n",
+ " \"\"\"Renders an image with description in HTML format.\n",
+ "\n",
+ " Args:\n",
+ " image (np.array): Image data as numpy array.\n",
+ " description (str): Description text to display alongside the image.\n",
+ "\n",
+ " Returns:\n",
+ " str: HTML formatted string for displaying the image and description.\n",
+ " \"\"\"\n",
+ " image = ((image + 1) / 2 * 255).astype(np.uint8) # [-1,1] -> [0, 255]\n",
+ " return f\"\"\"\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " \"\"\"\n",
+ "\n",
+ "def display_comparisons(predictions, validations):\n",
+ " \"\"\"Displays side-by-side comparisons of predictions and validations.\n",
+ "\n",
+ " Args:\n",
+ " predictions (list): List of tuples (image, description) for predictions.\n",
+ " validations (list): List of tuples (image, description) for validations.\n",
+ "\n",
+ " Prints:\n",
+ " Displays HTML output showing 10 randomly selected images with predictions and validations.\n",
+ " \"\"\"\n",
+ " html_out = \"\"\n",
+ "\n",
+ " # Select 10 random indices\n",
+ " num_comparisons = min(10, len(predictions), len(validations))\n",
+ " random_indices = random.sample(range(min(len(predictions), len(validations))), num_comparisons)\n",
+ "\n",
+ " for random_index in random_indices:\n",
+ " pred_image, pred_description = predictions[random_index]\n",
+ " val_image, val_description = validations[random_index]\n",
+ "\n",
+ " # Call render_example to get image content with description\n",
+ " pred_content = render_example(pred_image, f\"Prediction: {pred_description}\")\n",
+ " val_content = render_example(val_image, f\"Validation: {val_description}\")\n",
+ "\n",
+ " # Structure container with three columns and set widths\n",
+ " html_out += f\"\"\"\n",
+ "
\n",
+ "
Prediction: {pred_description}
\n",
+ "
{pred_content}
\n",
+ "
Validation: {val_description}
\n",
+ "
\n",
+ " \"\"\"\n",
+ " display(HTML(html_out))\n",
+ "\n",
+ "\n",
+ "# Generate predictions and validations\n",
+ "predictions = []\n",
+ "for image, description_pred in make_predictions(data_iterator(\"prediction\"), batch_size=4):\n",
+ " # Only append if both image and description_pred are not None\n",
+ " if image is not None and description_pred is not None:\n",
+ " predictions.append((image, description_pred))\n",
+ "\n",
+ "validations = []\n",
+ "for example in data_iterator(\"validation\"):\n",
+ " description = postprocess_tokens(example[\"text\"])\n",
+ " description = description[len(\"describe en\\n\"):] # Strip prefix describe the image\n",
+ " # Only append if image and description are not None\n",
+ " if example[\"image\"] is not None and description is not None:\n",
+ " validations.append((example[\"image\"], description))\n",
+ "\n",
+ "# Display predictions and validations side-by-side\n",
+ "print(\"Predictions and Validation Examples: 10 Randomly Selected Images \\n\")\n",
+ "display_comparisons(predictions, validations)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "k-8GGxGMjqjt"
+ },
+ "source": [
+ "# Conclusion\n",
+ "\n",
+ "This notebook fine-tuned PaliGemma on 3,692 image-description pairs from a diverse bird species dataset. Using 23 curated species with text descriptions, the model generally produced accurate descriptions similar to validation data. However, improvements are needed for better species identification accuracy, suggesting potential benefits from expanding training to include more of the 525 available species."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "e5HltsZsjqjt"
+ },
+ "source": [
+ "# Reference\n",
+ "- [Fine-tune PaliGemma with JAX](https://www.kaggle.com/code/nilaychauhan/fine-tune-paligemma-with-jax)\n",
+ "- [PaliGemma model README](https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/README.md)\n",
+ "- [BIRDS 525 SPECIES- IMAGE CLASSIFICATION](https://www.kaggle.com/datasets/gpiosenka/100-bird-species)"
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "name": "Finetune_PaliGemma_for_image_description.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
![](\"https://www.tensorflow.org/images/colab_logo_32px.png\"){kind=logo}
Run in Google Colab\n",
+ "