examples/python/llama.cpp/PromptAssember_with_AutoGGUFModel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![JohnSnowLabs](https://sparknlp.org/assets/images/logo.png)\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/JohnSnowLabs/spark-nlp/blob/master/examples/python/llama.cpp/PromptAssember_with_AutoGGUFModel.ipynb)\n",
    "\n",
    "# PromptAssembler with AutoGGUFModel\n",
    "\n",
    "Let's keep in mind a few things before we start 😊\n",
    "\n",
    "- llama.cpp support in the form of the `AutoGGUFModel` was introduced in `Spark NLP 5.5.0`, enabling quantized LLM inference on a wide range of devices. Please make sure you have upgraded to the latest Spark NLP release.\n",
    "- The `PromptAssembler` was introduced in `Spark NLP 5.5.1` to enable the construction of message prompts.\n",
    "\n",
    "This notebook will show you how you can construct your own message prompts for the AutoGGUFModel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install and Start Spark NLP\n",
    "\n",
    "- Let's install and setup Spark NLP (if running it Google Colab)\n",
    "- This part is pretty easy via our simple script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only execute this if you are on Google Colab\n",
    "! wget -q http://setup.johnsnowlabs.com/colab.sh -O - | bash"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start Spark with Spark NLP included via our simple `start()` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sparknlp\n",
    "\n",
    "# let's start Spark with Spark NLP with GPU enabled. If you don't have GPUs available remove this parameter.\n",
    "spark = sparknlp.start(gpu=True)\n",
    "print(sparknlp.version())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a `PromptAssembler` and use it to recreate the following conversation between a chatbot and a user:\n",
    "\n",
    "```\n",
    "SYSTEM: You are a helpful assistant.\n",
    "ASSISTANT: Hello there! How can I help you today?\n",
    "USER: I need help with organizing my room, give me some advice.\n",
    "```\n",
    "\n",
    "First we need to structure our messages in our Spark DataFrame correctly. For each row, the PromptAssembler expects an array of two-tuples. The first field should be the role and the second field the message. We will call this column `message`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------------------------------------------------------------------+\n",
      "|messages                                                                                                                                                        |\n",
      "+----------------------------------------------------------------------------------------------------------------------------------------------------------------+\n",
      "|[{system, You are a helpful assistant.}, {assistant, Hello there! How can I help you today?}, {user, I need help with organizing my room, give me some advice.}]|\n",
      "+----------------------------------------------------------------------------------------------------------------------------------------------------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    (\"system\", \"You are a helpful assistant.\"),\n",
    "    (\"assistant\", \"Hello there! How can I help you today?\"),\n",
    "    (\"user\", \"I need help with organizing my room, give me some advice.\"),\n",
    "]\n",
    "df = spark.createDataFrame([[messages]]).toDF(\"messages\")\n",
    "df.show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create the PromptAssembler to generate the prompts. We will use the template from [llama3.1 (extracted from the gguf file)](https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF?show_file_info=Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf).\n",
    "\n",
    "By default, the `addAssistant` parameter is set to `True`, so a assistant header will be appended to the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sparknlp.base import *\n",
    "\n",
    "\n",
    "template = (\n",
    "    \"{{- bos_token }} {%- if custom_tools is defined %} {%- set tools = custom_tools %} {%- \"\n",
    "    \"endif %} {%- if not tools_in_user_message is defined %} {%- set tools_in_user_message = true %} {%- \"\n",
    "    'endif %} {%- if not date_string is defined %} {%- set date_string = \"26 Jul 2024\" %} {%- endif %} '\n",
    "    \"{%- if not tools is defined %} {%- set tools = none %} {%- endif %} {#- This block extracts the \"\n",
    "    \"system message, so we can slot it into the right place. #} {%- if messages[0]['role'] == 'system' %}\"\n",
    "    \" {%- set system_message = messages[0]['content']|trim %} {%- set messages = messages[1:] %} {%- else\"\n",
    "    ' %} {%- set system_message = \"\" %} {%- endif %} {#- System message + builtin tools #} {{- '\n",
    "    '\"<|start_header_id|>system<|end_header_id|>\\\\n\\\\n\" }} {%- if builtin_tools is defined or tools is '\n",
    "    'not none %} {{- \"Environment: ipython\\\\n\" }} {%- endif %} {%- if builtin_tools is defined %} {{- '\n",
    "    '\"Tools: \" + builtin_tools | reject(\\'equalto\\', \\'code_interpreter\\') | join(\", \") + \"\\\\n\\\\n\"}} '\n",
    "    '{%- endif %} {{- \"Cutting Knowledge Date: December 2023\\\\n\" }} {{- \"Today Date: \" + date_string '\n",
    "    '+ \"\\\\n\\\\n\" }} {%- if tools is not none and not tools_in_user_message %} {{- \"You have access to '\n",
    "    'the following functions. To call a function, please respond with JSON for a function call.\" }} {{- '\n",
    "    '\\'Respond in the format {\"name\": function name, \"parameters\": dictionary of argument name and its'\n",
    "    ' value}.\\' }} {{- \"Do not use variables.\\\\n\\\\n\" }} {%- for t in tools %} {{- t | tojson(indent=4) '\n",
    "    '}} {{- \"\\\\n\\\\n\" }} {%- endfor %} {%- endif %} {{- system_message }} {{- \"<|eot_id|>\" }} {#- '\n",
    "    \"Custom tools are passed in a user message with some extra guidance #} {%- if tools_in_user_message \"\n",
    "    \"and not tools is none %} {#- Extract the first user message so we can plug it in here #} {%- if \"\n",
    "    \"messages | length != 0 %} {%- set first_user_message = messages[0]['content']|trim %} {%- set \"\n",
    "    'messages = messages[1:] %} {%- else %} {{- raise_exception(\"Cannot put tools in the first user '\n",
    "    \"message when there's no first user message!\\\") }} {%- endif %} {{- \"\n",
    "    \"'<|start_header_id|>user<|end_header_id|>\\\\n\\\\n' -}} {{- \\\"Given the following functions, please \"\n",
    "    'respond with a JSON for a function call \" }} {{- \"with its proper arguments that best answers the '\n",
    "    'given prompt.\\\\n\\\\n\" }} {{- \\'Respond in the format {\"name\": function name, \"parameters\": '\n",
    "    'dictionary of argument name and its value}.\\' }} {{- \"Do not use variables.\\\\n\\\\n\" }} {%- for t in '\n",
    "    'tools %} {{- t | tojson(indent=4) }} {{- \"\\\\n\\\\n\" }} {%- endfor %} {{- first_user_message + '\n",
    "    \"\\\"<|eot_id|>\\\"}} {%- endif %} {%- for message in messages %} {%- if not (message.role == 'ipython' \"\n",
    "    \"or message.role == 'tool' or 'tool_calls' in message) %} {{- '<|start_header_id|>' + message['role']\"\n",
    "    \" + '<|end_header_id|>\\\\n\\\\n'+ message['content'] | trim + '<|eot_id|>' }} {%- elif 'tool_calls' in \"\n",
    "    'message %} {%- if not message.tool_calls|length == 1 %} {{- raise_exception(\"This model only '\n",
    "    'supports single tool-calls at once!\") }} {%- endif %} {%- set tool_call = message.tool_calls[0]'\n",
    "    \".function %} {%- if builtin_tools is defined and tool_call.name in builtin_tools %} {{- \"\n",
    "    \"'<|start_header_id|>assistant<|end_header_id|>\\\\n\\\\n' -}} {{- \\\"<|python_tag|>\\\" + tool_call.name + \"\n",
    "    '\".call(\" }} {%- for arg_name, arg_val in tool_call.arguments | items %} {{- arg_name + \\'=\"\\' + '\n",
    "    'arg_val + \\'\"\\' }} {%- if not loop.last %} {{- \", \" }} {%- endif %} {%- endfor %} {{- \")\" }} {%- '\n",
    "    \"else %} {{- '<|start_header_id|>assistant<|end_header_id|>\\\\n\\\\n' -}} {{- '{\\\"name\\\": \\\"' + \"\n",
    "    'tool_call.name + \\'\", \\' }} {{- \\'\"parameters\": \\' }} {{- tool_call.arguments | tojson }} {{- \"}\" '\n",
    "    \"}} {%- endif %} {%- if builtin_tools is defined %} {#- This means we're in ipython mode #} {{- \"\n",
    "    '\"<|eom_id|>\" }} {%- else %} {{- \"<|eot_id|>\" }} {%- endif %} {%- elif message.role == \"tool\" '\n",
    "    'or message.role == \"ipython\" %} {{- \"<|start_header_id|>ipython<|end_header_id|>\\\\n\\\\n\" }} {%- '\n",
    "    \"if message.content is mapping or message.content is iterable %} {{- message.content | tojson }} {%- \"\n",
    "    'else %} {{- message.content }} {%- endif %} {{- \"<|eot_id|>\" }} {%- endif %} {%- endfor %} {%- if '\n",
    "    \"add_generation_prompt %} {{- '<|start_header_id|>assistant<|end_header_id|>\\\\n\\\\n' }} {%- endif %} \"\n",
    ")\n",
    "\n",
    "promptAssembler = (\n",
    "    PromptAssembler()\n",
    "    .setInputCol(\"messages\")\n",
    "    .setOutputCol(\"prompt\")\n",
    "    .setChatTemplate(template)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how the final prompt looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 8:>                                                          (0 + 1) / 1]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jsl-llama: Extracted 'libjllama.so' to '/tmp/libjllama.so'\n",
      "+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+\n",
      "|result                                                                                                                                                                                                                                                                                                                                                 |\n",
      "+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+\n",
      "|[<|start_header_id|>system<|end_header_id|>\\n\\nYou are a helpful assistant.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\nHello there! How can I help you today?<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\nI need help with organizing my room, give me some advice.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n]|\n",
      "+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "promptAssembler.transform(df).select(\"prompt.result\").show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can feed the prompt to a llama3.1 model loaded with AutoGGUFModel. Depending on your messages, you might need to the chat template or system prompt in the AutoGGUFModel. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sparknlp.annotator import AutoGGUFModel\n",
    "\n",
    "autoGGUFModel = (\n",
    "    AutoGGUFModel.loadSavedModel(\"path/to/llama3.1\", spark)\n",
    "    .setInputCols(\"prompt\")\n",
    "    .setOutputCol(\"completions\")\n",
    "    .setBatchSize(4)\n",
    "    .setNGpuLayers(99)\n",
    "    .setUseChatTemplate(False)  # Don't apply the chat template\n",
    "    .setSystemPrompt(\n",
    "        \"Your system prompt\"\n",
    "    )  # Set custom system prompt if not specified in the messages. Leave empty for default.\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "sparknlp_dev",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}