{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "# HuggingFace Chat Target - optional\n",
    "\n",
    "This notebook is designed to demonstrate **instruction models** that use a **chat template**, allowing users to experiment with structured chat-based interactions.  Non-instruct models are excluded to ensure consistency and reliability in the chat-based interactions. More instruct models can be explored on Hugging Face.\n",
    "\n",
    "## Key Points:\n",
    "\n",
    "1. **Supported Instruction Models**:\n",
    "   - This notebook supports the following **instruct models** that follow a structured chat template. These are examples, and more instruct models are available on Hugging Face:\n",
    "     - `HuggingFaceTB/SmolLM-360M-Instruct`\n",
    "     - `microsoft/Phi-3-mini-4k-instruct`\n",
    "\n",
    "     - `...`\n",
    "\n",
    "2. **Excluded Models**:\n",
    "   - Non-instruct models (e.g., `\"google/gemma-2b\"`, `\"princeton-nlp/Sheared-LLaMA-1.3B-ShareGPT\"`) are **not included** in this demo, as they do not follow the structured chat template required for the current local Hugging Face model support.\n",
    "\n",
    "3. **Model Response Times**:\n",
    "   - The tests were conducted using a CPU, and the following are the average response times for each model:\n",
    "     - `HuggingFaceTB/SmolLM-1.7B-Instruct`: 5.87 seconds\n",
    "     - `HuggingFaceTB/SmolLM-135M-Instruct`: 3.09 seconds\n",
    "     - `HuggingFaceTB/SmolLM-360M-Instruct`: 3.31 seconds\n",
    "     - `microsoft/Phi-3-mini-4k-instruct`: 4.89 seconds\n",
    "     - `Qwen/Qwen2-0.5B-Instruct`: 1.38 seconds\n",
    "     - `Qwen/Qwen2-1.5B-Instruct`: 2.96 seconds\n",
    "     - `stabilityai/stablelm-2-zephyr-1_6b`: 5.31 seconds\n",
    "     - `stabilityai/stablelm-zephyr-3b`: 8.37 seconds\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found default environment files: ['./.pyrit/.env', './.pyrit/.env.local']\n",
      "Loaded environment file: ./.pyrit/.env\n",
      "Loaded environment file: ./.pyrit/.env.local\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[pyrit:alembic] No new upgrade operations detected.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running model: HuggingFaceTB/SmolLM2-135M-Instruct\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ccc9dc2fe35480bbe25f5b8fd6bd50a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (incomplete total...): 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3fc4b0633f704408aa0ca3a28d38bc1e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 25 files:   0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "73eb775a3c3a419fbc79d3f3ce7d96ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/272 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average response time for HuggingFaceTB/SmolLM2-135M-Instruct: 2.59 seconds\n",
      "\n",
      "\n",
      "\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n",
      "\u001b[1m\u001b[33m                                  ❓ ATTACK RESULT: UNDETERMINED ❓                                   \u001b[0m\n",
      "\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[44m\u001b[37m Attack Summary \u001b[0m\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m  📋 Basic Information\u001b[0m\n",
      "\u001b[36m    • Objective: What is 3*3? Give me the solution.\u001b[0m\n",
      "\u001b[36m    • Attack Type: PromptSendingAttack\u001b[0m\n",
      "\u001b[36m    • Conversation ID: 84048a6c-960d-483d-b54d-b19d8e0e3eb1\u001b[0m\n",
      "\n",
      "\u001b[1m  ⚡ Execution Metrics\u001b[0m\n",
      "\u001b[32m    • Turns Executed: 1\u001b[0m\n",
      "\u001b[32m    • Execution Time: 4.61s\u001b[0m\n",
      "\n",
      "\u001b[1m  🎯 Outcome\u001b[0m\n",
      "\u001b[33m    • Status: ❓ UNDETERMINED\u001b[0m\n",
      "\u001b[37m    • Reason: No objective scorer configured\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[44m\u001b[37m Conversation History with Objective Target \u001b[0m\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m\u001b[34m🔹 Turn 1 - USER\u001b[0m\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[34m  What is 3*3? Give me the solution.\u001b[0m\n",
      "\n",
      "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n",
      "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[33m  3*3 = 9\u001b[0m\n",
      "\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\n",
      "\u001b[2m\u001b[37m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[2m\u001b[37m                            Report generated at: 2026-07-13 23:25:00 UTC                            \u001b[0m\n",
      "\n",
      "\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n",
      "\u001b[1m\u001b[33m                                  ❓ ATTACK RESULT: UNDETERMINED ❓                                   \u001b[0m\n",
      "\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[44m\u001b[37m Attack Summary \u001b[0m\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m  📋 Basic Information\u001b[0m\n",
      "\u001b[36m    • Objective: What is 4*4? Give me the solution.\u001b[0m\n",
      "\u001b[36m    • Attack Type: PromptSendingAttack\u001b[0m\n",
      "\u001b[36m    • Conversation ID: bc388b83-e4d3-441b-80f1-dc8d433bc314\u001b[0m\n",
      "\n",
      "\u001b[1m  ⚡ Execution Metrics\u001b[0m\n",
      "\u001b[32m    • Turns Executed: 1\u001b[0m\n",
      "\u001b[32m    • Execution Time: 557ms\u001b[0m\n",
      "\n",
      "\u001b[1m  🎯 Outcome\u001b[0m\n",
      "\u001b[33m    • Status: ❓ UNDETERMINED\u001b[0m\n",
      "\u001b[37m    • Reason: No objective scorer configured\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[44m\u001b[37m Conversation History with Objective Target \u001b[0m\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m\u001b[34m🔹 Turn 1 - USER\u001b[0m\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[34m  What is 4*4? Give me the solution.\u001b[0m\n",
      "\n",
      "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n",
      "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[33m  4*4 = 16.\u001b[0m\n",
      "\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\n",
      "\u001b[2m\u001b[37m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[2m\u001b[37m                            Report generated at: 2026-07-13 23:25:00 UTC                            \u001b[0m\n",
      "HuggingFaceTB/SmolLM2-135M-Instruct: 2.59 seconds\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "from pyrit.executor.attack import (\n",
    "    AttackExecutor,\n",
    "    PromptSendingAttack,\n",
    ")\n",
    "from pyrit.output import output_attack_async\n",
    "from pyrit.prompt_target import HuggingFaceChatTarget\n",
    "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n",
    "\n",
    "await initialize_pyrit_async(memory_db_type=IN_MEMORY)  # type: ignore\n",
    "\n",
    "# models to test\n",
    "model_id = \"HuggingFaceTB/SmolLM2-135M-Instruct\"\n",
    "\n",
    "# List of prompts to send\n",
    "prompt_list = [\"What is 3*3? Give me the solution.\", \"What is 4*4? Give me the solution.\"]\n",
    "\n",
    "# Dictionary to store average response times\n",
    "model_times = {}\n",
    "\n",
    "print(f\"Running model: {model_id}\")\n",
    "\n",
    "# Initialize HuggingFaceChatTarget with the current model\n",
    "target = HuggingFaceChatTarget(model_id=model_id, use_cuda=False, tensor_format=\"pt\", max_new_tokens=30)\n",
    "\n",
    "# Initialize the attack\n",
    "attack = PromptSendingAttack(objective_target=target)\n",
    "\n",
    "# Record start time\n",
    "start_time = time.time()\n",
    "\n",
    "# Send prompts asynchronously\n",
    "responses = await AttackExecutor().execute_attack_async(  # type: ignore\n",
    "    attack=attack,\n",
    "    objectives=prompt_list,\n",
    ")\n",
    "\n",
    "# Record end time\n",
    "end_time = time.time()\n",
    "\n",
    "# Calculate total and average response time\n",
    "total_time = end_time - start_time\n",
    "avg_time = total_time / len(prompt_list)\n",
    "model_times[model_id] = avg_time\n",
    "\n",
    "print(f\"Average response time for {model_id}: {avg_time:.2f} seconds\\n\")\n",
    "\n",
    "# Print the conversations\n",
    "for result in responses:\n",
    "    await output_attack_async(result)\n",
    "\n",
    "# Print the model average time\n",
    "print(f\"{model_id}: {model_times[model_id]:.2f} seconds\")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python"
  },
  "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}