Add Triton cell tags

Signed-off-by: Rishi Chandra <[email protected]>
NVIDIA · Oct 22, 2024 · 0fe696e · 0fe696e
1 parent 4940bf8
commit 0fe696e
Show file tree

Hide file tree

Showing 11 changed files with 734 additions and 151 deletions.
diff --git a/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/conditional_generation_tf.ipynb b/examples/ML+DL-Examples/Spark-DL/dl_inference/huggingface/conditional_generation_tf.ipynb
@@ -1039,7 +1039,11 @@
    "cell_type": "code",
    "execution_count": 35,
    "id": "b858cf85-82e6-41ef-905b-d8c5d6fea492",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [],
    "source": [
     "import os"
@@ -1049,7 +1053,11 @@
    "cell_type": "code",
    "execution_count": 36,
    "id": "05ce7c77-d562-45e8-89bb-cd656aba5a5f",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [],
    "source": [
     "%%bash\n",
@@ -1076,7 +1084,11 @@
    "cell_type": "code",
    "execution_count": 37,
    "id": "afd00b7e-8150-4c95-a2e4-037e9c90f92a",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1160,7 +1172,9 @@
    "execution_count": 38,
    "id": "1a997c33-5202-466d-8304-b8c30f32978f",
    "metadata": {
-    "tags": []
+    "tags": [
+     "TRITON"
+    ]
    },
    "outputs": [],
    "source": [
@@ -1175,7 +1189,11 @@
    "cell_type": "code",
    "execution_count": 39,
    "id": "9dea1875-6b95-4fc0-926d-a625a441b33d",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [],
    "source": [
     "# only use first N examples, since this is slow\n",
@@ -1186,7 +1204,11 @@
    "cell_type": "code",
    "execution_count": 40,
    "id": "5d6c54e7-534d-406f-b8e6-fd592efd0ab2",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [],
    "source": [
     "# only use first sentence and add prefix for conditional generation\n",
@@ -1201,7 +1223,11 @@
    "cell_type": "code",
    "execution_count": 41,
    "id": "dc1bbbe3-4232-49e5-80f6-99976524b73b",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [],
    "source": [
     "# only use first 100 rows, since generation takes a while\n",
@@ -1212,7 +1238,11 @@
    "cell_type": "code",
    "execution_count": 42,
    "id": "5d10c61c-6102-4d19-8dd6-0c7b5b65343e",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stdout",
@@ -1255,7 +1285,11 @@
    "cell_type": "code",
    "execution_count": 43,
    "id": "2e0907da-a5d9-4c3b-9db4-ce5e70ca9bb4",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [],
    "source": [
     "def triton_fn(triton_uri, model_name):\n",
@@ -1305,7 +1339,11 @@
    "cell_type": "code",
    "execution_count": 44,
    "id": "9308bdd7-6f67-484d-8b51-dd1e1b2960ba",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [],
    "source": [
     "generate = predict_batch_udf(partial(triton_fn, triton_uri=\"localhost:8001\", model_name=\"hf_generation_tf\"),\n",
@@ -1318,7 +1356,11 @@
    "cell_type": "code",
    "execution_count": 45,
    "id": "38484ffd-370d-492b-8ca4-9eff9f242a9f",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1354,7 +1396,11 @@
    "cell_type": "code",
    "execution_count": 46,
    "id": "ebcb6699-3ac2-4529-ab0f-fab0a5e792da",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1389,7 +1435,11 @@
    "cell_type": "code",
    "execution_count": 47,
    "id": "e2ed18ad-d00b-472c-b2c3-047932f2105d",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1424,7 +1474,11 @@
    "cell_type": "code",
    "execution_count": 48,
    "id": "0cd64a1c-beb8-47d5-ac6f-e8525bb61176",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1481,7 +1535,11 @@
    "cell_type": "code",
    "execution_count": 49,
    "id": "af70fed8-0f2b-4ea7-841c-476afdf9b1c0",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1500,7 +1558,11 @@
    "cell_type": "code",
    "execution_count": 50,
    "id": "ef075e10-e22c-4236-9e0b-cb47cf2d3d06",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stdout",
@@ -1543,7 +1605,11 @@
    "cell_type": "code",
    "execution_count": 51,
    "id": "2e7e4af8-b815-4375-b851-8368309ee8e1",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1578,7 +1644,11 @@
    "cell_type": "code",
    "execution_count": 52,
    "id": "7b0aefb0-a96b-4791-a23c-1ce9b24eb20c",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1613,7 +1683,11 @@
    "cell_type": "code",
    "execution_count": 53,
    "id": "1214b75b-a373-4579-b4c6-0cb8627da776",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1648,7 +1722,11 @@
    "cell_type": "code",
    "execution_count": 54,
    "id": "c9dbd21f-9e37-4221-b765-80ba8c80b884",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -1715,7 +1793,11 @@
    "cell_type": "code",
    "execution_count": 55,
    "id": "425d3b28-7705-45ba-8a18-ad34fc895219",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "TRITON"
+    ]
+   },
    "outputs": [
     {
      "name": "stderr",