how_to_count_token.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3d63c55d-a6fd-4c06-ad91-246fc6e606d5",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/UpstageAI/cookbook/blob/main/how_to_count_token.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>\n",
    "\n",
    "In this notebook, we'll see how to use Solar tokenizer to count the number of tokens in a message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8175bd06-4255-41f3-b17d-3849cc281792",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Install requirements\n",
    "!pip install -q openai tokenizers\n",
    "!pip install -q python-dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "140d2612-b6d8-4a78-b8a9-01a11c601495",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title set API key\n",
    "import os\n",
    "import getpass\n",
    "from pprint import pprint\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from IPython import get_ipython\n",
    "\n",
    "if \"google.colab\" in str(get_ipython()):\n",
    "    # Running in Google Colab. Please set the UPSTAGE_API_KEY in the Colab Secrets\n",
    "    from google.colab import userdata\n",
    "    os.environ[\"UPSTAGE_API_KEY\"] = userdata.get(\"UPSTAGE_API_KEY\")\n",
    "else:\n",
    "    # Running locally. Please set the UPSTAGE_API_KEY in the .env file\n",
    "    from dotenv import load_dotenv\n",
    "\n",
    "    load_dotenv()\n",
    "\n",
    "if \"UPSTAGE_API_KEY\" not in os.environ:\n",
    "    os.environ[\"UPSTAGE_API_KEY\"] = getpass.getpass(\"Enter your Upstage API key: \")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1b6250fd-17aa-4f6b-afbe-ce4e4a8574f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Set Solar client\n",
    "from openai import OpenAI\n",
    "from tokenizers import Tokenizer\n",
    "\n",
    "client = OpenAI(api_key=UPSTAGE_API_KEY, base_url=\"https://api.upstage.ai/v1/solar\")\n",
    "tokenizer = Tokenizer.from_pretrained(\"upstage/solar-1-mini-tokenizer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "11a0eff3-4652-4aab-bced-576e7fa8fe70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Define messages\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"Your name is Solar. You should answer politely.\"},\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"Hi!\",\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": \"How can I help you?\",\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"his tutorial demonstrates how to combine the semantic search results from a $vectorSearch query with full-text search results from a $search query by using reciprocal rank fusion. Reciprocal rank fusion is a way to combine results from different types of searches into a single result. This tutorial takes you through the following steps: Create an Atlas Vector Search index on the plot_embeddings field in the sample_mflix.embedded_movies collection.Create an Atlas Search index on the title field in the sample_mflix.embedded_movies collection. Combine and run a $vectorSearch query against the plot_embeddings field in the sample_mflix.embedded_movies collection with a $search query against the title field by using reciprocal rank fusion.PrerequisitesBefore you begin, ensure that your Atlas cluster meets the requirements described in the Prerequisites. NOTE Ensure that your Atlas cluster has enough memory to store both Atlas Search and Atlas Vector Search indexes and run performant queries. Create the Atlas Vector Search and Atlas Search Indexes This section demonstrates how to create the following indexes on the fields in the sample_mflix.embedded_movies collection: An Atlas Vector Search index on the plot_.\",\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3f5d4d66-cdb3-456e-9551-4ceb330c235f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title count messages tokens using tokenizer\n",
    "def count_tokens_from_messages(tokenizer, messages):\n",
    "    tokens_per_message = 5  # <|im_start|>{role}\\n{message}<|im_end|>\n",
    "    tokens_prefix = 1  # <|startoftext|>\n",
    "    tokens_suffix = 3  # <|im_start|>assistant\\n\n",
    "\n",
    "    num_tokens = 0\n",
    "    num_tokens += 1\n",
    "    for m in messages:\n",
    "        num_tokens += tokens_per_message\n",
    "        for _, v in m.items():\n",
    "            num_tokens += len(tokenizer.encode(v, add_special_tokens=False))\n",
    "\n",
    "    num_tokens += tokens_suffix\n",
    "    return num_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "78358a20-b8a1-402e-85eb-aac1196474f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Define messages\n",
    "messages = [{\"role\": \"system\", \"content\": \"Hello World!\"}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0632b2a6-5025-4088-80c2-510f59a9dd45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13 prompt tokens counted by the Solar tokenizer\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    count_tokens_from_messages(tokenizer, messages),\n",
    "    \"prompt tokens counted by the Solar tokenizer\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "3cf9d8fd-fbfc-4e71-a4e5-a1cb55df0cdc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "318 prompt tokens counted by the Solar.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# @title Count tokens using solar API\n",
    "response = client.chat.completions.create(\n",
    "    model=\"solar-1-mini-chat\", messages=messages, temperature=0\n",
    ")\n",
    "print(f\"{response.usage.prompt_tokens} prompt tokens counted by the Solar.\")\n",
    "print()"
   ]
  }
 ],
 "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",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}