examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/conditional_generation_torch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "8f6659b4-88da-4207-8d32-2674da5383a0",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
    "\n",
    "# PySpark DL Inference\n",
    "### Conditional generation with Huggingface\n",
    "\n",
    "In this notebook, we demonstrate distributed inference with the T5 transformer to perform sentence translation.  \n",
    "From: https://huggingface.co/docs/transformers/model_doc/t5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import T5Tokenizer, T5ForConditionalGeneration\n",
    "\n",
    "# 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": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n"
     ]
    }
   ],
   "source": [
    "tokenizer = T5Tokenizer.from_pretrained(\"google-t5/t5-small\")\n",
    "model = T5ForConditionalGeneration.from_pretrained(\"google-t5/t5-small\")\n",
    "\n",
    "task_prefix = \"translate English to German: \"\n",
    "\n",
    "lines = [\n",
    "    \"The house is wonderful\",\n",
    "    \"Welcome to NYC\",\n",
    "    \"HuggingFace is a company\"\n",
    "]\n",
    "\n",
    "input_sequences = [task_prefix + l for l in lines]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tokenizer(input_sequences,\n",
    "                      padding=True, \n",
    "                      return_tensors=\"pt\")\n",
    "\n",
    "outputs = model.generate(input_ids=inputs[\"input_ids\"], attention_mask=inputs[\"attention_mask\"], max_length=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Das Haus ist wunderbar',\n",
       " 'Willkommen in NYC',\n",
       " 'HuggingFace ist ein Unternehmen']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[tokenizer.decode(o, skip_special_tokens=True) for o in outputs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PySpark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1b8dae4a-3bfc-4430-b28a-7350db5efed4",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from pyspark.sql.types import *\n",
    "from pyspark import SparkConf\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import pandas_udf, col, struct\n",
    "from pyspark.ml.functions import predict_batch_udf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a93a1424-e483-4d37-a719-32fabee3f285",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import datasets\n",
    "from datasets import load_dataset\n",
    "datasets.disable_progress_bars()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the cluster environment to handle any platform-specific Spark configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": [
    "#### Create Spark Session\n",
    "\n",
    "For local standalone clusters, we'll connect to the cluster and create the Spark Session.  \n",
    "For CSP environments, Spark will either be preconfigured (Databricks) or we'll need to create the Spark Session (Dataproc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:34:55 WARN Utils: Your hostname, cb4ae00-lcedt resolves to a loopback address: 127.0.1.1; using 10.110.47.100 instead (on interface eno1)\n",
      "25/02/04 13:34:55 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "25/02/04 13:34:55 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    }
   ],
   "source": [
    "conf = SparkConf()\n",
    "\n",
    "if 'spark' not in globals():\n",
    "    if on_standalone:\n",
    "        import socket\n",
    "        conda_env = os.environ.get(\"CONDA_PREFIX\")\n",
    "        hostname = socket.gethostname()\n",
    "        conf.setMaster(f\"spark://{hostname}:7077\")\n",
    "        conf.set(\"spark.pyspark.python\", f\"{conda_env}/bin/python\")\n",
    "        conf.set(\"spark.pyspark.driver.python\", f\"{conda_env}/bin/python\")\n",
    "    elif on_dataproc:\n",
    "        conf.set(\"spark.executor.instances\", \"4\") # dataproc defaults to 2\n",
    "\n",
    "    conf.set(\"spark.executor.cores\", \"8\")\n",
    "    conf.set(\"spark.task.resource.gpu.amount\", \"0.125\")\n",
    "    conf.set(\"spark.executor.resource.gpu.amount\", \"1\")\n",
    "    conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", \"true\")\n",
    "    conf.set(\"spark.python.worker.reuse\", \"true\")\n",
    "\n",
    "conf.set(\"spark.sql.execution.arrow.maxRecordsPerBatch\", \"1000\")\n",
    "spark = SparkSession.builder.appName(\"spark-dl-examples\").config(conf=conf).getOrCreate()\n",
    "sc = spark.sparkContext"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f08c37a5-fb0c-45f6-8630-d2af67831641",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "Load the IMBD Movie Reviews dataset from Huggingface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f0ec30c9-365a-43c5-9c53-3497400ee548",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"imdb\", split=\"test\")\n",
    "dataset = dataset.to_pandas().drop(columns=\"label\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1e4269da-d2b3-46a5-9309-38a1ba825a47",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "#### Create PySpark DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "30dab34d-8e4b-4f30-b7c2-3dff49da018b",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StructType([StructField('text', StringType(), True)])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = spark.createDataFrame(dataset).repartition(8)\n",
    "df.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "55c33cc0-5dfb-449c-ae79-80972fb04405",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "25000"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "efd6d6d9-1c2c-4131-8df4-a3ef75c3fc57",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:35:02 WARN TaskSetManager: Stage 6 contains a task of very large size (4021 KiB). The maximum recommended task size is 1000 KiB.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Row(text=\"Anyone remember the first CKY, CKY2K etc..? Back when it was about making crazy cool stuff, rather than watching Bam Margera act like a douchebag, spoiled 5 year old, super/rock-star wannabe.<br /><br />The show used to be awesome, however, Bam's fame and wealth has led him to believe, that we now enjoy him acting childish and idiotic, more than actual cool stuff, that used to be in ex. CKY2K.<br /><br />The acts are so repetitive, there's like nothing new, except annoying stupidity and rehearsed comments... The only things we see is Bam Margera, so busy showing us how much he doesn't care, how much money he got or whatsoever.<br /><br />I really got nothing much left to say except, give us back CKY2K, cause Bam suck..<br /><br />I enjoy watching Steve-o, Knoxville etc. a thousand times more.\")]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.take(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "65a5b258-1634-441e-8b36-29777e54592d",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:35:02 WARN TaskSetManager: Stage 9 contains a task of very large size (4021 KiB). The maximum recommended task size is 1000 KiB.\n"
     ]
    }
   ],
   "source": [
    "data_path = \"spark-dl-datasets/imdb_test\"\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": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "89b909f4-5732-428b-ad61-9a6c5cf94df2",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "#### Load and preprocess DataFrame\n",
    "\n",
    "Define our preprocess function. We'll take the first sentence from each sample as our input for translation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "eb7e53d6-bbd0-48d2-a3be-36847275e2a9",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def preprocess(text: pd.Series, prefix: str = \"\") -> pd.Series:\n",
    "    @pandas_udf(\"string\")\n",
    "    def _preprocess(text: pd.Series) -> pd.Series:\n",
    "        return pd.Series([prefix + s.split(\".\")[0] for s in text])\n",
    "    return _preprocess(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "97eee1a4-9dc4-43b0-9578-6d7f8ff338bd",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------+\n",
      "|                                                                                                text|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "|The only reason I'm even giving this movie a 4 is because it was made in to an episode of Mystery...|\n",
      "|Awkward disaster mishmash has a team of scavengers coming across the overturned S.S. Poseidon, ho...|\n",
      "|Here is a fantastic concept for a film - a series of meteors crash into a small town and the resu...|\n",
      "|I walked out of the cinema having suffered this film after 30 mins. I left two friends pinned in ...|\n",
      "|A wildly uneven film where the major problem is the uneasy mix of comedy and thriller. To me, the...|\n",
      "|Leonard Rossiter and Frances de la Tour carry this film, not without a struggle, as the script wa...|\n",
      "|A good cast... A good idea but turns out it is flawed as hypnosis is not allowed as evidence in c...|\n",
      "|Yet again, I appear to be the only person on planet Earth who is capable of criticizing Japanese ...|\n",
      "|As a serious horror fan, I get that certain marketing ploys are used to sell movies, especially t...|\n",
      "|Upon writing this review I have difficulty trying to think of what to write about. Nothing much h...|\n",
      "|Simply awful. I'm including a spoiler warning here only because of including a coupla jokes from ...|\n",
      "|I am a fan of Ed Harris' work and I really had high expectations about this film. Having so good ...|\n",
      "|Well...I like Patricia Kaas. She is a beautiful lady and an extremely gifted and versatile singer...|\n",
      "|This is a new approach to comedy. It isn't funny.<br /><br />The joke is that this, in and of its...|\n",
      "|It's been mentioned by others the inane dialogue in this series and I agree.<br /><br />If Mom an...|\n",
      "|One of the most boring movies I've ever had to sit through, it's completely formulaic. Just a coo...|\n",
      "|This movie was playing on Lifetime Movie Network last month and I decided to check it out. I watc...|\n",
      "|1983's \"Frightmare\" is an odd little film. The director seems to be trying to combine the atmosph...|\n",
      "|'Felony' is a B-movie. No doubt about it.<br /><br />Of course, if you take a look at the cast li...|\n",
      "|This movie defines the word \"confused\". All the actors stay true to the script. More's the pity, ...|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Limit to N rows, since this can be slow\n",
    "df = spark.read.parquet(data_path).limit(512).repartition(8)\n",
    "df.show(truncate=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Append a prefix to tell the model to translate English to French:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fa14304d-b409-4d07-99ef-9da7c7c76158",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------+\n",
      "|                                                                                               input|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "|translate English to French: The only reason I'm even giving this movie a 4 is because it was mad...|\n",
      "|translate English to French: Awkward disaster mishmash has a team of scavengers coming across the...|\n",
      "|translate English to French: Here is a fantastic concept for a film - a series of meteors crash i...|\n",
      "|     translate English to French: I walked out of the cinema having suffered this film after 30 mins|\n",
      "|translate English to French: A wildly uneven film where the major problem is the uneasy mix of co...|\n",
      "|translate English to French: Leonard Rossiter and Frances de la Tour carry this film, not without...|\n",
      "|                                                            translate English to French: A good cast|\n",
      "|translate English to French: Yet again, I appear to be the only person on planet Earth who is cap...|\n",
      "|translate English to French: As a serious horror fan, I get that certain marketing ploys are used...|\n",
      "|translate English to French: Upon writing this review I have difficulty trying to think of what t...|\n",
      "|                                                           translate English to French: Simply awful|\n",
      "|translate English to French: I am a fan of Ed Harris' work and I really had high expectations abo...|\n",
      "|                                                                   translate English to French: Well|\n",
      "|                                       translate English to French: This is a new approach to comedy|\n",
      "|translate English to French: It's been mentioned by others the inane dialogue in this series and ...|\n",
      "|translate English to French: One of the most boring movies I've ever had to sit through, it's com...|\n",
      "|translate English to French: This movie was playing on Lifetime Movie Network last month and I de...|\n",
      "|                              translate English to French: 1983's \"Frightmare\" is an odd little film|\n",
      "|                                                  translate English to French: 'Felony' is a B-movie|\n",
      "|                                 translate English to French: This movie defines the word \"confused\"|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "input_df = df.select(preprocess(col(\"text\"), \"translate English to French: \").alias(\"input\")).cache()\n",
    "input_df.show(truncate=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bc9cbdd2-1ca6-48e4-a549-792b3726525b",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Inference using Spark DL API\n",
    "\n",
    "Distributed inference using the PySpark [predict_batch_udf](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.functions.predict_batch_udf.html#pyspark.ml.functions.predict_batch_udf):\n",
    "\n",
    "- predict_batch_fn uses PyTorch APIs to load the model and return a predict function which operates on numpy arrays \n",
    "- predict_batch_udf will convert the Spark DataFrame columns into numpy input batches for the predict function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "adb81177-442d-42ab-b86d-d8792201b4c8",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def predict_batch_fn():\n",
    "    import numpy as np\n",
    "    import torch\n",
    "    from transformers import T5ForConditionalGeneration, T5Tokenizer\n",
    "    from pyspark import TaskContext\n",
    "\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    print(f\"Initializing model on worker {TaskContext.get().partitionId()}, device {device}.\")\n",
    "    model = T5ForConditionalGeneration.from_pretrained(\"t5-small\").to(device)\n",
    "    tokenizer = T5Tokenizer.from_pretrained(\"t5-small\")\n",
    "\n",
    "    def predict(inputs):\n",
    "        flattened = np.squeeze(inputs).tolist()\n",
    "        inputs = tokenizer(flattened, \n",
    "                           padding=True,\n",
    "                           return_tensors=\"pt\").to(device)\n",
    "        outputs = model.generate(input_ids=inputs[\"input_ids\"],\n",
    "                                 attention_mask=inputs[\"attention_mask\"],\n",
    "                                 max_length=128)\n",
    "        string_outputs = np.array([tokenizer.decode(o, skip_special_tokens=True) for o in outputs])\n",
    "        print(\"predict: {}\".format(len(flattened)))\n",
    "        return string_outputs\n",
    "    \n",
    "    return predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "20aab3a1-2284-4c07-9ce1-a20cf54d88f3",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "generate = predict_batch_udf(predict_batch_fn,\n",
    "                             return_type=StringType(),\n",
    "                             batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a8d6f48e-09e7-4fc7-9d2f-1b68bc2976a7",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 24:=============================>                            (4 + 4) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 10.2 ms, sys: 5.05 ms, total: 15.2 ms\n",
      "Wall time: 7.41 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "preds = input_df.withColumn(\"preds\", generate(struct(\"input\")))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "abe2271d-0077-48f6-98b1-93524dd86447",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 27:=============================>                            (4 + 4) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.93 ms, sys: 1.98 ms, total: 5.91 ms\n",
      "Wall time: 4.08 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df.withColumn(\"preds\", generate(\"input\"))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "77623711-a742-4262-8839-16fc3ddd1af7",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 30:==============>                                           (2 + 6) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.85 ms, sys: 1.75 ms, total: 5.6 ms\n",
      "Wall time: 4.08 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df.withColumn(\"preds\", generate(col(\"input\")))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f339c654-52fd-4992-b054-188dfb260e5d",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|                                             input|                                             preds|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|translate English to French: The only reason I'...|La seule raison pour laquelle je donne même ce ...|\n",
      "|translate English to French: Awkward disaster m...|La mishmash d’Awkward a eu une équipe de scaven...|\n",
      "|translate English to French: Here is a fantasti...|Voici un concept fantastique pour un film : une...|\n",
      "|translate English to French: I walked out of th...|Je me suis rendu du cinéma après avoir subi ce ...|\n",
      "|translate English to French: A wildly uneven fi...|Un film extrêmement inégal où le problème majeu...|\n",
      "|translate English to French: Leonard Rossiter a...|Leonard Rossiter et Frances de la Tour mettent ...|\n",
      "|          translate English to French: A good cast|                                  Une bonne étoile|\n",
      "|translate English to French: Yet again, I appea...|Encore une fois, je semble être la seule person...|\n",
      "|translate English to French: As a serious horro...|En tant que grand fan d'horreur, je peux obteni...|\n",
      "|translate English to French: Upon writing this ...|la suite de cette étude, j'ai de la difficulté ...|\n",
      "|         translate English to French: Simply awful|                          Tout simplement terrible|\n",
      "|translate English to French: I am a fan of Ed H...|Je suis un fan de l'oeuvre d'Ed Harris et j'ai ...|\n",
      "|                 translate English to French: Well|                                           Eh bien|\n",
      "|translate English to French: This is a new appr...|  Il s’agit d’une nouvelle approche de la comédie.|\n",
      "|translate English to French: It's been mentione...|Il a été mentionné par d'autres le dialogue ina...|\n",
      "|translate English to French: One of the most bo...|Un des films les plus ennuyeux que je n'ai jama...|\n",
      "|translate English to French: This movie was pla...|Ce film jouait sur Lifetime Movie Network le mo...|\n",
      "|translate English to French: 1983's \"Frightmare...|Le film \"Frightmare\" de 1983 est un petit film ...|\n",
      "|translate English to French: 'Felony' is a B-movie|                       'Felony' est un mouvement B|\n",
      "|translate English to French: This movie defines...|                   Ce film définit le mot «confus»|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "preds.show(truncate=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try English to German:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_df2 = df.select(preprocess(col(\"text\"), \"translate English to German: \").alias(\"input\")).cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 36:==================================================>       (7 + 1) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.02 ms, sys: 705 μs, total: 6.73 ms\n",
      "Wall time: 4.24 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "preds = input_df2.withColumn(\"preds\", generate(struct(\"input\")))\n",
    "result = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 39:==============>                                           (2 + 6) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.12 ms, sys: 319 μs, total: 6.43 ms\n",
      "Wall time: 3.88 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df2.withColumn(\"preds\", generate(\"input\"))\n",
    "result = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 42:==============>                                           (2 + 6) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.03 ms, sys: 16 μs, total: 7.05 ms\n",
      "Wall time: 3.9 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df2.withColumn(\"preds\", generate(col(\"input\")))\n",
    "result = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|                                             input|                                             preds|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|translate English to German: The only reason I'...|Der einzige Grund, warum ich sogar diesen Film ...|\n",
      "|translate English to German: Awkward disaster m...|Awkward-Katastrophenmischmash hat ein Team von ...|\n",
      "|translate English to German: Here is a fantasti...|Hier ist ein fantastisches Konzept für einen Fi...|\n",
      "|translate English to German: I walked out of th...|Ich ging aus dem Kino, nachdem ich diesen Film ...|\n",
      "|translate English to German: A wildly uneven fi...|Ein völlig ungleicher Film, in dem das Hauptpro...|\n",
      "|translate English to German: Leonard Rossiter a...|Leonard Rossiter und Frances de la Tour tragen ...|\n",
      "|          translate English to German: A good cast|                                     Gutes Casting|\n",
      "|translate English to German: Yet again, I appea...|Ich scheine wieder einmal die einzige Person au...|\n",
      "|translate English to German: As a serious horro...|Als ernsthafter Horrorfan erhalte ich, dass bes...|\n",
      "|translate English to German: Upon writing this ...|Ich habe Schwierigkeiten, mich an die Regeln zu...|\n",
      "|         translate English to German: Simply awful|                               Einfach schrecklich|\n",
      "|translate English to German: I am a fan of Ed H...|Ich bin ein Fan von Ed Harris' Arbeit und hatte...|\n",
      "|                 translate English to German: Well|                                               Nun|\n",
      "|translate English to German: This is a new appr...|          Das ist ein neuer Ansatz für die Komödie|\n",
      "|translate English to German: It's been mentione...|Es wurde von anderen erwähnt, die unangenehme D...|\n",
      "|translate English to German: One of the most bo...|Einer der langwierigen Filme, die ich jemals du...|\n",
      "|translate English to German: This movie was pla...|Dieser Film spielte im letzten Monat auf Lifeti...|\n",
      "|translate English to German: 1983's \"Frightmare...|       1983 ist \"Frightmare\" ein merkwürdiger Film|\n",
      "|translate English to German: 'Felony' is a B-movie|                           'Felony' ist ein B-Film|\n",
      "|translate English to German: This movie defines...|         Dieser Film definiert das Wort \"verwirrt\"|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "preds.show(truncate=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a79a6f3a-cc34-46a4-aadd-16870423fffa",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "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",
    "<img src=\"../images/spark-pytriton.png\" alt=\"drawing\" width=\"700\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1e73757e-a451-4835-98e0-257ccf7a9025",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import the helper class from pytriton_utils.py:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "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": 30,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "71b1cb49-3d8f-4eeb-937a-c0c334bd2947",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def triton_server(ports):\n",
    "    import time\n",
    "    import signal\n",
    "    import numpy as np\n",
    "    import torch\n",
    "    from transformers import T5Tokenizer, T5ForConditionalGeneration\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",
    "\n",
    "    print(f\"SERVER: Initializing Conditional Generation model on worker {TaskContext.get().partitionId()}.\")\n",
    "    tokenizer = T5Tokenizer.from_pretrained(\"t5-small\")\n",
    "    model = T5ForConditionalGeneration.from_pretrained(\"t5-small\")\n",
    "    \n",
    "    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    print(f\"SERVER: Using {DEVICE} device.\")\n",
    "    model = model.to(DEVICE)\n",
    "\n",
    "    @batch\n",
    "    def _infer_fn(**inputs):\n",
    "        sentences = np.squeeze(inputs[\"text\"]).tolist()\n",
    "        print(f\"SERVER: Received batch of size {len(sentences)}\")\n",
    "        decoded_sentences = [s.decode(\"utf-8\") for s in sentences]\n",
    "        inputs = tokenizer(decoded_sentences,\n",
    "                        padding=True,\n",
    "                        return_tensors=\"pt\").to(DEVICE)\n",
    "        output_ids = model.generate(input_ids=inputs[\"input_ids\"],\n",
    "                                    attention_mask=inputs[\"attention_mask\"],\n",
    "                                    max_length=128)\n",
    "        outputs = np.array([[tokenizer.decode(o, skip_special_tokens=True)] for o in output_ids])\n",
    "        return {\n",
    "            \"translations\": outputs,\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=\"ConditionalGeneration\",\n",
    "            infer_func=_infer_fn,\n",
    "            inputs=[\n",
    "                Tensor(name=\"text\", dtype=object, shape=(-1,)),\n",
    "            ],\n",
    "            outputs=[\n",
    "                Tensor(name=\"translations\", 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",
    "            # The server manager sends SIGTERM to stop the server; this function ensures graceful cleanup.\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": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1bf14846-15a3-4bc8-b0c5-ce71680d3550",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "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": 31,
   "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": 32,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "5bf1fafc-d9c9-4fd7-901d-da97cf4ff496",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "model_name = \"ConditionalGeneration\"\n",
    "server_manager = TritonServerManager(num_nodes=num_nodes, model_name=model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-07 11:03:44,809 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-07 11:03:44,810 - INFO - Starting 1 servers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'cb4ae00-lcedt': (2020631, [7000, 7001, 7002])}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Returns {'hostname', (server_pid, [http_port, grpc_port, metrics_port])}\n",
    "server_manager.start_servers(triton_server)"
   ]
  },
  {
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "host_to_http_url = server_manager.host_to_http_url  # or server_manager.host_to_grpc_url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the Triton inference function, which returns a predict function for batch inference through the server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e203eb19-166d-4177-aa87-fd31b7e3c90e",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "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=240) 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[\"translations\"], -1)\n",
    "            return result_data\n",
    "        \n",
    "    return infer_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "be692f4a-cf86-4cf4-9530-7c62e479cacd",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "generate = predict_batch_udf(partial(triton_fn, model_name=model_name, host_to_url=host_to_http_url),\n",
    "                             return_type=StringType(),\n",
    "                             input_tensor_shapes=[[1]],\n",
    "                             batch_size=32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1b6b2a05-aea4-4e4d-a87d-0a6bd5ab554c",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "#### Load and preprocess DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a5e83230-5178-4fec-bba2-0e69be40e68c",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def preprocess(text: pd.Series, prefix: str = \"\") -> pd.Series:\n",
    "    @pandas_udf(\"string\")\n",
    "    def _preprocess(text: pd.Series) -> pd.Series:\n",
    "        return pd.Series([prefix + s.split(\".\")[0] for s in text])\n",
    "    return _preprocess(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "aad299b0-34bb-4edb-b1e4-cd0c82bb7455",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "df = spark.read.parquet(data_path).limit(512).repartition(8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7934a6fc-57bc-4104-a52c-076351e77cbe",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:35:39 WARN CacheManager: Asked to cache already cached data.\n"
     ]
    }
   ],
   "source": [
    "input_df = df.select(preprocess(col(\"text\"), \"translate English to French: \").alias(\"input\")).cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "0f6229ef-01c8-43c9-a259-c5df6a18d689",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 51:====================================>                     (5 + 3) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.09 ms, sys: 4.41 ms, total: 9.5 ms\n",
      "Wall time: 4.96 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "preds = input_df.withColumn(\"preds\", generate(struct(\"input\")))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "5a543b4c-8b29-4f61-9773-2639bbc7f728",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 54:===========================================>              (6 + 2) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.4 ms, sys: 1.12 ms, total: 6.52 ms\n",
      "Wall time: 4.41 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df.withColumn(\"preds\", generate(\"input\"))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "4c0cfc4e-ef0a-435e-9fdf-72b72b6def93",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 57:===========================================>              (6 + 2) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.59 ms, sys: 1.79 ms, total: 6.38 ms\n",
      "Wall time: 4.55 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df.withColumn(\"preds\", generate(col(\"input\")))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2d756e2e-8b60-43cb-b5f9-e27de11be24d",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|                                             input|                                             preds|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|translate English to French: The only reason I'...|La seule raison pour laquelle je donne même ce ...|\n",
      "|translate English to French: Awkward disaster m...|La mishmash d’Awkward a eu une équipe de scaven...|\n",
      "|translate English to French: Here is a fantasti...|Voici un concept fantastique pour un film : une...|\n",
      "|translate English to French: I walked out of th...|Je me suis rendu du cinéma après avoir subi ce ...|\n",
      "|translate English to French: A wildly uneven fi...|Un film extrêmement inégal où le problème majeu...|\n",
      "|translate English to French: Leonard Rossiter a...|Leonard Rossiter et Frances de la Tour mettent ...|\n",
      "|          translate English to French: A good cast|                                  Une bonne étoile|\n",
      "|translate English to French: Yet again, I appea...|Encore une fois, je semble être la seule person...|\n",
      "|translate English to French: As a serious horro...|En tant que grand fan d'horreur, je peux obteni...|\n",
      "|translate English to French: Upon writing this ...|la suite de cette étude, j'ai de la difficulté ...|\n",
      "|         translate English to French: Simply awful|                          Tout simplement terrible|\n",
      "|translate English to French: I am a fan of Ed H...|Je suis un fan de l'oeuvre d'Ed Harris et j'ai ...|\n",
      "|                 translate English to French: Well|                                           Eh bien|\n",
      "|translate English to French: This is a new appr...|  Il s’agit d’une nouvelle approche de la comédie.|\n",
      "|translate English to French: It's been mentione...|Il a été mentionné par d'autres le dialogue ina...|\n",
      "|translate English to French: One of the most bo...|Un des films les plus ennuyeux que je n'ai jama...|\n",
      "|translate English to French: This movie was pla...|Ce film jouait sur Lifetime Movie Network le mo...|\n",
      "|translate English to French: 1983's \"Frightmare...|Le film \"Frightmare\" de 1983 est un petit film ...|\n",
      "|translate English to French: 'Felony' is a B-movie|                       'Felony' est un mouvement B|\n",
      "|translate English to French: This movie defines...|                   Ce film définit le mot «confus»|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "preds.show(truncate=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "86ae68d4-57da-41d9-91b4-625ef9465d60",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "#### Shut down servers on each executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:35:53,794 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-04 13:35:58,983 - INFO - Sucessfully stopped 1 servers.                 \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[True]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "server_manager.stop_servers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "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": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "008c3e50-d321-4431-a9ab-919b35d1b042",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "environmentMetadata": null,
   "language": "python",
   "notebookMetadata": {
    "mostRecentlyExecutedCommandWithImplicitDF": {
     "commandId": 421988607303514,
     "dataframes": [
      "_sqldf"
     ]
    },
    "pythonIndentUnit": 4
   },
   "notebookName": "spark-triton-db.ipynb",
   "widgets": {}
  },
  "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": 4
}