# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
# Licensed under the Apache License, Version 2.0 (the “License”);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an “AS IS” BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
from __future__ import annotations
import base64
from abc import ABC, abstractmethod
from io import BytesIO
from math import ceil
from typing import TYPE_CHECKING, List, Optional
from anthropic import Anthropic
from PIL import Image
from camel.types import ModelType, OpenAIImageType, OpenAIVisionDetailType
if TYPE_CHECKING:
from mistral_common.protocol.instruct.request import ( # type:ignore[import-not-found]
ChatCompletionRequest,
)
from camel.messages import OpenAIMessage
LOW_DETAIL_TOKENS = 85
FIT_SQUARE_PIXELS = 2048
SHORTEST_SIDE_PIXELS = 768
SQUARE_PIXELS = 512
SQUARE_TOKENS = 170
EXTRA_TOKENS = 85
[docs]
def messages_to_prompt(messages: List[OpenAIMessage], model: ModelType) -> str:
r"""Parse the message list into a single prompt following model-specific
formats.
Args:
messages (List[OpenAIMessage]): Message list with the chat history
in OpenAI API format.
model (ModelType): Model type for which messages will be parsed.
Returns:
str: A single prompt summarizing all the messages.
"""
system_message = messages[0]["content"]
ret: str
if model in [
ModelType.LLAMA_2,
ModelType.LLAMA_3,
ModelType.GROQ_LLAMA_3_8B,
ModelType.GROQ_LLAMA_3_70B,
]:
# reference: https://github.com/facebookresearch/llama/blob/cfc3fc8c1968d390eb830e65c63865e980873a06/llama/generation.py#L212
seps = [" ", " </s><s>"]
role_map = {"user": "[INST]", "assistant": "[/INST]"}
system_prompt = f"[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n"
ret = ""
for i, msg in enumerate(messages[1:]):
role = role_map[msg["role"]]
content = msg["content"]
if content:
if not isinstance(content, str):
raise ValueError(
"Currently multimodal context is not "
"supported by the token counter."
)
if i == 0:
ret += system_prompt + content
else:
ret += role + " " + content + seps[i % 2]
else:
ret += role
return ret
elif model in [ModelType.VICUNA, ModelType.VICUNA_16K]:
seps = [" ", "</s>"]
role_map = {"user": "USER", "assistant": "ASSISTANT"}
system_prompt = f"{system_message}"
ret = system_prompt + seps[0]
for i, msg in enumerate(messages[1:]):
role = role_map[msg["role"]]
content = msg["content"]
if not isinstance(content, str):
raise ValueError(
"Currently multimodal context is not "
"supported by the token counter."
)
if content:
ret += role + ": " + content + seps[i % 2]
else:
ret += role + ":"
return ret
elif model == ModelType.GLM_4_OPEN_SOURCE:
system_prompt = f"[gMASK]<sop><|system|>\n{system_message}"
ret = system_prompt
for msg in messages[1:]:
role = msg["role"]
content = msg["content"]
if not isinstance(content, str):
raise ValueError(
"Currently multimodal context is not "
"supported by the token counter."
)
if content:
ret += "<|" + role + "|>" + "\n" + content
else:
ret += "<|" + role + "|>" + "\n"
return ret
elif model == ModelType.QWEN_2:
system_prompt = f"<|im_start|>system\n{system_message}<|im_end|>"
ret = system_prompt + "\n"
for msg in messages[1:]:
role = msg["role"]
content = msg["content"]
if not isinstance(content, str):
raise ValueError(
"Currently multimodal context is not "
"supported by the token counter."
)
if content:
ret += (
'<|im_start|>'
+ role
+ '\n'
+ content
+ '<|im_end|>'
+ '\n'
)
else:
ret += '<|im_start|>' + role + '\n'
return ret
elif model == ModelType.GROQ_MIXTRAL_8_7B:
# Mistral/Mixtral format
system_prompt = f"<s>[INST] {system_message} [/INST]\n"
ret = system_prompt
for msg in messages[1:]:
if msg["role"] == "user":
ret += f"[INST] {msg['content']} [/INST]\n"
elif msg["role"] == "assistant":
ret += f"{msg['content']}</s>\n"
if not isinstance(msg['content'], str):
raise ValueError(
"Currently multimodal context is not "
"supported by the token counter."
)
return ret.strip()
elif model in [ModelType.GROQ_GEMMA_7B_IT, ModelType.GROQ_GEMMA_2_9B_IT]:
# Gemma format
ret = f"<bos>{system_message}\n"
for msg in messages:
if msg["role"] == "user":
ret += f"Human: {msg['content']}\n"
elif msg["role"] == "assistant":
ret += f"Assistant: {msg['content']}\n"
if not isinstance(msg['content'], str):
raise ValueError(
"Currently multimodal context is not supported by the token counter."
)
ret += "<eos>"
return ret
else:
raise ValueError(f"Invalid model type: {model}")
[docs]
def get_model_encoding(value_for_tiktoken: str):
r"""Get model encoding from tiktoken.
Args:
value_for_tiktoken: Model value for tiktoken.
Returns:
tiktoken.Encoding: Model encoding.
"""
import tiktoken
try:
encoding = tiktoken.encoding_for_model(value_for_tiktoken)
except KeyError:
if value_for_tiktoken in [
ModelType.O1_MINI.value,
ModelType.O1_PREVIEW.value,
]:
encoding = tiktoken.get_encoding("o200k_base")
else:
print("Model not found. Using cl100k_base encoding.")
encoding = tiktoken.get_encoding("cl100k_base")
return encoding
[docs]
class BaseTokenCounter(ABC):
r"""Base class for token counters of different kinds of models."""
[docs]
@abstractmethod
def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int:
r"""Count number of tokens in the provided message list.
Args:
messages (List[OpenAIMessage]): Message list with the chat history
in OpenAI API format.
Returns:
int: Number of tokens in the messages.
"""
pass
[docs]
class OpenSourceTokenCounter(BaseTokenCounter):
def __init__(self, model_type: ModelType, model_path: str):
r"""Constructor for the token counter for open-source models.
Args:
model_type (ModelType): Model type for which tokens will be
counted.
model_path (str): The path to the model files, where the tokenizer
model should be located.
"""
# Use a fast Rust-based tokenizer if it is supported for a given model.
# If a fast tokenizer is not available for a given model,
# a normal Python-based tokenizer is returned instead.
from transformers import AutoTokenizer
try:
tokenizer = AutoTokenizer.from_pretrained(
model_path,
use_fast=True,
)
except TypeError:
tokenizer = AutoTokenizer.from_pretrained(
model_path,
use_fast=False,
)
except Exception:
raise ValueError(
f"Invalid `model_path` ({model_path}) is provided. "
"Tokenizer loading failed."
)
self.tokenizer = tokenizer
self.model_type = model_type
[docs]
def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int:
r"""Count number of tokens in the provided message list using
loaded tokenizer specific for this type of model.
Args:
messages (List[OpenAIMessage]): Message list with the chat history
in OpenAI API format.
Returns:
int: Number of tokens in the messages.
"""
prompt = messages_to_prompt(messages, self.model_type)
input_ids = self.tokenizer(prompt).input_ids
return len(input_ids)
[docs]
class OpenAITokenCounter(BaseTokenCounter):
def __init__(self, model: ModelType):
r"""Constructor for the token counter for OpenAI models.
Args:
model (ModelType): Model type for which tokens will be counted.
"""
self.model: str = model.value_for_tiktoken
self.model_type = model
self.tokens_per_message: int
self.tokens_per_name: int
if self.model == "gpt-3.5-turbo-0301":
# Every message follows <|start|>{role/name}\n{content}<|end|>\n
self.tokens_per_message = 4
# If there's a name, the role is omitted
self.tokens_per_name = -1
elif ("gpt-3.5-turbo" in self.model) or ("gpt-4" in self.model):
self.tokens_per_message = 3
self.tokens_per_name = 1
elif "o1" in self.model:
self.tokens_per_message = 2
self.tokens_per_name = 1
else:
# flake8: noqa :E501
raise NotImplementedError(
"Token counting for OpenAI Models is not presently "
f"implemented for model {model}. "
"See https://github.com/openai/openai-python/blob/main/chatml.md "
"for information on how messages are converted to tokens. "
"See https://platform.openai.com/docs/models/gpt-4"
"or https://platform.openai.com/docs/models/gpt-3-5"
"for information about openai chat models."
)
self.encoding = get_model_encoding(self.model)
[docs]
def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int:
r"""Count number of tokens in the provided message list with the
help of package tiktoken.
Args:
messages (List[OpenAIMessage]): Message list with the chat history
in OpenAI API format.
Returns:
int: Number of tokens in the messages.
"""
num_tokens = 0
for message in messages:
num_tokens += self.tokens_per_message
for key, value in message.items():
if not isinstance(value, list):
num_tokens += len(self.encoding.encode(str(value)))
else:
for item in value:
if item["type"] == "text":
num_tokens += len(
self.encoding.encode(str(item["text"]))
)
elif item["type"] == "image_url":
image_str: str = item["image_url"]["url"]
detail = item["image_url"]["detail"]
image_prefix_format = "data:image/{};base64,"
image_prefix: Optional[str] = None
for image_type in list(OpenAIImageType):
# Find the correct image format
image_prefix = image_prefix_format.format(
image_type.value
)
if image_prefix in image_str:
break
assert isinstance(image_prefix, str)
encoded_image = image_str.split(image_prefix)[1]
image_bytes = BytesIO(
base64.b64decode(encoded_image)
)
image = Image.open(image_bytes)
num_tokens += self._count_tokens_from_image(
image, OpenAIVisionDetailType(detail)
)
if key == "name":
num_tokens += self.tokens_per_name
# every reply is primed with <|start|>assistant<|message|>
num_tokens += 3
return num_tokens
def _count_tokens_from_image(
self, image: Image.Image, detail: OpenAIVisionDetailType
) -> int:
r"""Count image tokens for OpenAI vision model. An :obj:`"auto"`
resolution model will be treated as :obj:`"high"`. All images with
:obj:`"low"` detail cost 85 tokens each. Images with :obj:`"high"` detail
are first scaled to fit within a 2048 x 2048 square, maintaining their
aspect ratio. Then, they are scaled such that the shortest side of the
image is 768px long. Finally, we count how many 512px squares the image
consists of. Each of those squares costs 170 tokens. Another 85 tokens are
always added to the final total. For more details please refer to `OpenAI
vision docs <https://platform.openai.com/docs/guides/vision>`_
Args:
image (PIL.Image.Image): Image to count number of tokens.
detail (OpenAIVisionDetailType): Image detail type to count
number of tokens.
Returns:
int: Number of tokens for the image given a detail type.
"""
if detail == OpenAIVisionDetailType.LOW:
return LOW_DETAIL_TOKENS
width, height = image.size
if width > FIT_SQUARE_PIXELS or height > FIT_SQUARE_PIXELS:
scaling_factor = max(width, height) / FIT_SQUARE_PIXELS
width = int(width / scaling_factor)
height = int(height / scaling_factor)
scaling_factor = min(width, height) / SHORTEST_SIDE_PIXELS
scaled_width = int(width / scaling_factor)
scaled_height = int(height / scaling_factor)
h = ceil(scaled_height / SQUARE_PIXELS)
w = ceil(scaled_width / SQUARE_PIXELS)
total = EXTRA_TOKENS + SQUARE_TOKENS * h * w
return total
[docs]
class AnthropicTokenCounter(BaseTokenCounter):
def __init__(self, model_type: ModelType):
r"""Constructor for the token counter for Anthropic models.
Args:
model_type (ModelType): Model type for which tokens will be
counted.
"""
self.model_type = model_type
self.client = Anthropic()
self.tokenizer = self.client.get_tokenizer()
[docs]
def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int:
r"""Count number of tokens in the provided message list using
loaded tokenizer specific for this type of model.
Args:
messages (List[OpenAIMessage]): Message list with the chat history
in OpenAI API format.
Returns:
int: Number of tokens in the messages.
"""
num_tokens = 0
for message in messages:
content = str(message["content"])
num_tokens += self.client.count_tokens(content)
return num_tokens
[docs]
class GeminiTokenCounter(BaseTokenCounter):
def __init__(self, model_type: ModelType):
r"""Constructor for the token counter for Gemini models."""
import google.generativeai as genai
self.model_type = model_type
self._client = genai.GenerativeModel(self.model_type.value)
[docs]
def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int:
r"""Count number of tokens in the provided message list using
loaded tokenizer specific for this type of model.
Args:
messages (List[OpenAIMessage]): Message list with the chat history
in OpenAI API format.
Returns:
int: Number of tokens in the messages.
"""
converted_messages = []
for message in messages:
role = message.get('role')
if role == 'assistant':
role_to_gemini = 'model'
else:
role_to_gemini = 'user'
converted_message = {
"role": role_to_gemini,
"parts": message.get("content"),
}
converted_messages.append(converted_message)
return self._client.count_tokens(converted_messages).total_tokens
[docs]
class LiteLLMTokenCounter:
def __init__(self, model_type: str):
r"""Constructor for the token counter for LiteLLM models.
Args:
model_type (str): Model type for which tokens will be counted.
"""
self.model_type = model_type
self._token_counter = None
self._completion_cost = None
@property
def token_counter(self):
if self._token_counter is None:
from litellm import token_counter
self._token_counter = token_counter
return self._token_counter
@property
def completion_cost(self):
if self._completion_cost is None:
from litellm import completion_cost
self._completion_cost = completion_cost
return self._completion_cost
[docs]
def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int:
r"""Count number of tokens in the provided message list using
the tokenizer specific to this type of model.
Args:
messages (List[OpenAIMessage]): Message list with the chat history
in LiteLLM API format.
Returns:
int: Number of tokens in the messages.
"""
return self.token_counter(model=self.model_type, messages=messages)
[docs]
def calculate_cost_from_response(self, response: dict) -> float:
r"""Calculate the cost of the given completion response.
Args:
response (dict): The completion response from LiteLLM.
Returns:
float: The cost of the completion call in USD.
"""
return self.completion_cost(completion_response=response)
[docs]
class MistralTokenCounter(BaseTokenCounter):
def __init__(self, model_type: ModelType):
r"""Constructor for the token counter for Mistral models.
Args:
model_type (ModelType): Model type for which tokens will be
counted.
"""
from mistral_common.tokens.tokenizers.mistral import ( # type:ignore[import-not-found]
MistralTokenizer,
)
self.model_type = model_type
# Determine the model type and set the tokenizer accordingly
model_name = (
"codestral-22b"
if self.model_type
in {ModelType.MISTRAL_CODESTRAL, ModelType.MISTRAL_CODESTRAL_MAMBA}
else self.model_type.value
)
self.tokenizer = MistralTokenizer.from_model(model_name)
[docs]
def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int:
r"""Count number of tokens in the provided message list using
loaded tokenizer specific for this type of model.
Args:
messages (List[OpenAIMessage]): Message list with the chat history
in OpenAI API format.
Returns:
int: Total number of tokens in the messages.
"""
total_tokens = 0
for msg in messages:
tokens = self.tokenizer.encode_chat_completion(
self._convert_response_from_openai_to_mistral(msg)
).tokens
total_tokens += len(tokens)
return total_tokens
def _convert_response_from_openai_to_mistral(
self, openai_msg: OpenAIMessage
) -> ChatCompletionRequest:
r"""Convert an OpenAI message to a Mistral ChatCompletionRequest.
Args:
openai_msg (OpenAIMessage): An individual message with OpenAI
format.
Returns:
ChatCompletionRequest: The converted message in Mistral's request
format.
"""
from mistral_common.protocol.instruct.request import (
ChatCompletionRequest, # type:ignore[import-not-found]
)
mistral_request = ChatCompletionRequest( # type: ignore[type-var]
model=self.model_type.value,
messages=[openai_msg],
)
return mistral_request
