AI21 Jamba Instruct fixes (#3254)

* ai21 labs jamba instruct samples * fix code quality issues * update files --------- Co-authored-by: Shail Shah <shaisha@microsoft.com>
Azure · Jun 21, 2024 · 6081e56 · 6081e56
1 parent 332153b
commit 6081e56
Show file tree

Hide file tree

Showing 5 changed files with 1,064 additions and 0 deletions.
diff --git a/sdk/python/foundation-models/ai21-labs/jamba-instruct/ai21_client.ipynb b/sdk/python/foundation-models/ai21-labs/jamba-instruct/ai21_client.ipynb
@@ -0,0 +1,201 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Use AI21-Jamba-Instruct through Azure AI Models-as-a-Service\n",
+    "\n",
+    "Use AI21's client to consume Jamba-Instruct deployments in Azure AI and Azure ML through serverless API endpoints delivered through Models-as-a-Service (MaaS).\n",
+    "\n",
+    "> Review the documentation for Jamba-Instruct for AI Studio and for ML Studio for details on how to provision inference endpoints, regional availability, pricing and inference schema reference."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Prerequisites\n",
+    "\n",
+    "Before we start, there are certain steps we need to take to deploy the models:\n",
+    "\n",
+    "* Register for a valid Azure account with subscription \n",
+    "* Make sure you have access to [Azure AI Studio](https://learn.microsoft.com/en-us/azure/ai-studio/what-is-ai-studio?tabs=home)\n",
+    "* Create a project and resource group\n",
+    "* Select `AI21-Jamba-Instruct`.\n",
+    "\n",
+    "    > Notice that some models may not be available in all the regions in Azure AI and Azure Machine Learning. On those cases, you can create a workspace or project in the region where the models are available and then consume it with a connection from a different one. To learn more about using connections see [Consume models with connections](https://learn.microsoft.com/en-us/azure/ai-studio/how-to/deployments-connections)\n",
+    "\n",
+    "* Deploy with \"Pay-as-you-go\"\n",
+    "\n",
+    "Once deployed successfully, you should be assigned for an API endpoint and a security key for inference.\n",
+    "\n",
+    "For more information, you should consult Azure's official documentation [here](https://aka.ms/ai21-jamba-instruct-azure-ai-studio-docs) for model deployment and inference.\n",
+    "\n",
+    "To complete this tutorial, you will need to: \n",
+    "\n",
+    "* Install `ai21`:\n",
+    "\n",
+    "    ```bash\n",
+    "    pip install -U \"ai21>=2.6.0\"\n",
+    "    ```\n",
+    "* If it's not working on your first go, try restarting the kernel and then run the pip install again."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Example\n",
+    "\n",
+    "The following is an example about how to use `ai21`'s client on Azure and leveraging this for AI21-Jamba-Instruct through MaaS."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install -U \"ai21>=2.6.0\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "is_executing": true
+    },
+    "name": "imports"
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from ai21 import AI21AzureClient\n",
+    "from ai21.models.chat import ChatMessage"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To use `ai21`, create a client and configure it as follows:\n",
+    "\n",
+    "- `endpoint`: Use the endpoint URL from your deployment. Include `/v1` at the end of the endpoint.\n",
+    "- `api_key`: Use your API key."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "name": "chat_client"
+   },
+   "outputs": [],
+   "source": [
+    "client = AI21AzureClient(base_url=\"<your-maas-endpoint>\", api_key=\"<your-api-key>\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Alternatively, you can set an environment variable for your API key:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.environ[\"AI21_API_KEY\"] = \"<your-api-key>\"\n",
+    "client = AI21AzureClient(\n",
+    "    base_url=\"<your-maas-endpoint>\", api_key=os.environ.get(\"AI21_API_KEY\")\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Use the client to create chat completions requests:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "messages = [\n",
+    "    ChatMessage(content=\"You are a helpful assistant\", role=\"system\"),\n",
+    "    ChatMessage(\n",
+    "        content=\"Who is the most renowned architect in the world? Provide a short poem that describes their work in the style of Shakespeare with Iambic pentimeter and a rhythm pattern of ABABCC\",\n",
+    "        role=\"user\",\n",
+    "    ),\n",
+    "]\n",
+    "\n",
+    "chat_completions = client.chat.completions.create(\n",
+    "    model=\"jamba-instruct\",\n",
+    "    messages=messages,\n",
+    "    temperature=1.0,  # Setting =1 allows for greater variability per API call.\n",
+    "    top_p=1.0,  # Setting =1 allows full sample of tokens to be considered per API call.\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The generated text can be accessed as follows:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "name": "chat_response"
+   },
+   "outputs": [],
+   "source": [
+    "print(chat_completions.to_json())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Aditional resources\n",
+    "\n",
+    "Here are some additional reference:  \n",
+    "\n",
+    "* [Plan and manage costs (marketplace)](https://learn.microsoft.com/azure/ai-studio/how-to/costs-plan-manage#monitor-costs-for-models-offered-through-the-azure-marketplace)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "jupyter",
+   "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.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}