Use dataset streaming, cleanup diagram (#497)

Use streaming to avoid saving entire Huggingface dataset to disk for large datasets. Updated diagram for clarity regarding client/server interaction. --------- Signed-off-by: Rishi Chandra <[email protected]>
NVIDIA · Feb 21, 2025 · efd0ed6 · efd0ed6
1 parent 13ff674
commit efd0ed6
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 23 deletions.
diff --git a/examples/ML+DL-Examples/Spark-DL/dl_inference/README.md b/examples/ML+DL-Examples/Spark-DL/dl_inference/README.md
@@ -39,22 +39,22 @@ In this simple case, the `predict_batch_fn` will use TensorFlow APIs to load the
 
 #### Notebook List
 
-Below is a full list of the notebooks with links to the examples they are based on. All notebooks have been saved with sample outputs for quick browsing.  
+Below is a full list of the notebooks and their links. All notebooks have been saved with sample outputs for quick browsing.  
 
 |   | Framework  | Notebook Name | Description | Link
 | ------------- | ------------- | ------------- | ------------- | ------------- 
-| 1 | HuggingFace | DeepSeek-R1 | LLM batch inference using the DeepSeek-R1-Distill-Llama reasoning model to solve word problems. | [Link](https://huggingface.co/deepseek-ai/DeepSeek-R1)
-| 2 | HuggingFace | Qwen-2.5-7b | LLM batch inference using the Qwen-2.5-7b model for text summarization. | [Link](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct)
-| 3 | HuggingFace | Gemma-7b | LLM batch inference using the Google Gemma-7b model for code comprehension tasks. | [Link](https://huggingface.co/google/gemma-7b-it)
-| 4 | HuggingFace | Sentence Transformers | Sentence embeddings using SentenceTransformers in Torch. | [Link](https://huggingface.co/sentence-transformers)
-| 5+6 | HuggingFace | Conditional Generation | Sentence translation using the T5 text-to-text transformer (Torch and Tensorflow). | [Link](https://huggingface.co/docs/transformers/model_doc/t5#t5) 
-| 7+8 | HuggingFace | Pipelines | Sentiment analysis using Huggingface pipelines (Torch and Tensorflow). | [Link](https://huggingface.co/docs/transformers/quicktour#pipeline-usage)
-| 9 | PyTorch | Image Classification | Training a model to predict clothing categories in FashionMNIST, and deploying with Torch-TensorRT accelerated inference. | [Link](https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html)
-| 10 | PyTorch | Housing Regression | Training and deploying a model to predict housing prices in the California Housing Dataset, and deploying with Torch-TensorRT accelerated inference. | [Link](https://github.com/christianversloot/machine-learning-articles/blob/main/how-to-create-a-neural-network-for-regression-with-pytorch.md)
-| 11 | Tensorflow | Image Classification | Training and deploying a model to predict hand-written digits in MNIST. | [Link](https://github.com/tensorflow/docs/blob/master/site/en/tutorials/keras/save_and_load.ipynb)
-| 12 | Tensorflow | Keras Preprocessing | Training and deploying a model with preprocessing layers to predict likelihood of pet adoption in the PetFinder mini dataset. | [Link](https://github.com/tensorflow/docs/blob/master/site/en/tutorials/structured_data/preprocessing_layers.ipynb)
-| 13 | Tensorflow | Keras Resnet50 | Deploying ResNet-50 to perform flower recognition from flower images. | [Link](https://docs.databricks.com/en/_extras/notebooks/source/deep-learning/keras-metadata.html)
-| 14 | Tensorflow | Text Classification | Training and deploying a model to perform sentiment analysis on the IMDB dataset. | [Link](https://github.com/tensorflow/docs/blob/master/site/en/tutorials/keras/text_classification.ipynb)
+| 1 | HuggingFace | DeepSeek-R1 | LLM batch inference using the DeepSeek-R1-Distill-Llama reasoning model to solve word problems. | [Link](huggingface/deepseek-r1_torch.ipynb)
+| 2 | HuggingFace | Qwen-2.5-7b | LLM batch inference using the Qwen-2.5-7b model for text summarization. | [Link](huggingface/qwen-2.5-7b_torch.ipynb)
+| 3 | HuggingFace | Gemma-7b | LLM batch inference using the Google Gemma-7b model for code comprehension tasks. | [Link](huggingface/gemma-7b_torch.ipynb)
+| 4 | HuggingFace | Sentence Transformers | Sentence embeddings using SentenceTransformers in Torch. | [Link](huggingface/sentence_transformers_torch.ipynb)
+| 5+6 | HuggingFace | Conditional Generation | Sentence translation using the T5 text-to-text transformer (Torch and Tensorflow). | [Torch Link](huggingface/conditional_generation_torch.ipynb), [TF Link](huggingface/conditional_generation_tf.ipynb)
+| 7+8 | HuggingFace | Pipelines | Sentiment analysis using Huggingface pipelines (Torch and Tensorflow). | [Torch Link](huggingface/pipelines_torch.ipynb), [TF Link](huggingface/pipelines_tf.ipynb)
+| 9 | PyTorch | Image Classification | Training a model to predict clothing categories in FashionMNIST, and deploying with Torch-TensorRT accelerated inference. | [Link](pytorch/image_classification_torch.ipynb)
+| 10 | PyTorch | Housing Regression | Training and deploying a model to predict housing prices in the California Housing Dataset, and deploying with Torch-TensorRT accelerated inference. | [Link](pytorch/housing_regression_torch.ipynb)
+| 11 | Tensorflow | Image Classification | Training and deploying a model to predict hand-written digits in MNIST. | [Link](tensorflow/image_classification_tf.ipynb)
+| 12 | Tensorflow | Keras Preprocessing | Training and deploying a model with preprocessing layers to predict likelihood of pet adoption in the PetFinder mini dataset. | [Link](tensorflow/keras_preprocessing_tf.ipynb)
+| 13 | Tensorflow | Keras Resnet50 | Deploying ResNet-50 to perform flower recognition from flower images. | [Link](tensorflow/keras_resnet50_tf.ipynb)
+| 14 | Tensorflow | Text Classification | Training and deploying a model to perform sentiment analysis on the IMDB dataset. | [Link](tensorflow/text_classification_tf.ipynb)
 
 
 ## Running Locally

diff --git a/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/deepseek-r1_torch.ipynb b/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/deepseek-r1_torch.ipynb
@@ -6,11 +6,11 @@
    "source": [
     "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
     "\n",
-    "# PySpark LLM Inference: DeepSeek-R1\n",
+    "# PySpark LLM Inference: DeepSeek-R1 Reasoning Q/A\n",
     "\n",
     "In this notebook, we demonstrate distributed batch inference with [DeepSeek-R1](https://github.com/deepseek-ai/DeepSeek-R1), using open weights on Huggingface.\n",
     "\n",
-    "We use [DeepSeek-R1-Distill-Llama-8B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-8B) as demonstration. DeepSeek's distilled models are based on open-source LLMs (such as Llama/Qwen), and are fine-tuned using samples generated by DeepSeek-R1 to perform multi-step reasoning tasks.\n",
+    "We use [DeepSeek-R1-Distill-Llama-8B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-8B) as demonstration. DeepSeek's distilled models are based on open-source LLMs (such as Llama/Qwen), and are fine-tuned using samples generated by DeepSeek-R1. We'll show how to use the model to reason through word problems.\n",
     "\n",
     "**Note:** Running this model on GPU with 16-bit precision requires **~18GB** of GPU RAM. Make sure your instances have sufficient GPU capacity."
    ]
@@ -261,6 +261,7 @@
    "outputs": [],
    "source": [
     "import os\n",
+    "import pandas as pd\n",
     "import datasets\n",
     "from datasets import load_dataset\n",
     "datasets.disable_progress_bars()"
@@ -330,7 +331,7 @@
    "source": [
     "#### Load DataFrame\n",
     "\n",
-    "Load the Orca Math Word Problems dataset from Huggingface and store in a Spark Dataframe."
+    "Load the first 500 samples of the [Orca Math Word Problems dataset](https://huggingface.co/datasets/microsoft/orca-math-word-problems-200k) from Huggingface and store in a Spark Dataframe."
    ]
   },
   {
@@ -339,8 +340,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dataset = load_dataset(\"microsoft/orca-math-word-problems-200k\", split=\"train[:1%]\")\n",
-    "dataset = dataset.to_pandas()[\"question\"]"
+    "dataset = load_dataset(\"microsoft/orca-math-word-problems-200k\", split=\"train\", streaming=True)\n",
+    "dataset = pd.Series([sample[\"question\"] for sample in dataset.take(500)])"
    ]
   },
   {

diff --git a/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/gemma-7b_torch.ipynb b/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/gemma-7b_torch.ipynb
@@ -6,11 +6,11 @@
    "source": [
     "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
     "\n",
-    "# PySpark LLM Inference: Gemma-7b\n",
+    "# PySpark LLM Inference: Gemma-7b Code Comprehension\n",
     "\n",
     "In this notebook, we demonstrate distributed inference with the Google [Gemma-7b-instruct](https://huggingface.co/google/gemma-7b-it) LLM, using open-weights on Huggingface.\n",
     "\n",
-    "The Gemma-7b-instruct is an instruction-fine-tuned version of the Gemma-7b base model.\n",
+    "The Gemma-7b-instruct is an instruction-fine-tuned version of the Gemma-7b base model. We'll show how to use the model to perform code comprehension tasks.\n",
     "\n",
     "**Note:** Running this model on GPU with 16-bit precision requires **~18 GB** of GPU RAM. Make sure your instances have sufficient GPU capacity."
    ]
@@ -200,6 +200,7 @@
    "outputs": [],
    "source": [
     "import os\n",
+    "import pandas as pd\n",
     "import datasets\n",
     "from datasets import load_dataset\n",
     "datasets.disable_progress_bars()"
@@ -269,7 +270,7 @@
    "source": [
     "#### Load DataFrame\n",
     "\n",
-    "Load the code comprehension dataset from Huggingface and store in a Spark Dataframe."
+    "Load the first 500 samples of the [Code Comprehension dataset](https://huggingface.co/datasets/imbue/code-comprehension) from Huggingface and store in a Spark Dataframe."
    ]
   },
   {
@@ -278,8 +279,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dataset = load_dataset(\"imbue/code-comprehension\", split=\"train[:1%]\")\n",
-    "dataset = dataset.to_pandas()[\"question\"]"
+    "dataset = load_dataset(\"imbue/code-comprehension\", split=\"train\", streaming=True)\n",
+    "dataset = pd.Series([sample[\"question\"] for sample in dataset.take(500)])"
    ]
   },
   {

diff --git a/examples/ML+DL-Examples/Spark-DL/dl_inference/images/spark-pytriton.png b/examples/ML+DL-Examples/Spark-DL/dl_inference/images/spark-pytriton.png