diff --git a/examples/ML+DL-Examples/Spark-DL/dl_inference/README.md b/examples/ML+DL-Examples/Spark-DL/dl_inference/README.md
index 451256c2..ea696c48 100644
--- a/examples/ML+DL-Examples/Spark-DL/dl_inference/README.md
+++ b/examples/ML+DL-Examples/Spark-DL/dl_inference/README.md
@@ -43,17 +43,18 @@ Below is a full list of the notebooks with links to the examples they are based
| | Framework | Notebook Name | Description | Link
| ------------- | ------------- | ------------- | ------------- | -------------
-| 1 | HuggingFace | DeepSeek-R1 | LLM batch inference using the DeepSeek-R1-Distill-Llama reasoning model. | [Link](https://huggingface.co/deepseek-ai/DeepSeek-R1)
-| 2 | HuggingFace | Gemma-7b | LLM batch inference using the lightweight Google Gemma-7b model. | [Link](https://huggingface.co/google/gemma-7b-it)
-| 3 | HuggingFace | Sentence Transformers | Sentence embeddings using SentenceTransformers in Torch. | [Link](https://huggingface.co/sentence-transformers)
-| 4+5 | HuggingFace | Conditional Generation | Sentence translation using the T5 text-to-text transformer for both Torch and Tensorflow. | [Link](https://huggingface.co/docs/transformers/model_doc/t5#t5)
-| 6+7 | HuggingFace | Pipelines | Sentiment analysis using Huggingface pipelines for both Torch and Tensorflow. | [Link](https://huggingface.co/docs/transformers/quicktour#pipeline-usage)
-| 8 | 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)
-| 9 | 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)
-| 10 | 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)
-| 11 | 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)
-| 12 | 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)
-| 13 | 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](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)
## Running Locally
diff --git a/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/qwen-2.5-7b_torch.ipynb b/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/qwen-2.5-7b_torch.ipynb
new file mode 100644
index 00000000..faf5c51d
--- /dev/null
+++ b/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/qwen-2.5-7b_torch.ipynb
@@ -0,0 +1,923 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "![](\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\")
\n",
+ "\n",
+ "# PySpark LLM Inference: Qwen-2.5 Text Summarization\n",
+ "\n",
+ "In this notebook, we demonstrate distributed batch inference with [Qwen-2.5](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct), using open weights on Huggingface.\n",
+ "\n",
+ "The Qwen-2.5-7b-instruct is an instruction-fine-tuned version of the Qwen-2.5-7b base model. We'll show how to use the model to perform text summarization.\n",
+ "\n",
+ "**Note:** Running this model on GPU with 16-bit precision requires **~16GB** of GPU RAM. Make sure your instances have sufficient GPU capacity."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "The dataset we'll use requires Zstandard compression."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 16,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Collecting zstandard\n",
+ " Downloading zstandard-0.23.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.0 kB)\n",
+ "Downloading zstandard-0.23.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.4 MB)\n",
+ "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.4/5.4 MB\u001b[0m \u001b[31m66.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+ "\u001b[?25hInstalling collected packages: zstandard\n",
+ "Successfully installed zstandard-0.23.0\n",
+ "Note: you may need to restart the kernel to use updated packages.\n"
+ ]
+ }
+ ],
+ "source": [
+ "%pip install zstandard"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Manually enable Huggingface tokenizer parallelism to avoid disabling with PySpark parallelism.\n",
+ "# See (https://github.com/huggingface/transformers/issues/5486) for more info. \n",
+ "import os\n",
+ "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Check the cluster environment to handle any platform-specific configurations."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "on_databricks = os.environ.get(\"DATABRICKS_RUNTIME_VERSION\", False)\n",
+ "on_dataproc = os.environ.get(\"DATAPROC_IMAGE_VERSION\", False)\n",
+ "on_standalone = not (on_databricks or on_dataproc)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "For cloud environments, set the huggingface cache dir to DBFS/GCS so that executors can load the model from cache."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "if on_databricks:\n",
+ " hf_home = \"/dbfs/FileStore/hf_home\"\n",
+ " dbutils.fs.mkdirs(hf_home)\n",
+ " os.environ[\"HF_HOME\"] = hf_home\n",
+ "elif on_dataproc:\n",
+ " hf_home = \"/mnt/gcs/hf_home\"\n",
+ " os.mkdir(hf_home) if not os.path.exists(hf_home) else None\n",
+ " os.environ[\"HF_HOME\"] = hf_home"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Warmup: Running locally\n",
+ "\n",
+ "**Note**: If the driver node does not have sufficient GPU capacity, proceed to the PySpark section."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 16,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "352b738e1a2442b0a997467aaf6eb0ad",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Loading checkpoint shards: 0%| | 0/4 [00:00system\\n\"),\n",
+ " lit(system_prompt),\n",
+ " lit(\"<|im_end|>\\n<|im_start|>user\\n\"),\n",
+ " col(\"value\"),\n",
+ " lit(\"<|im_end|>\\n<|im_start|>assistant\\n\")\n",
+ " ).alias(\"prompt\")\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "<|im_start|>system\n",
+ "You are a knowledgeable AI assistant. Your job is to create a 2-3 sentence summary \n",
+ "of a research abstract that captures the main objective, methodology, and key findings, using clear \n",
+ "language while preserving technical accuracy and quantitative results.<|im_end|>\n",
+ "<|im_start|>user\n",
+ "Epidemiology of hypoxaemia in children with acute lower respiratory infection.\n",
+ "To determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age. Systematic review of the published literature. Out-patient clinics, emergency departments and hospitalisation wards in 23 health centres from 10 countries. Cohort studies reporting the frequency of hypoxaemia in children under 5 years of age with ALRI, and the association between hypoxaemia and the risk of dying. Prevalence of hypoxaemia measured in children with ARI and relative risks for the association between the severity of illness and the frequency of hypoxaemia, and between hypoxaemia and the risk of dying. Seventeen published studies were found that included 4,021 children under 5 with acute respiratory infections (ARI) and reported the prevalence of hypoxaemia. Out-patient children and those with a clinical diagnosis of upper ARI had a low risk of hypoxaemia (pooled estimate of 6% to 9%). The prevalence increased to 31% and to 43% in patients in emergency departments and in cases with clinical pneumonia, respectively, and it was even higher among hospitalised children (47%) and in those with radiographically confirmed pneumonia (72%). The cumulated data also suggest that hypoxaemia is more frequent in children living at high altitude. Three papers reported an association between hypoxaemia and death, with relative risks varying between 1.4 and 4.6. Papers describing predictors of hypoxaemia have focused on clinical signs for detecting hypoxaemia rather than on identifying risk factors for developing this complication. Hypoxaemia is a common and potentially lethal complication of ALRI in children under 5, particularly among those with severe disease and those living at high altitude. Given the observed high prevalence of hypoxaemia and its likely association with increased mortality, efforts should be made to improve the detection of hypoxaemia and to provide oxygen earlier to more children with severe ALRI.<|im_end|>\n",
+ "<|im_start|>assistant\n",
+ "\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(df.take(1)[0].prompt)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data_path = \"spark-dl-datasets/pubmed_abstracts\"\n",
+ "if on_databricks:\n",
+ " dbutils.fs.mkdirs(\"/FileStore/spark-dl-datasets\")\n",
+ " data_path = \"dbfs:/FileStore/\" + data_path\n",
+ "\n",
+ "df.write.mode(\"overwrite\").parquet(data_path)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Using Triton Inference Server\n",
+ "In this section, we demonstrate integration with the [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server), an open-source, GPU-accelerated serving solution for DL. \n",
+ "We use [PyTriton](https://github.com/triton-inference-server/pytriton), a Flask-like framework that handles client/server communication with the Triton server. \n",
+ "\n",
+ "The process looks like this:\n",
+ "- Distribute a PyTriton task across the Spark cluster, instructing each node to launch a Triton server process.\n",
+ "- Define a Triton inference function, which contains a client that binds to the local server on a given node and sends inference requests.\n",
+ "- Wrap the Triton inference function in a predict_batch_udf to launch parallel inference requests using Spark.\n",
+ "- Finally, distribute a shutdown signal to terminate the Triton server processes on each node.\n",
+ "\n",
+ "![\"drawing\"](\"../images/spark-pytriton.png\")
"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from functools import partial"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Import the helper class from pytriton_utils.py:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 15,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sc.addPyFile(\"pytriton_utils.py\")\n",
+ "\n",
+ "from pytriton_utils import TritonServerManager"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Define the Triton Server function:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 16,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def triton_server(ports):\n",
+ " import time\n",
+ " import signal\n",
+ " import torch\n",
+ " import numpy as np\n",
+ " from pytriton.decorators import batch\n",
+ " from pytriton.model_config import DynamicBatcher, ModelConfig, Tensor\n",
+ " from pytriton.triton import Triton, TritonConfig\n",
+ " from pyspark import TaskContext\n",
+ " from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+ "\n",
+ " print(f\"SERVER: Initializing model on worker {TaskContext.get().partitionId()}.\")\n",
+ " model = AutoModelForCausalLM.from_pretrained(\n",
+ " \"Qwen/Qwen2.5-7B-Instruct\",\n",
+ " torch_dtype=torch.bfloat16,\n",
+ " device_map=\"auto\"\n",
+ " )\n",
+ " tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-7B-Instruct\", padding_side=\"left\")\n",
+ "\n",
+ " @batch\n",
+ " def _infer_fn(**inputs):\n",
+ " prompts = np.squeeze(inputs[\"prompts\"]).tolist()\n",
+ " print(f\"SERVER: Received batch of size {len(prompts)}\")\n",
+ " decoded_prompts = [p.decode(\"utf-8\") for p in prompts]\n",
+ " tokenized_inputs = tokenizer(decoded_prompts, padding=True, return_tensors=\"pt\").to(model.device)\n",
+ " generated_ids = model.generate(**tokenized_inputs, max_new_tokens=256)\n",
+ " outputs = tokenizer.batch_decode(generated_ids[:, tokenized_inputs.input_ids.shape[1]:], skip_special_tokens = True)\n",
+ " return {\n",
+ " \"outputs\": np.array(outputs).reshape(-1, 1)\n",
+ " }\n",
+ "\n",
+ " workspace_path = f\"/tmp/triton_{time.strftime('%m_%d_%M_%S')}\"\n",
+ " triton_conf = TritonConfig(http_port=ports[0], grpc_port=ports[1], metrics_port=ports[2])\n",
+ " with Triton(config=triton_conf, workspace=workspace_path) as triton:\n",
+ " triton.bind(\n",
+ " model_name=\"qwen-2.5\",\n",
+ " infer_func=_infer_fn,\n",
+ " inputs=[\n",
+ " Tensor(name=\"prompts\", dtype=object, shape=(-1,)),\n",
+ " ],\n",
+ " outputs=[\n",
+ " Tensor(name=\"outputs\", dtype=object, shape=(-1,)),\n",
+ " ],\n",
+ " config=ModelConfig(\n",
+ " max_batch_size=64,\n",
+ " batcher=DynamicBatcher(max_queue_delay_microseconds=5000), # 5ms\n",
+ " ),\n",
+ " strict=True,\n",
+ " )\n",
+ "\n",
+ " def _stop_triton(signum, frame):\n",
+ " print(\"SERVER: Received SIGTERM. Stopping Triton server.\")\n",
+ " triton.stop()\n",
+ "\n",
+ " signal.signal(signal.SIGTERM, _stop_triton)\n",
+ "\n",
+ " print(\"SERVER: Serving inference\")\n",
+ " triton.serve()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Start Triton servers"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "**Specify the number of nodes in the cluster.** \n",
+ "Following the README, the example standalone cluster uses 1 node. The example Databricks/Dataproc cluster scripts use 4 nodes by default. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 17,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Change based on cluster setup\n",
+ "num_nodes = 1 if on_standalone else 4"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "The `TritonClusterManager` will handle the lifecycle of Triton server instances across the Spark cluster:\n",
+ "- Find available ports for HTTP/gRPC/metrics\n",
+ "- Deploy a server on each node via stage-level scheduling\n",
+ "- Gracefully shutdown servers across nodes"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 18,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "model_name = \"qwen-2.5\"\n",
+ "server_manager = TritonServerManager(num_nodes=num_nodes, model_name=model_name)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 19,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "2025-02-16 11:49:25,237 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
+ "2025-02-16 11:49:25,239 - INFO - Starting 1 servers.\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ " \r"
+ ]
+ },
+ {
+ "data": {
+ "text/plain": [
+ "{'cb4ae00-lcedt': (3490378, [7000, 7001, 7002])}"
+ ]
+ },
+ "execution_count": 19,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "# Returns {'hostname', (server_pid, [http_port, grpc_port, metrics_port])}\n",
+ "server_manager.start_servers(triton_server, wait_retries=24) # allow up to 2 minutes for model loading"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Define client function"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Get the hostname -> url mapping from the server manager:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 20,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "host_to_grpc_url = server_manager.host_to_grpc_url # or server_manager.host_to_http_url"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 21,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def triton_fn(model_name, host_to_url):\n",
+ " import socket\n",
+ " import numpy as np\n",
+ " from pytriton.client import ModelClient\n",
+ "\n",
+ " url = host_to_url[socket.gethostname()]\n",
+ " print(f\"Connecting to Triton model {model_name} at {url}.\")\n",
+ "\n",
+ " def infer_batch(inputs):\n",
+ " with ModelClient(url, model_name, inference_timeout_s=500) as client:\n",
+ " flattened = np.squeeze(inputs).tolist()\n",
+ " # Encode batch\n",
+ " encoded_batch = [[text.encode(\"utf-8\")] for text in flattened]\n",
+ " encoded_batch_np = np.array(encoded_batch, dtype=np.bytes_)\n",
+ " # Run inference\n",
+ " result_data = client.infer_batch(encoded_batch_np)\n",
+ " result_data = np.squeeze(result_data[\"outputs\"], -1)\n",
+ " return result_data\n",
+ " \n",
+ " return infer_batch"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 22,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "generate = predict_batch_udf(partial(triton_fn, model_name=model_name, host_to_url=host_to_grpc_url),\n",
+ " return_type=StringType(),\n",
+ " input_tensor_shapes=[[1]],\n",
+ " batch_size=8)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Load DataFrame\n",
+ "\n",
+ "We'll parallelize over a small set of prompts for demonstration."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 23,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "df = spark.read.parquet(data_path).limit(64).repartition(8)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Run Inference"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 24,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "[Stage 10:=====================> (3 + 5) / 8]\r"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "CPU times: user 10.5 ms, sys: 6.63 ms, total: 17.1 ms\n",
+ "Wall time: 23.7 s\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ " \r"
+ ]
+ }
+ ],
+ "source": [
+ "%%time\n",
+ "# first pass caches model/fn\n",
+ "preds = df.withColumn(\"outputs\", generate(col(\"prompt\")))\n",
+ "results = preds.collect()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 25,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "[Stage 16:=====================> (3 + 5) / 8]\r"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "CPU times: user 8.1 ms, sys: 4.47 ms, total: 12.6 ms\n",
+ "Wall time: 21.7 s\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ " \r"
+ ]
+ }
+ ],
+ "source": [
+ "%%time\n",
+ "preds = df.withColumn(\"outputs\", generate(col(\"prompt\")))\n",
+ "results = preds.collect()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 28,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Q: <|im_start|>system\n",
+ "You are a knowledgeable AI assistant. Your job is to create a 2-3 sentence summary \n",
+ "of a research abstract that captures the main objective, methodology, and key findings, using clear \n",
+ "language while preserving technical accuracy and quantitative results.<|im_end|>\n",
+ "<|im_start|>user\n",
+ "Oral health promotion evaluation--time for development.\n",
+ "Increasing emphasis is now being placed upon the evaluation of health service interventions to demonstrate their effects. A series of effectiveness reviews of the oral health education and promotion literature has demonstrated that many of these interventions are poorly and inadequately evaluated. It is therefore difficult to determine the effectiveness of many interventions. Based upon developments from the field of health promotion research this paper explores options for improving the quality of oral health promotion evaluation. It is essential that the methods and measures used in the evaluation of oral health promotion are appropriate to the intervention. For many oral health promotion interventions clinical measures and methods of evaluation may not be appropriate. This paper outlines an evaluation framework which can be used to assess the range of effects of oral health promotion programmes. Improving the quality of oral health promotion evaluation is a shared responsibility between researchers and those involved in the provision of programmes. The provision of adequate resources and training are essential requirements for this to be successfully achieved.<|im_end|>\n",
+ "<|im_start|>assistant\n",
+ " \n",
+ "\n",
+ "A: This research aims to improve the evaluation of oral health promotion programs by developing an appropriate framework. It explores how methods and measures should align with the specific nature of these interventions, emphasizing that both researchers and program providers must collaborate to ensure adequate resources and training are available for high-quality evaluations. \n",
+ "\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(f\"Q: {results[0].prompt} \\n\")\n",
+ "print(f\"A: {results[0].outputs} \\n\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Shut down server on each executor"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 29,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "2025-02-16 11:51:42,365 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "2025-02-16 11:51:47,609 - INFO - Sucessfully stopped 1 servers. \n"
+ ]
+ },
+ {
+ "data": {
+ "text/plain": [
+ "[True]"
+ ]
+ },
+ "execution_count": 29,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "server_manager.stop_servers()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 30,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "if not on_databricks: # on databricks, spark.stop() puts the cluster in a bad state\n",
+ " spark.stop()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "spark-dl-torch",
+ "language": "python",
+ "name": "python3"
+ },
+ "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.11"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}