adapting the script multi_label_classification.py to be Backend-Agnostic (#2052)

* adapting the script multi_label_classification.py to be Backend-Agnostic

* addressing comments

Author: Humbulani1234

examples/nlp/ipynb/multi_label_classification.ipynb

@@ -10,7 +10,7 @@
     "\n",
     "**Author:** [Sayak Paul](https://twitter.com/RisingSayak), [Soumik Rakshit](https://github.com/soumik12345)<br>\n",
     "**Date created:** 2020/09/25<br>\n",
-    "**Last modified:** 2020/12/23<br>\n",
+    "**Last modified:** 2025/02/27<br>\n",
     "**Description:** Implementing a large-scale multi-label text classification model."
    ]
   },
@@ -49,19 +49,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
    "outputs": [],
    "source": [
-    "from tensorflow.keras import layers\n",
-    "from tensorflow import keras\n",
-    "import tensorflow as tf\n",
+    "import os\n",
+    "\n",
+    "os.environ[\"KERAS_BACKEND\"] = \"jax\"  # or tensorflow, or torch\n",
+    "\n",
+    "import keras\n",
+    "from keras import layers, ops\n",
     "\n",
     "from sklearn.model_selection import train_test_split\n",
-    "from ast import literal_eval\n",
     "\n",
+    "from ast import literal_eval\n",
     "import matplotlib.pyplot as plt\n",
     "import pandas as pd\n",
     "import numpy as np"
@@ -81,7 +84,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -106,7 +109,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -127,7 +130,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -148,7 +151,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -178,7 +181,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -203,7 +206,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -235,7 +238,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -275,14 +278,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
    "outputs": [],
    "source": [
+    "# For RaggedTensor\n",
+    "import tensorflow as tf\n",
+    "\n",
     "terms = tf.ragged.constant(train_df[\"terms\"].values)\n",
-    "lookup = tf.keras.layers.StringLookup(output_mode=\"multi_hot\")\n",
+    "lookup = layers.StringLookup(output_mode=\"multi_hot\")\n",
     "lookup.adapt(terms)\n",
     "vocab = lookup.get_vocabulary()\n",
     "\n",
@@ -294,7 +300,8 @@
     "\n",
     "\n",
     "print(\"Vocabulary:\\n\")\n",
-    "print(vocab)\n"
+    "print(vocab)\n",
+    ""
    ]
   },
   {
@@ -310,7 +317,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -337,7 +344,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -361,7 +368,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -380,7 +387,8 @@
     "        (dataframe[\"summaries\"].values, label_binarized)\n",
     "    )\n",
     "    dataset = dataset.shuffle(batch_size * 10) if is_train else dataset\n",
-    "    return dataset.batch(batch_size)\n"
+    "    return dataset.batch(batch_size)\n",
+    ""
    ]
   },
   {
@@ -394,7 +402,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -416,7 +424,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -450,7 +458,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -460,7 +468,8 @@
     "vocabulary = set()\n",
     "train_df[\"summaries\"].str.lower().str.split().apply(vocabulary.update)\n",
     "vocabulary_size = len(vocabulary)\n",
-    "print(vocabulary_size)\n"
+    "print(vocabulary_size)\n",
+    ""
    ]
   },
   {
@@ -475,7 +484,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -498,7 +507,8 @@
     ").prefetch(auto)\n",
     "test_dataset = test_dataset.map(\n",
     "    lambda text, label: (text_vectorizer(text), label), num_parallel_calls=auto\n",
-    ").prefetch(auto)\n"
+    ").prefetch(auto)\n",
+    ""
    ]
   },
   {
@@ -535,7 +545,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -550,7 +560,8 @@
     "            layers.Dense(lookup.vocabulary_size(), activation=\"sigmoid\"),\n",
     "        ]  # More on why \"sigmoid\" has been used here in a moment.\n",
     "    )\n",
-    "    return shallow_mlp_model\n"
+    "    return shallow_mlp_model\n",
+    ""
    ]
   },
   {
@@ -582,7 +593,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -635,7 +646,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
@@ -676,20 +687,40 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {
     "colab_type": "code"
    },
    "outputs": [],
    "source": [
-    "# Create a model for inference.\n",
-    "model_for_inference = keras.Sequential([text_vectorizer, shallow_mlp_model])\n",
     "\n",
-    "# Create a small dataset just for demoing inference.\n",
-    "inference_dataset = make_dataset(test_df.sample(100), is_train=False)\n",
+    "# We create a custom Model to override the predict method so\n",
+    "# that it first vectorizes text data\n",
+    "class ModelEndtoEnd(keras.Model):\n",
+    "\n",
+    "    def predict(self, inputs):\n",
+    "        indices = text_vectorizer(inputs)\n",
+    "        return super().predict(indices)\n",
+    "\n",
+    "\n",
+    "def get_inference_model(model):\n",
+    "    inputs = shallow_mlp_model.inputs\n",
+    "    outputs = shallow_mlp_model.outputs\n",
+    "    end_to_end_model = ModelEndtoEnd(inputs, outputs, name=\"end_to_end_model\")\n",
+    "    end_to_end_model.compile(\n",
+    "        optimizer=\"adam\", loss=\"binary_crossentropy\", metrics=[\"accuracy\"]\n",
+    "    )\n",
+    "    return end_to_end_model\n",
+    "\n",
+    "\n",
+    "model_for_inference = get_inference_model(shallow_mlp_model)\n",
+    "\n",
+    "# Create a small dataset just for demonstrating inference.\n",
+    "inference_dataset = make_dataset(test_df.sample(2), is_train=False)\n",
     "text_batch, label_batch = next(iter(inference_dataset))\n",
     "predicted_probabilities = model_for_inference.predict(text_batch)\n",
     "\n",
+    "\n",
     "# Perform inference.\n",
     "for i, text in enumerate(text_batch[:5]):\n",
     "    label = label_batch[i].numpy()[None, ...]\n",
@@ -731,16 +762,15 @@
     "tackle the multi-label binarization part and inverse-transforming the processed labels\n",
     "to the original form.\n",
     "\n",
-    "Thanks [Cingis Kratochvil](https://github.com/cumbalik) for suggesting and extending\n",
-    "this code example by the binary accuracy."
+    "Thanks to [Cingis Kratochvil](https://github.com/cumbalik) for suggesting and extending this code example by introducing binary accuracy as the evaluation metric."
    ]
   }
  ],
  "metadata": {
   "accelerator": "GPU",
   "colab": {
    "collapsed_sections": [],
-   "name": "..\\examples\\nlp\\multi_label_classification",
+   "name": "multi_label_classification",
    "private_outputs": false,
    "provenance": [],
    "toc_visible": true
@@ -765,4 +795,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}