https://github.com/spellcraftai/oaib
Use the OpenAI Batch tool to make async batch requests to the OpenAI API.
https://github.com/spellcraftai/oaib
Last synced: 12 months ago
JSON representation
Use the OpenAI Batch tool to make async batch requests to the OpenAI API.
- Host: GitHub
- URL: https://github.com/spellcraftai/oaib
- Owner: SpellcraftAI
- License: mit
- Created: 2024-02-10T15:35:29.000Z (over 2 years ago)
- Default Branch: dev
- Last Pushed: 2024-03-02T22:26:29.000Z (over 2 years ago)
- Last Synced: 2025-03-24T17:28:22.519Z (about 1 year ago)
- Language: Python
- Homepage:
- Size: 85.9 KB
- Stars: 96
- Watchers: 3
- Forks: 3
- Open Issues: 6
-
Metadata Files:
- Readme: README.ipynb
- License: LICENSE
Awesome Lists containing this project
README
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# OpenAI Batcher (`oaib`)\n",
"\n",
"A Python library for making rate-limited, async batch requests to the OpenAI\n",
"API.\n",
"\n",
"- Use `Batch` to make batch requests as quickly as possible given TPM/RPM\n",
" limits.\n",
"- Use `Auto` to automatically read your rate limits from OpenAI's response\n",
" headers, and run the job as fast as possible."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Set OpenAI API Key"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ['OPENAI_API_KEY'] = input()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Using the `Batch` class\n",
"\n",
"You can mix and match endpoints as needed for regular `Batch`, but it's not\n",
"recommended as the rate limits for different endpoints/models will differ. For\n",
"maximum efficiency, use `Auto`."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "1b639118a1114517bb39abdae101cfcf",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/20 [00:00, ?req/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "691da1fccbf1427ca8be81548594d5d6",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"RPM: 0%| | 0/100"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "74c6cefa195a4bf89a743ff6891e207b",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"TPM: 0%| | 0/1000"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Run took 20.02s.\n",
"\n"
]
},
{
"data": {
"text/html": [
"
\n",
"\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"\n",
"\n",
" \n",
" \n",
" \n",
" endpoint\n",
" model\n",
" messages\n",
" result\n",
" \n",
" \n",
" \n",
" \n",
" 0\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzMf7oXHlLWpsISVpAaUPxps0g5...\n",
" \n",
" \n",
" 1\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzNeFY9TLN71o2FMEqGssHIrOQq...\n",
" \n",
" \n",
" 2\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzNDiD3ikFtBZ4hHXWEsLONHUeS...\n",
" \n",
" \n",
" 3\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzP9dccsvrGsOR3X5HgmHqsR2fm...\n",
" \n",
" \n",
" 4\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzQV7ZnIoXccx9R8dIfS4rdPd0U...\n",
" \n",
" \n",
" 5\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzRVQvp3wwmEvbFzNPtrXBmcOhR...\n",
" \n",
" \n",
" 6\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzRSw7iTCLs0uu8fWZwDcaPGB0s...\n",
" \n",
" \n",
" 7\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzS6D1gACsJW6JXvuS42N4lQLh7...\n",
" \n",
" \n",
" 8\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzUFmpIzWjKsNGnlLvZW3DhF752...\n",
" \n",
" \n",
" 9\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzU6Mg6Zk4BC5uelndmHjmGAQ0I...\n",
" \n",
" \n",
" 10\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzV5K4OEk80dDuSwohiTualLOoO...\n",
" \n",
" \n",
" 11\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzXCcmbcuy1EQPskJrfN5po1Ix9...\n",
" \n",
" \n",
" 12\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzX4UfaBPDO3fF8vMO1dsQ6tfiT...\n",
" \n",
" \n",
" 13\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzZme2VBNhckfItEZRJqmpC3E53...\n",
" \n",
" \n",
" 14\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzaYddjxrFpmYUDjMHXlDPgS7G4...\n",
" \n",
" \n",
" 15\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzcTLFClvwjtPNNyv3KJ2xvZ8dQ...\n",
" \n",
" \n",
" 16\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzcuWlMUwqXj0AeMQkZqzNIQvJo...\n",
" \n",
" \n",
" 17\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzepfAhFnZj1AlnVSyNHHMyZ0mK...\n",
" \n",
" \n",
" 18\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzfdULvZx6OyUZpARpTMNKYzfKx...\n",
" \n",
" \n",
" 19\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVzfyEru2f67JyunYCdPMaC0fD2H...\n",
" \n",
" \n",
"\n",
""
],
"text/plain": [
" endpoint model \\\n",
"0 chat.completions.create gpt-3.5-turbo \n",
"1 chat.completions.create gpt-3.5-turbo \n",
"2 chat.completions.create gpt-3.5-turbo \n",
"3 chat.completions.create gpt-3.5-turbo \n",
"4 chat.completions.create gpt-3.5-turbo \n",
"5 chat.completions.create gpt-3.5-turbo \n",
"6 chat.completions.create gpt-3.5-turbo \n",
"7 chat.completions.create gpt-3.5-turbo \n",
"8 chat.completions.create gpt-3.5-turbo \n",
"9 chat.completions.create gpt-3.5-turbo \n",
"10 chat.completions.create gpt-3.5-turbo \n",
"11 chat.completions.create gpt-3.5-turbo \n",
"12 chat.completions.create gpt-3.5-turbo \n",
"13 chat.completions.create gpt-3.5-turbo \n",
"14 chat.completions.create gpt-3.5-turbo \n",
"15 chat.completions.create gpt-3.5-turbo \n",
"16 chat.completions.create gpt-3.5-turbo \n",
"17 chat.completions.create gpt-3.5-turbo \n",
"18 chat.completions.create gpt-3.5-turbo \n",
"19 chat.completions.create gpt-3.5-turbo \n",
"\n",
" messages \\\n",
"0 [{'role': 'user', 'content': 'say hello'}] \n",
"1 [{'role': 'user', 'content': 'say hello'}] \n",
"2 [{'role': 'user', 'content': 'say hello'}] \n",
"3 [{'role': 'user', 'content': 'say hello'}] \n",
"4 [{'role': 'user', 'content': 'say hello'}] \n",
"5 [{'role': 'user', 'content': 'say hello'}] \n",
"6 [{'role': 'user', 'content': 'say hello'}] \n",
"7 [{'role': 'user', 'content': 'say hello'}] \n",
"8 [{'role': 'user', 'content': 'say hello'}] \n",
"9 [{'role': 'user', 'content': 'say hello'}] \n",
"10 [{'role': 'user', 'content': 'say hello'}] \n",
"11 [{'role': 'user', 'content': 'say hello'}] \n",
"12 [{'role': 'user', 'content': 'say hello'}] \n",
"13 [{'role': 'user', 'content': 'say hello'}] \n",
"14 [{'role': 'user', 'content': 'say hello'}] \n",
"15 [{'role': 'user', 'content': 'say hello'}] \n",
"16 [{'role': 'user', 'content': 'say hello'}] \n",
"17 [{'role': 'user', 'content': 'say hello'}] \n",
"18 [{'role': 'user', 'content': 'say hello'}] \n",
"19 [{'role': 'user', 'content': 'say hello'}] \n",
"\n",
" result \n",
"0 {'id': 'chatcmpl-8qVzMf7oXHlLWpsISVpAaUPxps0g5... \n",
"1 {'id': 'chatcmpl-8qVzNeFY9TLN71o2FMEqGssHIrOQq... \n",
"2 {'id': 'chatcmpl-8qVzNDiD3ikFtBZ4hHXWEsLONHUeS... \n",
"3 {'id': 'chatcmpl-8qVzP9dccsvrGsOR3X5HgmHqsR2fm... \n",
"4 {'id': 'chatcmpl-8qVzQV7ZnIoXccx9R8dIfS4rdPd0U... \n",
"5 {'id': 'chatcmpl-8qVzRVQvp3wwmEvbFzNPtrXBmcOhR... \n",
"6 {'id': 'chatcmpl-8qVzRSw7iTCLs0uu8fWZwDcaPGB0s... \n",
"7 {'id': 'chatcmpl-8qVzS6D1gACsJW6JXvuS42N4lQLh7... \n",
"8 {'id': 'chatcmpl-8qVzUFmpIzWjKsNGnlLvZW3DhF752... \n",
"9 {'id': 'chatcmpl-8qVzU6Mg6Zk4BC5uelndmHjmGAQ0I... \n",
"10 {'id': 'chatcmpl-8qVzV5K4OEk80dDuSwohiTualLOoO... \n",
"11 {'id': 'chatcmpl-8qVzXCcmbcuy1EQPskJrfN5po1Ix9... \n",
"12 {'id': 'chatcmpl-8qVzX4UfaBPDO3fF8vMO1dsQ6tfiT... \n",
"13 {'id': 'chatcmpl-8qVzZme2VBNhckfItEZRJqmpC3E53... \n",
"14 {'id': 'chatcmpl-8qVzaYddjxrFpmYUDjMHXlDPgS7G4... \n",
"15 {'id': 'chatcmpl-8qVzcTLFClvwjtPNNyv3KJ2xvZ8dQ... \n",
"16 {'id': 'chatcmpl-8qVzcuWlMUwqXj0AeMQkZqzNIQvJo... \n",
"17 {'id': 'chatcmpl-8qVzepfAhFnZj1AlnVSyNHHMyZ0mK... \n",
"18 {'id': 'chatcmpl-8qVzfdULvZx6OyUZpARpTMNKYzfKx... \n",
"19 {'id': 'chatcmpl-8qVzfyEru2f67JyunYCdPMaC0fD2H... "
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from oaib import Batch\n",
"\n",
"# Use low rate limits for this example.\n",
"batch = Batch(rpm=100, tpm=1_000, workers=5)\n",
"\n",
"# Creating a batch with 20 chat completions.\n",
"for i in range(20):\n",
" await batch.add(\n",
" \"chat.completions.create\", \n",
" model=\"gpt-3.5-turbo\", \n",
" messages=[{\"role\": \"user\", \"content\": \"say hello\"}]\n",
" )\n",
"\n",
"await batch.run()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Using the `Auto` class\n",
"\n",
"Automatically use the given TPM/RPM rate limits provided by OpenAI API\n",
"responses."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "67d8275d5a2f4f9cbd2ab2eef1fd1c0d",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/1000 [00:00, ?req/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "034506ef02c841f59533536f425e5eb6",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"RPM: 0%| | 0/500"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "64443f4ef537420389abd7ceb2485cef",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"TPM: 0%| | 0/10000"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Run took 12.58s.\n",
"\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"\n",
"\n",
" \n",
" \n",
" \n",
" endpoint\n",
" model\n",
" messages\n",
" result\n",
" \n",
" \n",
" \n",
" \n",
" 0\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVxz211nsiQQSY2k54r4r141UX83...\n",
" \n",
" \n",
" 1\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVy0TrnGax3XqkDOlEiImosPnvIL...\n",
" \n",
" \n",
" 2\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVxzMdsF5v1je6iAmvjK7BCZKQna...\n",
" \n",
" \n",
" 3\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVy0KiRWYPtQ099p3b1k1HfYBYwT...\n",
" \n",
" \n",
" 4\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVy0M7l6Fg0mvWCpOXomVSqV6Pow...\n",
" \n",
" \n",
" ...\n",
" ...\n",
" ...\n",
" ...\n",
" ...\n",
" \n",
" \n",
" 995\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVyBemEPgb6lV5Opnu8X9UQ7T9iZ...\n",
" \n",
" \n",
" 996\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVyBBlh4PMbI9qtca80UyMbrOGAF...\n",
" \n",
" \n",
" 997\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVyBWveGsMeSLS5SzUbXIiMLvGaS...\n",
" \n",
" \n",
" 998\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVyBbBnX5SleJWSDIHqZ8lS0y15V...\n",
" \n",
" \n",
" 999\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'say hello'}]\n",
" {'id': 'chatcmpl-8qVyBYTizFbBIu3RQ9XlvKDlnSuEG...\n",
" \n",
" \n",
"\n",
"1000 rows × 4 columns
\n",
""
],
"text/plain": [
" endpoint model \\\n",
"0 chat.completions.create gpt-3.5-turbo \n",
"1 chat.completions.create gpt-3.5-turbo \n",
"2 chat.completions.create gpt-3.5-turbo \n",
"3 chat.completions.create gpt-3.5-turbo \n",
"4 chat.completions.create gpt-3.5-turbo \n",
".. ... ... \n",
"995 chat.completions.create gpt-3.5-turbo \n",
"996 chat.completions.create gpt-3.5-turbo \n",
"997 chat.completions.create gpt-3.5-turbo \n",
"998 chat.completions.create gpt-3.5-turbo \n",
"999 chat.completions.create gpt-3.5-turbo \n",
"\n",
" messages \\\n",
"0 [{'role': 'user', 'content': 'say hello'}] \n",
"1 [{'role': 'user', 'content': 'say hello'}] \n",
"2 [{'role': 'user', 'content': 'say hello'}] \n",
"3 [{'role': 'user', 'content': 'say hello'}] \n",
"4 [{'role': 'user', 'content': 'say hello'}] \n",
".. ... \n",
"995 [{'role': 'user', 'content': 'say hello'}] \n",
"996 [{'role': 'user', 'content': 'say hello'}] \n",
"997 [{'role': 'user', 'content': 'say hello'}] \n",
"998 [{'role': 'user', 'content': 'say hello'}] \n",
"999 [{'role': 'user', 'content': 'say hello'}] \n",
"\n",
" result \n",
"0 {'id': 'chatcmpl-8qVxz211nsiQQSY2k54r4r141UX83... \n",
"1 {'id': 'chatcmpl-8qVy0TrnGax3XqkDOlEiImosPnvIL... \n",
"2 {'id': 'chatcmpl-8qVxzMdsF5v1je6iAmvjK7BCZKQna... \n",
"3 {'id': 'chatcmpl-8qVy0KiRWYPtQ099p3b1k1HfYBYwT... \n",
"4 {'id': 'chatcmpl-8qVy0M7l6Fg0mvWCpOXomVSqV6Pow... \n",
".. ... \n",
"995 {'id': 'chatcmpl-8qVyBemEPgb6lV5Opnu8X9UQ7T9iZ... \n",
"996 {'id': 'chatcmpl-8qVyBBlh4PMbI9qtca80UyMbrOGAF... \n",
"997 {'id': 'chatcmpl-8qVyBWveGsMeSLS5SzUbXIiMLvGaS... \n",
"998 {'id': 'chatcmpl-8qVyBbBnX5SleJWSDIHqZ8lS0y15V... \n",
"999 {'id': 'chatcmpl-8qVyBYTizFbBIu3RQ9XlvKDlnSuEG... \n",
"\n",
"[1000 rows x 4 columns]"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from oaib import Auto\n",
"\n",
"# Automatically set rate limits.\n",
"batch = Auto(workers=8)\n",
"\n",
"# Fetch 1,000 chat completions as quickly as possible, setting rate limits\n",
"# automatically from OpenAI's response headers.\n",
"for i in range(1000):\n",
" await batch.add(\n",
" \"chat.completions.create\", \n",
" model=\"gpt-3.5-turbo\", \n",
" messages=[{\"role\": \"user\", \"content\": \"say hello\"}]\n",
" )\n",
"\n",
"# TPM/RPM progress bars will show as full the whole time for Auto runs due to a\n",
"# tqdm.notebook bug, but the percentage indicators are correct.\n",
"await batch.run()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Metadata and Index\n",
"\n",
"You can add custom metadata to your observations with `add(metadata={...}`, and\n",
"set the index for the DataFrame with `Batch(index=[...])` and\n",
"`Auto(index=[...])`."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "6a1c33001c494ac1a8f949d97949c6c6",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/15 [00:00, ?req/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "fa73772722a4493eae74a26bbeb50102",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"RPM: 0%| | 0/1000"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "45d1aec16e9d4e4c93fa8eedfc8687ef",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"TPM: 0%| | 0/10000"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Run took 4.74s.\n",
"\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"\n",
"\n",
" \n",
" \n",
" \n",
" \n",
" endpoint\n",
" model\n",
" messages\n",
" result\n",
" \n",
" \n",
" difficulty\n",
" i\n",
" \n",
" \n",
" \n",
" \n",
" \n",
" \n",
" \n",
" \n",
" easy\n",
" 0\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: easy...\n",
" {'id': 'chatcmpl-8rdiFkLTnjs4LbX2ZUKyb1UaRcPwH...\n",
" \n",
" \n",
" 1\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: easy...\n",
" {'id': 'chatcmpl-8rdiFtEcz6CEvO8K9jZpAdbFaaJ3o...\n",
" \n",
" \n",
" 2\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: easy...\n",
" {'id': 'chatcmpl-8rdiFEmy4TfO4iR3aJbh4u9lgAD1W...\n",
" \n",
" \n",
" 3\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: easy...\n",
" {'id': 'chatcmpl-8rdiFVw5YuHY8WhuNqJjyJcFO9Mhs...\n",
" \n",
" \n",
" 4\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: easy...\n",
" {'id': 'chatcmpl-8rdiGoYqWA5wH3xoFKVl8pRGGTuDZ...\n",
" \n",
" \n",
" hard\n",
" 0\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: hard...\n",
" {'id': 'chatcmpl-8rdiG8ZAUhcsBOtgtPuOjblzEo14a...\n",
" \n",
" \n",
" 1\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: hard...\n",
" {'id': 'chatcmpl-8rdiGgl9uEe4ASQgt5uzwMpeJhnU6...\n",
" \n",
" \n",
" 2\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: hard...\n",
" {'id': 'chatcmpl-8rdiGfVkiZqrE1p2TxeKnCxc2zzUb...\n",
" \n",
" \n",
" 3\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: hard...\n",
" {'id': 'chatcmpl-8rdiGdKxmYC6mS4QSuiW3HjsMRW05...\n",
" \n",
" \n",
" 4\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: hard...\n",
" {'id': 'chatcmpl-8rdiGqTe3MVy8FGJ6qtVAjxmMODy6...\n",
" \n",
" \n",
" medium\n",
" 0\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: medi...\n",
" {'id': 'chatcmpl-8rdiGsk7ohvIVwzuFMfO3cH9zNLW4...\n",
" \n",
" \n",
" 1\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: medi...\n",
" {'id': 'chatcmpl-8rdiG64Y66W8CZZ9MI4xiVNHheHiF...\n",
" \n",
" \n",
" 2\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: medi...\n",
" {'id': 'chatcmpl-8rdiGGmXfXx0uRuKafeKODGQe42vz...\n",
" \n",
" \n",
" 3\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: medi...\n",
" {'id': 'chatcmpl-8rdiG95XCozUiFXY7ryA4rSfzbwRk...\n",
" \n",
" \n",
" 4\n",
" chat.completions.create\n",
" gpt-3.5-turbo\n",
" [{'role': 'user', 'content': 'difficulty: medi...\n",
" {'id': 'chatcmpl-8rdiGi6nF5FLMdPIYrkrYgI15Yfcw...\n",
" \n",
" \n",
"\n",
""
],
"text/plain": [
" endpoint model \\\n",
"difficulty i \n",
"easy 0 chat.completions.create gpt-3.5-turbo \n",
" 1 chat.completions.create gpt-3.5-turbo \n",
" 2 chat.completions.create gpt-3.5-turbo \n",
" 3 chat.completions.create gpt-3.5-turbo \n",
" 4 chat.completions.create gpt-3.5-turbo \n",
"hard 0 chat.completions.create gpt-3.5-turbo \n",
" 1 chat.completions.create gpt-3.5-turbo \n",
" 2 chat.completions.create gpt-3.5-turbo \n",
" 3 chat.completions.create gpt-3.5-turbo \n",
" 4 chat.completions.create gpt-3.5-turbo \n",
"medium 0 chat.completions.create gpt-3.5-turbo \n",
" 1 chat.completions.create gpt-3.5-turbo \n",
" 2 chat.completions.create gpt-3.5-turbo \n",
" 3 chat.completions.create gpt-3.5-turbo \n",
" 4 chat.completions.create gpt-3.5-turbo \n",
"\n",
" messages \\\n",
"difficulty i \n",
"easy 0 [{'role': 'user', 'content': 'difficulty: easy... \n",
" 1 [{'role': 'user', 'content': 'difficulty: easy... \n",
" 2 [{'role': 'user', 'content': 'difficulty: easy... \n",
" 3 [{'role': 'user', 'content': 'difficulty: easy... \n",
" 4 [{'role': 'user', 'content': 'difficulty: easy... \n",
"hard 0 [{'role': 'user', 'content': 'difficulty: hard... \n",
" 1 [{'role': 'user', 'content': 'difficulty: hard... \n",
" 2 [{'role': 'user', 'content': 'difficulty: hard... \n",
" 3 [{'role': 'user', 'content': 'difficulty: hard... \n",
" 4 [{'role': 'user', 'content': 'difficulty: hard... \n",
"medium 0 [{'role': 'user', 'content': 'difficulty: medi... \n",
" 1 [{'role': 'user', 'content': 'difficulty: medi... \n",
" 2 [{'role': 'user', 'content': 'difficulty: medi... \n",
" 3 [{'role': 'user', 'content': 'difficulty: medi... \n",
" 4 [{'role': 'user', 'content': 'difficulty: medi... \n",
"\n",
" result \n",
"difficulty i \n",
"easy 0 {'id': 'chatcmpl-8rdiFkLTnjs4LbX2ZUKyb1UaRcPwH... \n",
" 1 {'id': 'chatcmpl-8rdiFtEcz6CEvO8K9jZpAdbFaaJ3o... \n",
" 2 {'id': 'chatcmpl-8rdiFEmy4TfO4iR3aJbh4u9lgAD1W... \n",
" 3 {'id': 'chatcmpl-8rdiFVw5YuHY8WhuNqJjyJcFO9Mhs... \n",
" 4 {'id': 'chatcmpl-8rdiGoYqWA5wH3xoFKVl8pRGGTuDZ... \n",
"hard 0 {'id': 'chatcmpl-8rdiG8ZAUhcsBOtgtPuOjblzEo14a... \n",
" 1 {'id': 'chatcmpl-8rdiGgl9uEe4ASQgt5uzwMpeJhnU6... \n",
" 2 {'id': 'chatcmpl-8rdiGfVkiZqrE1p2TxeKnCxc2zzUb... \n",
" 3 {'id': 'chatcmpl-8rdiGdKxmYC6mS4QSuiW3HjsMRW05... \n",
" 4 {'id': 'chatcmpl-8rdiGqTe3MVy8FGJ6qtVAjxmMODy6... \n",
"medium 0 {'id': 'chatcmpl-8rdiGsk7ohvIVwzuFMfO3cH9zNLW4... \n",
" 1 {'id': 'chatcmpl-8rdiG64Y66W8CZZ9MI4xiVNHheHiF... \n",
" 2 {'id': 'chatcmpl-8rdiGGmXfXx0uRuKafeKODGQe42vz... \n",
" 3 {'id': 'chatcmpl-8rdiG95XCozUiFXY7ryA4rSfzbwRk... \n",
" 4 {'id': 'chatcmpl-8rdiGi6nF5FLMdPIYrkrYgI15Yfcw... "
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from oaib import Batch\n",
"\n",
"n = 5\n",
"batch = Batch(rpm=1000, tpm=10000, workers=5, index=[\"difficulty\", \"i\"])\n",
"difficulties = [\"easy\", \"medium\", \"hard\"]\n",
"\n",
"for difficulty in difficulties:\n",
" for i in range(n):\n",
" await batch.add(\n",
" \"chat.completions.create\",\n",
" metadata={\"difficulty\": difficulty, \"i\": i},\n",
" model=\"gpt-3.5-turbo\",\n",
" messages=[{\n",
" \"role\": \"user\", \n",
" \"content\": f\"difficulty: {difficulty}. write a math problem.\"\n",
" }]\n",
" )\n",
"\n",
"await batch.run()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"## Notes\n",
"\n",
"1. It is not possible to perfectly guarantee the tokens per minute limit is not\n",
" breached because we cannot know the total token usage until the response\n",
" comes back.\n",
" \n",
" Use the `safety` param to set the rate limit tolerance. By default it is set\n",
" to 10%, and will wait until the predicted TPM (the current TPM plus the\n",
" average number of tokens per request) drops below 90% of the limit.\n",
"\n",
" \n",
"2. By default, important logs are stored at `oaib.txt`. This can be disabled\n",
" using `loglevel=0`.\n",
"\n",
"3. There's an error with TPM/RPM progress bar display in Jupyter Notebooks for\n",
" the `Auto` class only. This is caused by a `tqdm.notebook` bug where only the\n",
" initial totals (here, our limits) are used to calculate the width of the bar,\n",
" and the `Auto` class updates these values only after the first request. The\n",
" text percentage displays are accurate."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.0.0"
}
},
"nbformat": 4,
"nbformat_minor": 2
}