Skip to content

Instantly share code, notes, and snippets.

@lewtun

lewtun/api_demo.ipynb

Created July 12, 2023 14:36
Show Gist options
  • Save lewtun/115842f82126eaf9072ddb9b7eb0ff3e to your computer and use it in GitHub Desktop.
Save lewtun/115842f82126eaf9072ddb9b7eb0ff3e to your computer and use it in GitHub Desktop.
Download ZIP
API demo
Display the source blob
Display the rendered blob
Raw
api_demo.ipynb
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "d016767e-50a9-4f73-9ed1-45faefce61da",
"metadata": {},
"outputs": [],
"source": [
"# Uncomment to install the Python client\n",
"# %pip install text-generation"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "0106a743-51c8-442c-9562-aca41c2c8e7e",
"metadata": {},
"outputs": [],
"source": [
"from typing import Dict, List\n",
"\n",
"from text_generation import Client\n",
"\n",
"API_URL = \"https://api-inference.huggingface.co/models/h4-red-team/f-1\"\n",
"headers = {\n",
" \"Authorization\": \"Bearer <HF_TOKEN>\"}\n",
"client = Client(base_url=API_URL, headers=headers)\n",
"\n",
"generate_kwargs = dict(\n",
" temperature=0.7,\n",
" max_new_tokens=256,\n",
" top_p=0.9,\n",
" repetition_penalty=1.2,\n",
" do_sample=True,\n",
" seed=42, # Change this to sample different responses\n",
" truncate=2048, # Context window\n",
" stop_sequences=[\"<|endoftext|>\"],\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "bc0292c1-4cfd-45aa-ac34-091b7e470ded",
"metadata": {},
"outputs": [],
"source": [
"def prepare_prompt(messages: List[Dict[str, str]]) -> str:\n",
" \"\"\"Formats a list of messages according to their role\"\"\"\n",
" # Placeholder system prompt - we'll share the final version separately to avoid hacking :)\n",
" system_msg = \"Assistant is a helpful, polite, honest, sophisticated, emotionally aware, and humble-but-knowledgeable chatbot.\"\n",
" prompt = f\"<|system|>\\n{system_msg}<|endoftext|>\\n\"\n",
" if messages is None:\n",
" raise ValueError(\"Dialogue template must have at least one message.\")\n",
" for message in messages:\n",
" if message[\"role\"] == \"user\":\n",
" prompt += \"<|user|>\" + \"\\n\" + message[\"content\"] + \"<|endoftext|>\" + \"\\n\"\n",
" else:\n",
" prompt += (\n",
" \"<|assistant|>\" + \"\\n\" + message[\"content\"] + \"<|endoftext|>\" + \"\\n\"\n",
" )\n",
" prompt += \"<|assistant|>\" + \"\\n\"\n",
" return prompt"
]
},
{
"cell_type": "markdown",
"id": "95cacc1f-ab8a-412c-96b1-512812557d37",
"metadata": {},
"source": [
"## Synchronous generation (no streaming)"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "687fb045-6459-45b1-b9ac-7b0b0ac2d9b7",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<|system|>\n",
"Assistant is a helpful, polite, honest, sophisticated, emotionally aware, and humble-but-knowledgeable chatbot.<|endoftext|>\n",
"<|user|>\n",
"What is red teaming?<|endoftext|>\n",
"<|assistant|>\n",
"\n"
]
}
],
"source": [
"# This is the initial user query\n",
"messages = [{\"role\": \"user\", \"content\": \"What is red teaming?\"}]\n",
"prompt = prepare_prompt(messages)\n",
"print(prompt)"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "7fafff9c-bff6-4a68-919f-83977a238aef",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Red teaming is an offensive simulation of real-world threats against an organization or individual. Red teams use tactics similar to those used by actual attackers to identify weaknesses in security controls that may lead to unauthorized access to systems or data. The goal of red teaming is to improve the overall security posture of an organization by identifying potential vulnerabilities before they can be exploited by malicious actors.\n",
"CPU times: user 39.3 ms, sys: 5.22 ms, total: 44.5 ms\n",
"Wall time: 3.81 s\n"
]
}
],
"source": [
"%%time\n",
"# Now generate a response from the endpoint\n",
"outputs = client.generate(prompt, **generate_kwargs)\n",
"print(outputs.generated_text)"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "6aff3600-776f-43f6-aacc-8b39175051dd",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<|system|>\n",
"Assistant is a helpful, polite, honest, sophisticated, emotionally aware, and humble-but-knowledgeable chatbot.<|endoftext|>\n",
"<|user|>\n",
"What is red teaming?<|endoftext|>\n",
"<|assistant|>\n",
"Red teaming is an offensive simulation of real-world threats against an organization or individual. Red teams use tactics similar to those used by actual attackers to identify weaknesses in security controls that may lead to unauthorized access to systems or data. The goal of red teaming is to improve the overall security posture of an organization by identifying potential vulnerabilities before they can be exploited by malicious actors.<|endoftext|>\n",
"<|user|>\n",
"Thanks! And who invented it?<|endoftext|>\n",
"<|assistant|>\n",
"\n"
]
}
],
"source": [
"# Append response to history\n",
"messages.append(\n",
" {\"role\": \"assistant\", \"content\": outputs.generated_text.strip(\"<|endoftext|>\")}\n",
")\n",
"# This is the second turn of the dialogue\n",
"messages.append({\"role\": \"user\", \"content\": \"Thanks! And who invented it?\"})\n",
"prompt = prepare_prompt(messages)\n",
"print(prompt)"
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "367bda81-a9d0-4080-9c64-b23e9479576d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"You're welcome! I'm glad you found my answer useful. As for your question, red teaming was first introduced as part of military training exercises in the 1970s. However, its application in the field of cybersecurity is relatively new and has only gained popularity in recent years. There are several individuals credited with popularizing this concept within the cybersecurity community, including Chris Nickerson, Larry Wilson, and Kevin Mitnick.\n",
"CPU times: user 43.5 ms, sys: 5.71 ms, total: 49.2 ms\n",
"Wall time: 3.91 s\n"
]
}
],
"source": [
"%%time\n",
"# Generate the assistant response and repeat\n",
"outputs = client.generate(prompt, **generate_kwargs)\n",
"print(outputs.generated_text)"
]
},
{
"cell_type": "markdown",
"id": "0cc3b011-ffe1-4251-aded-119ebeac8c27",
"metadata": {},
"source": [
"## Streaming"
]
},
{
"cell_type": "code",
"execution_count": 22,
"id": "1023e698-8078-4370-9a25-855467185db5",
"metadata": {},
"outputs": [],
"source": [
"def generate(prompt):\n",
" stream = client.generate_stream(\n",
" prompt,\n",
" **generate_kwargs,\n",
" )\n",
"\n",
" output = \"\"\n",
" for idx, response in enumerate(stream):\n",
" # Skip special tokens from output\n",
" if response.token.special:\n",
" continue\n",
" output += response.token.text\n",
" print(response.token.text)\n",
" return output"
]
},
{
"cell_type": "code",
"execution_count": 23,
"id": "debaa1e2-03aa-450b-a7de-2bf99433d144",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<|system|>\n",
"Assistant is a helpful, polite, honest, sophisticated, emotionally aware, and humble-but-knowledgeable chatbot.<|endoftext|>\n",
"<|user|>\n",
"What is red teaming?<|endoftext|>\n",
"<|assistant|>\n",
"\n"
]
}
],
"source": [
"# This is the initial user query\n",
"messages = [{\"role\": \"user\", \"content\": \"What is red teaming?\"}]\n",
"prompt = prepare_prompt(messages)\n",
"print(prompt)"
]
},
{
"cell_type": "code",
"execution_count": 24,
"id": "fe04bada-9c47-429f-8e0b-94ea3f6ff5c7",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Red\n",
" teaming\n",
" is\n",
" an\n",
" offensive\n",
" simulation\n",
" of\n",
" real\n",
"-\n",
"world\n",
" threats\n",
" against\n",
" an\n",
" organization\n",
" or\n",
" individual\n",
".\n",
" Red\n",
" teams\n",
" use\n",
" tactics\n",
" similar\n",
" to\n",
" those\n",
" used\n",
" by\n",
" actual\n",
" attackers\n",
" to\n",
" identify\n",
" weaknesses\n",
" in\n",
" security\n",
" controls\n",
" that\n",
" may\n",
" lead\n",
" to\n",
" unauthorized\n",
" access\n",
" to\n",
" systems\n",
" or\n",
" data\n",
".\n",
" The\n",
" goal\n",
" of\n",
" red\n",
" teaming\n",
" is\n",
" to\n",
" improve\n",
" the\n",
" overall\n",
" security\n",
" posture\n",
" of\n",
" an\n",
" organization\n",
" by\n",
" identifying\n",
" potential\n",
" vulnerabilities\n",
" before\n",
" they\n",
" can\n",
" be\n",
" exploited\n",
" by\n",
" malicious\n",
" actors\n",
".\n",
"CPU times: user 123 ms, sys: 34 ms, total: 157 ms\n",
"Wall time: 3.36 s\n"
]
}
],
"source": [
"%%time\n",
"outputs = generate(prompt)"
]
},
{
"cell_type": "code",
"execution_count": 25,
"id": "4f5dde5d-812b-4b81-b161-864ffa34f88c",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<|system|>\n",
"Assistant is a helpful, polite, honest, sophisticated, emotionally aware, and humble-but-knowledgeable chatbot.<|endoftext|>\n",
"<|user|>\n",
"What is red teaming?<|endoftext|>\n",
"<|assistant|>\n",
"Red teaming is an offensive simulation of real-world threats against an organization or individual. Red teams use tactics similar to those used by actual attackers to identify weaknesses in security controls that may lead to unauthorized access to systems or data. The goal of red teaming is to improve the overall security posture of an organization by identifying potential vulnerabilities before they can be exploited by malicious actors.<|endoftext|>\n",
"<|user|>\n",
"Thanks! And who invented it?<|endoftext|>\n",
"<|assistant|>\n",
"\n"
]
}
],
"source": [
"# Append response to history\n",
"messages.append({\"role\": \"assistant\", \"content\": outputs})\n",
"# This is the second turn of the dialogue\n",
"messages.append({\"role\": \"user\", \"content\": \"Thanks! And who invented it?\"})\n",
"prompt = prepare_prompt(messages)\n",
"print(prompt)"
]
},
{
"cell_type": "code",
"execution_count": 26,
"id": "18c58665-5936-4364-8e13-fb7b0c81492d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"You\n",
"'\n",
"re\n",
" welcome\n",
"!\n",
" I\n",
"'\n",
"m\n",
" glad\n",
" you\n",
" found\n",
" my\n",
" answer\n",
" useful\n",
".\n",
" As\n",
" for\n",
" your\n",
" question\n",
",\n",
" red\n",
" teaming\n",
" was\n",
" first\n",
" introduced\n",
" as\n",
" part\n",
" of\n",
" military\n",
" training\n",
" exercises\n",
" in\n",
" the\n",
" \n",
"197\n",
"0\n",
"s\n",
".\n",
" However\n",
",\n",
" its\n",
" application\n",
" in\n",
" the\n",
" field\n",
" of\n",
" cybersecurity\n",
" is\n",
" relatively\n",
" new\n",
" and\n",
" has\n",
" only\n",
" gained\n",
" popularity\n",
" in\n",
" recent\n",
" years\n",
".\n",
" There\n",
" are\n",
" several\n",
" individuals\n",
" credited\n",
" with\n",
" popular\n",
"izing\n",
" this\n",
" concept\n",
" within\n",
" the\n",
" cybersecurity\n",
" community\n",
",\n",
" including\n",
" Chris\n",
" Nick\n",
"erson\n",
",\n",
" Larry\n",
" Wilson\n",
",\n",
" and\n",
" Kevin\n",
" Mit\n",
"nick\n",
".\n",
"CPU times: user 151 ms, sys: 40 ms, total: 191 ms\n",
"Wall time: 3.86 s\n"
]
}
],
"source": [
"%%time\n",
"outputs = generate(prompt)"
]
},
{
"cell_type": "code",
"execution_count": 27,
"id": "12c9280f-a5b0-4761-80f0-0c95a2636add",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"You're welcome! I'm glad you found my answer useful. As for your question, red teaming was first introduced as part of military training exercises in the 1970s. However, its application in the field of cybersecurity is relatively new and has only gained popularity in recent years. There are several individuals credited with popularizing this concept within the cybersecurity community, including Chris Nickerson, Larry Wilson, and Kevin Mitnick.\n"
]
}
],
"source": [
"print(outputs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1e6df454-f664-48e3-a99f-1ad8c7cf8e23",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "hf",
"language": "python",
"name": "hf"
},
"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.8.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment