feat: improve local model support and doc #96

README.md

@@ -205,10 +205,14 @@ https://github.com/OpenBMB/AgentVerse/assets/11704492/4d07da68-f942-4205-b558-f1
   - [Simulation](#simulation)
     - [Framework Required Modules](#framework-required-modules)
     - [CLI Example](#cli-example)
-    - [GUI Example (Local)](#gui-example-local)
+    - [GUI Example](#gui-example)
   - [Task-Solving](#task-solving)
     - [Framework Required Modules](#framework-required-modules-1)
     - [CLI Example](#cli-example-1)
+  - [Local Model Support](#local-model-support)
+    - [1. Install the Additional Dependencies](#1-install-the-additional-dependencies)
+    - [2. Launch the Local Server](#2-launch-the-local-server)
+    - [3. Modify the Config File](#3-modify-the-config-file)
 - [AgentVerse Showcases](#agentverse-showcases)
   - [Simulation Showcases](#simulation-showcases)
   - [Task-Solving Showcases](#task-solving-showcases)
@@ -282,7 +286,7 @@ You can create a multi-agent environments provided by us. Using the classroom sc
 agentverse-simulation --task simulation/nlp_classroom_9players
 ```
 
-### GUI Example (Local)
+### GUI Example
 
 We also provide a local website demo for this environment. You can launch it with
 
@@ -338,6 +342,33 @@ We have provided more tasks in `agentverse/tasks/tasksolving/tool_using/` that s
 
 Also, you can take a look at `agentverse/tasks/tasksolving` for more experiments we have done in our paper.
 
+## Local Model Support
+### 1. Install the Additional Dependencies
+If you want to use local models such as LLaMA, you need to additionally install some other dependencies:
+```bash
+pip install -r requirements_local.txt
+```
+
+### 2. Launch the Local Server
+Then modify the `MODEL_PATH` and `MODEL_NAME` according to your need to launch the local server with the following command:
+```bash
+bash scripts/run_local_model_server.sh
+```
+The script will launch a service for Llama 7B chat model.
+The `MODEL_NAME` in AgentVerse currently supports several models including `llama-2-7b-chat-hf`, `llama-2-13b-chat-hf`, `llama-2-70b-chat-hf`, `vicuna-7b-v1.5`, and `vicuna-13b-v1.5`. If you wish to integrate additional models that are [compatible with FastChat](https://github.com/lm-sys/FastChat/blob/main/docs/model_support.md), you need to:
+1. Add the new `MODEL_NAME` into the `LOCAL_LLMS` within `agentverse/llms/__init__.py`. Furthermore, establish
+2. Add the mapping from the new `MODEL_NAME` to its corresponding Huggingface identifier in the `LOCAL_LLMS_MAPPING` within the `agentverse/llms/__init__.py` file.
+
+### 3. Modify the Config File
+In your config file, set the `llm_type` to `local` and `model` to the `MODEL_NAME`. For example
+```yaml
+llm:
+  llm_type: local
+  model: llama-2-7b-chat-hf
+  ...
+```
+
+You can refer to `agentverse/tasks/tasksolving/commongen/llama-2-7b-chat-hf/config.yaml` for a more detailed example.
 
 # AgentVerse Showcases
 

agentverse/llms/__init__.py

@@ -1,7 +1,20 @@
 from agentverse.registry import Registry
 
 llm_registry = Registry(name="LLMRegistry")
-LOCAL_LLMS = ["llama-2-7b-chat-hf"]
+LOCAL_LLMS = [
+    "llama-2-7b-chat-hf",
+    "llama-2-13b-chat-hf",
+    "llama-2-70b-chat-hf",
+    "vicuna-7b-v1.5",
+    "vicuna-13b-v1.5",
+]
+LOCAL_LLMS_MAPPING = {
+    "llama-2-7b-chat-hf": "meta-llama/Llama-2-7b-chat-hf",
+    "llama-2-13b-chat-hf": "meta-llama/Llama-2-13b-chat-hf",
+    "llama-2-70b-chat-hf": "meta-llama/Llama-2-70b-chat-hf",
+    "vicuna-7b-v1.5": "lmsys/vicuna-7b-v1.5",
+    "vicuna-13b-v1.5": "lmsys/vicuna-13b-v1.5",
+}
 
 from .base import BaseLLM, BaseChatModel, BaseCompletionModel, LLMResult
 from .openai import OpenAIChat

agentverse/llms/openai.py

@@ -97,7 +97,7 @@ class OpenAIChatArgs(BaseModelArgs):
 @llm_registry.register("gpt-35-turbo")
 @llm_registry.register("gpt-3.5-turbo")
 @llm_registry.register("gpt-4")
-@llm_registry.register("llama-2-7b-chat-hf")
+@llm_registry.register("local")
 class OpenAIChat(BaseChatModel):
     args: OpenAIChatArgs = Field(default_factory=OpenAIChatArgs)
 

agentverse/llms/utils/token_counter.py

@@ -4,7 +4,7 @@
 from typing import List, Union, Dict
 from agentverse.logging import logger
 from agentverse.message import Message
-from agentverse.llms import LOCAL_LLMS
+from agentverse.llms import LOCAL_LLMS, LOCAL_LLMS_MAPPING
 
 
 def count_string_tokens(prompt: str = "", model: str = "gpt-3.5-turbo") -> int:
@@ -27,9 +27,10 @@ def count_message_tokens(
         tokens_per_message = 3
         tokens_per_name = 1
         encoding_model = "gpt-4"
-    elif model in LOCAL_LLMS:
+    elif model.lower() in LOCAL_LLMS or model in LOCAL_LLMS:
         from transformers import AutoTokenizer
-        encoding = AutoTokenizer.from_pretrained(model)
+
+        encoding = AutoTokenizer.from_pretrained(LOCAL_LLMS_MAPPING[model.lower()])
     else:
         raise NotImplementedError(
             f"count_message_tokens() is not implemented for model {model}.\n"

agentverse/tasks/tasksolving/commongen/llama-2-7b-chat-hf/config.yaml

@@ -96,8 +96,8 @@ agents:
     memory:
       memory_type: chat_history
     llm:
-      llm_type: llama-2-7b-chat-hf
-      model: "llama-2-7b-chat-hf"
+      llm_type: local
+      model: llama-2-7b-chat-hf
       temperature: 0
       max_tokens: 512
     output_parser:
@@ -113,8 +113,8 @@ agents:
     memory:
       memory_type: chat_history
     llm:
-      llm_type: llama-2-7b-chat-hf
-      model: "llama-2-7b-chat-hf"
+      llm_type: local
+      model: llama-2-7b-chat-hf
       temperature: 0
       max_tokens: 1024
     output_parser:
@@ -138,8 +138,8 @@ agents:
     memory:
       memory_type: chat_history
     llm:
-      llm_type: llama-2-7b-chat-hf
-      model: "llama-2-7b-chat-hf"
+      llm_type: local
+      model: llama-2-7b-chat-hf
       temperature: 0
       max_tokens: 1024
     output_parser:
@@ -154,7 +154,7 @@ agents:
     memory:
       memory_type: chat_history
     llm:
-      llm_type: llama-2-7b-chat-hf
+      llm_type: local
       model: llama-2-7b-chat-hf
       temperature: 0
       max_tokens: 1024
@@ -172,7 +172,7 @@ agents:
     memory:
       memory_type: chat_history
     llm:
-      llm_type: llama-2-7b-chat-hf
+      llm_type: local
       model: llama-2-7b-chat-hf
       temperature: 0.3
       max_tokens: 1024
@@ -189,8 +189,8 @@ agents:
     memory:
       memory_type: chat_history
     llm:
-      llm_type: llama-2-7b-chat-hf
-      model: "llama-2-7b-chat-hf"
+      llm_type: local
+      model: llama-2-7b-chat-hf
       temperature: 0
       max_tokens: 1024
     output_parser:

scripts/run_local_model_server.sh

@@ -1,9 +1,8 @@
-:<<COMMENT
-See https://github.com/lm-sys/FastChat/blob/main/docs/openai_api.md for more usages.
-COMMENT
+# See https://github.com/lm-sys/FastChat/blob/main/docs/openai_api.md for more usages.
+
 # export CUDA_VISIBLE_DEVICES=0
-MODEL_PATH="path_to_the_downloaded_model_dir"
-MODEL_NAME="name_of_the_model"
+MODEL_PATH="meta-llama/Llama-2-7b-chat-hf"   # path_to_the_downloaded_model_dir
+MODEL_NAME="llama-2-7b-chat-hf"              # name_of_the_model
 python3 -m fastchat.serve.controller & \
 python3 -m fastchat.serve.multi_model_worker \
     --model-path ${MODEL_PATH} \