run-llama--llama_index/docs/examples/llm/nvidia_text_completion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NVIDIA LLM Text Completion API\n",
    "\n",
    "The `llama-index-llms-nvidia` package extends the `NVIDIA` class to support the `/completions` API for code completion models such as:\n",
    "\n",
    "- `bigcode/starcoder2-7b`\n",
    "- `bigcode/starcoder2-15b`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet llama-index-llms-nvidia"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "**To get started:**\n",
    "\n",
    "1. Create a free account with [NVIDIA](https://build.nvidia.com/), which hosts NVIDIA AI Foundation models.\n",
    "\n",
    "2. Click on your model of choice.\n",
    "\n",
    "3. Under Input select the Python tab, and click `Get API Key`. Then click `Generate Key`.\n",
    "\n",
    "4. Copy and save the generated key as NVIDIA_API_KEY. From there, you should have access to the endpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "# del os.environ['NVIDIA_API_KEY']  ## delete key and reset\n",
    "if os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    print(\"Valid NVIDIA_API_KEY already in environment. Delete to reset\")\n",
    "else:\n",
    "    nvapi_key = getpass.getpass(\"NVAPI Key (starts with nvapi-): \")\n",
    "    assert nvapi_key.startswith(\n",
    "        \"nvapi-\"\n",
    "    ), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# llama-parse is async-first, running the async code in a notebook requires the use of nest_asyncio\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with the NVIDIA API Catalog\n",
    "\n",
    "### Usage of the `use_chat_completions` argument\n",
    "\n",
    "Set `None` (default) to decide per-invocation whether to use `/chat/completions` or `/completions` endpoints with query keyword arguments.\n",
    "\n",
    "- Set `False` to use the `/completions` endpoint.\n",
    "- Set `True` to use the `/chat/completions` endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.nvidia import NVIDIA\n",
    "\n",
    "llm = NVIDIA(model=\"bigcode/starcoder2-15b\", use_chat_completions=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Available models\n",
    "\n",
    "Use `is_chat_model` to filter available text completion models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print([model for model in llm.available_models if model.is_chat_model])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with NVIDIA NIMs\n",
    "\n",
    "In addition to connecting to hosted [NVIDIA NIMs](https://ai.nvidia.com), this connector can be used to connect to local NIM instances. This helps you take your applications local when necessary.\n",
    "\n",
    "For instructions on how to set up local NIM instances, refer to [NVIDIA NIM](https://developer.nvidia.com/nim)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.nvidia import NVIDIA\n",
    "\n",
    "# Connect to a NIM running at localhost:8080\n",
    "llm = NVIDIA(base_url=\"http://localhost:8080/v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Complete: `.complete()`\n",
    "\n",
    "We can use `.complete()`/`.acomplete()` (which takes a string) to prompt a response from the selected model.\n",
    "\n",
    "Let's use our default model for this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(llm.complete(\"# Function that does quicksort:\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, LlamaIndex returns a `CompletionResponse`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Async Complete: `.acomplete()`\n",
    "\n",
    "There is also an async implementation which can be leveraged in the same way!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await llm.acomplete(\"# Function that does quicksort:\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = llm.stream_complete(prompt=\"# Reverse string in python:\", max_tokens=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for t in x:\n",
    "    print(t.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Async Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = await llm.astream_complete(\n",
    "    prompt=\"# Reverse program in python:\", max_tokens=512\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async for t in x:\n",
    "    print(t.delta, end=\"\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}