Add gpt-3.5-turbo-16k support (camel-ai#171)

Author: lightaime

camel/agents/chat_agent.py

@@ -23,11 +23,7 @@
 from camel.messages import ChatMessage, MessageType, SystemMessage
 from camel.models import BaseModelBackend, ModelFactory
 from camel.typing import ModelType, RoleType
-from camel.utils import (
-    get_model_token_limit,
-    num_tokens_from_messages,
-    openai_api_key_required,
-)
+from camel.utils import num_tokens_from_messages, openai_api_key_required
 
 
 @dataclass(frozen=True)
@@ -84,11 +80,11 @@ def __init__(
         self.model: ModelType = (model if model is not None else
                                  ModelType.GPT_3_5_TURBO)
         self.model_config: ChatGPTConfig = model_config or ChatGPTConfig()
-        self.model_token_limit: int = get_model_token_limit(self.model)
         self.message_window_size: Optional[int] = message_window_size
 
         self.model_backend: BaseModelBackend = ModelFactory.create(
             self.model, self.model_config.__dict__)
+        self.model_token_limit: int = self.model_backend.token_limit
 
         self.terminated: bool = False
         self.init_messages()

camel/models/base_model.py

@@ -14,11 +14,24 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List
 
+from camel.typing import ModelType
+
 
 class BaseModelBackend(ABC):
     r"""Base class for different model backends.
     May be OpenAI API, a local LLM, a stub for unit tests, etc."""
 
+    def __init__(self, model_type: ModelType,
+                 model_config_dict: Dict[str, Any]) -> None:
+        r"""Constructor for the model backend.
+
+        Args:
+            model_type (ModelType): Model for which a backend is created.
+            model_config_dict (Dict[str, Any]): A config dictionary.
+        """
+        self.model_type = model_type
+        self.model_config_dict = model_config_dict
+
     @abstractmethod
     def run(self, messages: List[Dict]) -> Dict[str, Any]:
         r"""Runs the query to the backend model.
@@ -35,3 +48,11 @@ def run(self, messages: List[Dict]) -> Dict[str, Any]:
             Dict[str, Any]: All backends must return a dict in OpenAI format.
         """
         pass
+
+    @property
+    def token_limit(self) -> int:
+        r"""Returns the maximum token limit for a given model.
+        Returns:
+            int: The maximum token limit for the given model.
+        """
+        return self.model_type.token_limit

camel/models/model_factory.py

@@ -42,7 +42,10 @@ def create(model_type: ModelType,
         """
         model_class: Any
         if model_type in {
-                ModelType.GPT_3_5_TURBO, ModelType.GPT_4, ModelType.GPT_4_32k
+                ModelType.GPT_3_5_TURBO,
+                ModelType.GPT_3_5_TURBO_16K,
+                ModelType.GPT_4,
+                ModelType.GPT_4_32k,
         }:
             model_class = OpenAIModel
         elif model_type == ModelType.STUB:

camel/models/openai_model.py

@@ -28,18 +28,16 @@ def __init__(self, model_type: ModelType,
         Args:
             model_type (ModelType): Model for which a backend is created,
                 one of GPT_* series.
-            model_config_dict (Dict[str, Any]): a dictionary that will
+            model_config_dict (Dict[str, Any]): A dictionary that will
                 be fed into openai.ChatCompletion.create().
         """
-        super().__init__()
-        self.model_type = model_type
-        self.model_config_dict = model_config_dict
+        super().__init__(model_type, model_config_dict)
 
     def run(self, messages: List[Dict]) -> Dict[str, Any]:
         r"""Run inference of OpenAI chat completion.
 
         Args:
-            messages (List[Dict]): message list with the chat history
+            messages (List[Dict]): Message list with the chat history
                 in OpenAI API format.
 
         Returns:

camel/models/stub_model.py

@@ -14,14 +14,17 @@
 from typing import Any, Dict, List
 
 from camel.models import BaseModelBackend
+from camel.typing import ModelType
 
 
 class StubModel(BaseModelBackend):
     r"""A dummy model used for unit tests."""
+    model_type = ModelType.STUB
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, model_type: ModelType,
+                 model_config_dict: Dict[str, Any]) -> None:
         r"""All arguments are unused for the dummy model."""
-        super().__init__()
+        pass
 
     def run(self, messages: List[Dict]) -> Dict[str, Any]:
         r"""Run fake inference by returning a fixed string.

camel/typing.py

@@ -24,14 +24,34 @@ class RoleType(Enum):
 
 class ModelType(Enum):
     GPT_3_5_TURBO = "gpt-3.5-turbo"
+    GPT_3_5_TURBO_16K = "gpt-3.5-turbo-16k"
     GPT_4 = "gpt-4"
     GPT_4_32k = "gpt-4-32k"
     STUB = "stub"
 
     @property
-    def value_for_tiktoken(self):
+    def value_for_tiktoken(self) -> str:
         return self.value if self.name != "STUB" else "gpt-3.5-turbo"
 
+    @property
+    def token_limit(self) -> int:
+        r"""Returns the maximum token limit for a given model.
+        Returns:
+            int: The maximum token limit for the given model.
+        """
+        if self is ModelType.GPT_3_5_TURBO:
+            return 4096
+        elif self is ModelType.GPT_3_5_TURBO_16K:
+            return 16384
+        elif self is ModelType.GPT_4:
+            return 8192
+        elif self is ModelType.GPT_4_32k:
+            return 32768
+        elif self is ModelType.STUB:
+            return 4096
+        else:
+            raise ValueError("Unknown model type")
+
 
 class TaskType(Enum):
     AI_SOCIETY = "ai_society"

camel/utils.py

@@ -29,27 +29,28 @@
 
 def count_tokens_openai_chat_models(
     messages: List[OpenAIMessage],
-    encoding: Any,
+    encoding: tiktoken.Encoding,
+    tokens_per_message: int,
+    tokens_per_name: int,
 ) -> int:
     r"""Counts the number of tokens required to generate an OpenAI chat based
     on a given list of messages.
 
     Args:
         messages (List[OpenAIMessage]): The list of messages.
-        encoding (Any): The encoding method to use.
+        encoding (tiktoken.Encoding): The encoding method to use.
 
     Returns:
         int: The number of tokens required.
     """
     num_tokens = 0
     for message in messages:
-        # message follows <im_start>{role/name}\n{content}<im_end>\n
-        num_tokens += 4
+        num_tokens += tokens_per_message
         for key, value in message.items():
             num_tokens += len(encoding.encode(value))
             if key == "name":  # if there's a name, the role is omitted
-                num_tokens += -1  # role is always 1 token
-    num_tokens += 2  # every reply is primed with <im_start>assistant
+                num_tokens += tokens_per_name
+    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
     return num_tokens
 
 
@@ -81,11 +82,26 @@ def num_tokens_from_messages(
     except KeyError:
         encoding = tiktoken.get_encoding("cl100k_base")
 
-    if model in {
-            ModelType.GPT_3_5_TURBO, ModelType.GPT_4, ModelType.GPT_4_32k,
-            ModelType.STUB
-    }:
-        return count_tokens_openai_chat_models(messages, encoding)
+    if model.value_for_tiktoken.startswith("gpt-3.5-turbo"):
+        # Every message follows <|start|>{role/name}\n{content}<|end|>\n
+        tokens_per_message = 4
+        # If there's a name, the role is omitted
+        tokens_per_name = -1
+        return count_tokens_openai_chat_models(
+            messages,
+            encoding,
+            tokens_per_message,
+            tokens_per_name,
+        )
+    elif model.value_for_tiktoken.startswith("gpt-4"):
+        tokens_per_message = 3
+        tokens_per_name = 1
+        return count_tokens_openai_chat_models(
+            messages,
+            encoding,
+            tokens_per_message,
+            tokens_per_name,
+        )
     else:
         raise NotImplementedError(
             f"`num_tokens_from_messages`` is not presently implemented "
@@ -97,27 +113,6 @@ def num_tokens_from_messages(
             f"for information about openai chat models.")
 
 
-def get_model_token_limit(model: ModelType) -> int:
-    r"""Returns the maximum token limit for a given model.
-
-    Args:
-        model (ModelType): The type of the model.
-
-    Returns:
-        int: The maximum token limit for the given model.
-    """
-    if model == ModelType.GPT_3_5_TURBO:
-        return 4096
-    elif model == ModelType.GPT_4:
-        return 8192
-    elif model == ModelType.GPT_4_32k:
-        return 32768
-    elif model == ModelType.STUB:
-        return 4096
-    else:
-        raise ValueError("Unknown model type")
-
-
 def openai_api_key_required(func: F) -> F:
     r"""Decorator that checks if the OpenAI API key is available in the
     environment variables.

test/agents/test_chat_agent.py

@@ -18,7 +18,6 @@
 from camel.generators import SystemMessageGenerator
 from camel.messages import ChatMessage, SystemMessage
 from camel.typing import ModelType, RoleType, TaskType
-from camel.utils import get_model_token_limit
 
 parametrize = pytest.mark.parametrize('model', [
     ModelType.STUB,
@@ -28,7 +27,7 @@
 
 
 @parametrize
-def test_chat_agent(model):
+def test_chat_agent(model: ModelType):
 
     model_config = ChatGPTConfig()
     system_msg = SystemMessageGenerator(
@@ -54,7 +53,7 @@ def test_chat_agent(model):
     assert assistant_response.info['id'] is not None
 
     assistant.reset()
-    token_limit = get_model_token_limit(model)
+    token_limit = assistant.model_token_limit
     user_msg = ChatMessage(role_name="Patient", role_type=RoleType.USER,
                            meta_dict=dict(), role="user",
                            content="token" * (token_limit + 1))

test/messages/test_message_base.py

@@ -51,7 +51,7 @@ def test_base_message_contains_operator(base_message: BaseMessage):
 def test_base_message_token_len(base_message: BaseMessage):
     token_len = base_message.token_len()
     assert isinstance(token_len, int)
-    assert token_len == 9
+    assert token_len == 10
 
 
 def test_extract_text_and_code_prompts():

test/models/test_model_factory.py

@@ -18,9 +18,10 @@
 from camel.typing import ModelType
 
 parametrize = pytest.mark.parametrize('model', [
-    ModelType.STUB,
     pytest.param(ModelType.GPT_3_5_TURBO, marks=pytest.mark.model_backend),
+    pytest.param(ModelType.GPT_3_5_TURBO_16K, marks=pytest.mark.model_backend),
     pytest.param(ModelType.GPT_4, marks=pytest.mark.model_backend),
+    ModelType.STUB,
 ])
 
 
@@ -30,18 +31,12 @@ def test_model_factory(model):
     model_inst = ModelFactory.create(model, model_config_dict)
     messages = [
         {
-            'role': 'system',
-            'content': 'You can make a task more specific.'
+            "role": "system",
+            "content": "Initialize system",
         },
         {
-            'role':
-            'user',
-            'content': ('Here is a task that Python Programmer will help '
-                        'Stock Trader to complete: Develop a trading bot '
-                        'for the stock market.\nPlease make it more specific.'
-                        ' Be creative and imaginative.\nPlease reply with '
-                        'the specified task in 50 words or less. '
-                        'Do not add anything else.')
+            "role": "user",
+            "content": "Hello",
         },
     ]
     response = model_inst.run(messages)