Adding LLM test notebook

clinical-data-mining · Nov 25, 2024 · 314b6aa · 314b6aa
1 parent 3bb1f7b
commit 314b6aa
Show file tree

Hide file tree

Showing 2 changed files with 163 additions and 0 deletions.
diff --git a/docs/reference/images/compute_cluster.png b/docs/reference/images/compute_cluster.png
diff --git a/examples/databricks-llama-test.ipynb b/examples/databricks-llama-test.ipynb
@@ -0,0 +1,163 @@
+{
+ "cells": [
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "# Databricks LLM Test\n",
+    "This test notebook:\n",
+    "- Creates a client object with Databricks using a user token\n",
+    "- Prompts an LLM of a user's choice (`SERVING_MODEL`)\n",
+    "- Returns the content of the message\n",
+    "\n",
+    "Notes: \n",
+    "- Compute cluster info: `CDSI ML Cluster`\n",
+    "- The `SERVING_MODEL` must be a serving endpoint in Databricks. This process in done in the [Databricks UI](https://msk-mode-test.cloud.databricks.com/ml/endpoints/)\n",
+    "- When testing is completed, MAKE SURE THE `CDSI ML Cluster` IS NO LONGER RUNNING. Unexpected costs will result even if cluster is idle!\n",
+    "\n",
+    "!['Compute cluster'](../docs/reference/images/compute_cluster.png)\n",
+    "    "
+   ],
+   "id": "5f0f96e52ae54ff5"
+  },
+  {
+   "cell_type": "code",
+   "id": "1660925dae63fc66",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-25T23:39:55.269786Z",
+     "start_time": "2024-11-25T23:39:55.266304Z"
+    }
+   },
+   "source": "from openai import OpenAI\n",
+   "outputs": [],
+   "execution_count": 6
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-25T23:39:55.361127Z",
+     "start_time": "2024-11-25T23:39:55.358009Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "DATABRICKS_TOKEN = '<YOUR-TOKEN-HERE>'\n",
+    "SERVING_MODEL = \"meta_llama_3_8b_instruct_cdm\"\n",
+    "MAX_TOKENS = 256\n"
+   ],
+   "id": "f841f23f1f9b3923",
+   "outputs": [],
+   "execution_count": 7
+  },
+  {
+   "metadata": {
+    "collapsed": true,
+    "ExecuteTime": {
+     "end_time": "2024-11-25T23:39:55.479054Z",
+     "start_time": "2024-11-25T23:39:55.465727Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "# Create client with Databricks\n",
+    "client = OpenAI(\n",
+    "    api_key=DATABRICKS_TOKEN,\n",
+    "    base_url=\"https://msk-mode-test.cloud.databricks.com/serving-endpoints\"\n",
+    ")\n",
+    "\n"
+   ],
+   "id": "initial_id",
+   "outputs": [],
+   "execution_count": 8
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-25T23:39:55.623657Z",
+     "start_time": "2024-11-25T23:39:55.620263Z"
+    }
+   },
+   "cell_type": "code",
+   "source": " USER_PROMPT = \"What are the sites of disease based on this piece of text: IMPRESSION: Since CT scan of DATE, Predominantly upper lobe and superior segment lower lobe ground glass nodules are not seen on this chest radiograph. No consolidation or effusion.\"",
+   "id": "feab31d6aad17116",
+   "outputs": [],
+   "execution_count": 9
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-25T23:39:57.198005Z",
+     "start_time": "2024-11-25T23:39:55.714907Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "chat_completion = client.chat.completions.create(\n",
+    "    messages=[\n",
+    "        {\n",
+    "            \"role\": \"system\",\n",
+    "            \"content\": \"You are an AI assistant\"\n",
+    "        },\n",
+    "        {\n",
+    "            \"role\": \"user\",\n",
+    "            \"content\": USER_PROMPT\n",
+    "        }\n",
+    "    ],\n",
+    "    model=SERVING_MODEL,\n",
+    "    max_tokens=MAX_TOKENS\n",
+    ")\n",
+    "\n",
+    "print(chat_completion.choices[0].message.content)"
+   ],
+   "id": "e1c8106b1bfbb9fd",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "According to the text, the site of disease is the upper lobe and superior segment of the lower lobe of the lungs, specifically:\n",
+      "\n",
+      "* Predominantly upper lobe\n",
+      "* Superior segment lower lobe (indicating a specific region within the lower lobe)\n"
+     ]
+    }
+   ],
+   "execution_count": 10
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-25T23:39:57.404001Z",
+     "start_time": "2024-11-25T23:39:57.397328Z"
+    }
+   },
+   "cell_type": "code",
+   "source": "",
+   "id": "8e6985ffc3d5f21f",
+   "outputs": [],
+   "execution_count": null
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "conda-env-cdm-fongc2",
+   "language": "python",
+   "name": "conda-env-cdm-fongc2"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}