Upload folder using huggingface_hub

README.md

@@ -1,12 +1,6 @@
 ---
-title: Test Fastchat
-emoji: 👁
-colorFrom: pink
-colorTo: green
+title: test_fastchat
+app_file: gradio_web_server.py
 sdk: gradio
 sdk_version: 4.27.0
-app_file: app.py
-pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

api_provider.py

@@ -0,0 +1,454 @@
+"""Call API providers."""
+
+import json
+import os
+import random
+import time
+
+import requests
+
+from fastchat.utils import build_logger
+
+
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
+
+
+def get_api_provider_stream_iter(
+    conv,
+    model_name,
+    model_api_dict,
+    temperature,
+    top_p,
+    max_new_tokens,
+):
+    if model_api_dict["api_type"] == "openai":
+        prompt = conv.to_openai_api_messages()
+        stream_iter = openai_api_stream_iter(
+            model_api_dict["model_name"],
+            prompt,
+            temperature,
+            top_p,
+            max_new_tokens,
+            api_base=model_api_dict["api_base"],
+            api_key=model_api_dict["api_key"],
+        )
+    elif model_api_dict["api_type"] == "anthropic":
+        prompt = conv.get_prompt()
+        stream_iter = anthropic_api_stream_iter(
+            model_name, prompt, temperature, top_p, max_new_tokens
+        )
+    elif model_api_dict["api_type"] == "gemini":
+        stream_iter = gemini_api_stream_iter(
+            model_api_dict["model_name"],
+            conv,
+            temperature,
+            top_p,
+            max_new_tokens,
+            api_key=model_api_dict["api_key"],
+        )
+    elif model_api_dict["api_type"] == "bard":
+        prompt = conv.to_openai_api_messages()
+        stream_iter = bard_api_stream_iter(
+            model_api_dict["model_name"],
+            prompt,
+            temperature,
+            top_p,
+            api_key=model_api_dict["api_key"],
+        )
+    elif model_api_dict["api_type"] == "mistral":
+        prompt = conv.to_openai_api_messages()
+        stream_iter = mistral_api_stream_iter(
+            model_name, prompt, temperature, top_p, max_new_tokens
+        )
+    elif model_api_dict["api_type"] == "nvidia":
+        prompt = conv.to_openai_api_messages()
+        stream_iter = nvidia_api_stream_iter(
+            model_name,
+            prompt,
+            temperature,
+            top_p,
+            max_new_tokens,
+            model_api_dict["api_base"],
+        )
+    elif model_api_dict["api_type"] == "ai2":
+        prompt = conv.to_openai_api_messages()
+        stream_iter = ai2_api_stream_iter(
+            model_name,
+            model_api_dict["model_name"],
+            prompt,
+            temperature,
+            top_p,
+            max_new_tokens,
+            api_base=model_api_dict["api_base"],
+            api_key=model_api_dict["api_key"],
+        )
+    else:
+        raise NotImplementedError()
+
+    return stream_iter
+
+
+def openai_api_stream_iter(
+    model_name,
+    messages,
+    temperature,
+    top_p,
+    max_new_tokens,
+    api_base=None,
+    api_key=None,
+):
+    import openai
+
+    api_key = api_key or os.environ["OPENAI_API_KEY"]
+
+    if "azure" in model_name:
+        client = openai.AzureOpenAI(
+            api_version="2023-07-01-preview",
+            azure_endpoint=api_base or "https://api.openai.com/v1",
+            api_key=api_key,
+        )
+    else:
+        client = openai.OpenAI(
+            base_url=api_base or "https://api.openai.com/v1", api_key=api_key
+        )
+
+    if model_name == "gpt-4-turbo":
+        model_name = "gpt-4-1106-preview"
+
+    # Make requests
+    gen_params = {
+        "model": model_name,
+        "prompt": messages,
+        "temperature": temperature,
+        "top_p": top_p,
+        "max_new_tokens": max_new_tokens,
+    }
+    logger.info(f"==== request ====\n{gen_params}")
+
+    res = client.chat.completions.create(
+        model=model_name,
+        messages=messages,
+        temperature=temperature,
+        max_tokens=max_new_tokens,
+        stream=True,
+    )
+    text = ""
+    for chunk in res:
+        if len(chunk.choices) > 0:
+            text += chunk.choices[0].delta.content or ""
+            data = {
+                "text": text,
+                "error_code": 0,
+            }
+            yield data
+
+
+def anthropic_api_stream_iter(model_name, prompt, temperature, top_p, max_new_tokens):
+    import anthropic
+
+    c = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
+
+    # Make requests
+    gen_params = {
+        "model": model_name,
+        "prompt": prompt,
+        "temperature": temperature,
+        "top_p": top_p,
+        "max_new_tokens": max_new_tokens,
+    }
+    logger.info(f"==== request ====\n{gen_params}")
+
+    res = c.completions.create(
+        prompt=prompt,
+        stop_sequences=[anthropic.HUMAN_PROMPT],
+        max_tokens_to_sample=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        model=model_name,
+        stream=True,
+    )
+    text = ""
+    for chunk in res:
+        text += chunk.completion
+        data = {
+            "text": text,
+            "error_code": 0,
+        }
+        yield data
+
+
+def gemini_api_stream_iter(
+    model_name, conv, temperature, top_p, max_new_tokens, api_key=None
+):
+    import google.generativeai as genai  # pip install google-generativeai
+
+    if api_key is None:
+        api_key = os.environ["GEMINI_API_KEY"]
+    genai.configure(api_key=api_key)
+
+    generation_config = {
+        "temperature": temperature,
+        "max_output_tokens": max_new_tokens,
+        "top_p": top_p,
+    }
+    params = {
+        "model": model_name,
+        "prompt": conv,
+    }
+    params.update(generation_config)
+    logger.info(f"==== request ====\n{params}")
+
+    safety_settings = [
+        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
+        {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
+        {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
+    ]
+    model = genai.GenerativeModel(
+        model_name=model_name,
+        generation_config=generation_config,
+        safety_settings=safety_settings,
+    )
+    history = []
+    for role, message in conv.messages[:-2]:
+        history.append({"role": role, "parts": message})
+    convo = model.start_chat(history=history)
+    response = convo.send_message(conv.messages[-2][1], stream=True)
+
+    try:
+        text = ""
+        for chunk in response:
+            text += chunk.text
+            data = {
+                "text": text,
+                "error_code": 0,
+            }
+            yield data
+    except Exception as e:
+        logger.error(f"==== error ====\n{e}")
+        reason = chunk.candidates
+        yield {
+            "text": f"**API REQUEST ERROR** Reason: {reason}.",
+            "error_code": 1,
+        }
+
+
+def bard_api_stream_iter(model_name, conv, temperature, top_p, api_key=None):
+    del top_p  # not supported
+    del temperature  # not supported
+
+    if api_key is None:
+        api_key = os.environ["BARD_API_KEY"]
+
+    # convert conv to conv_bard
+    conv_bard = []
+    for turn in conv:
+        if turn["role"] == "user":
+            conv_bard.append({"author": "0", "content": turn["content"]})
+        elif turn["role"] == "assistant":
+            conv_bard.append({"author": "1", "content": turn["content"]})
+        else:
+            raise ValueError(f"Unsupported role: {turn['role']}")
+
+    params = {
+        "model": model_name,
+        "prompt": conv_bard,
+    }
+    logger.info(f"==== request ====\n{params}")
+
+    try:
+        res = requests.post(
+            f"https://generativelanguage.googleapis.com/v1beta2/models/{model_name}:generateMessage?key={api_key}",
+            json={
+                "prompt": {
+                    "messages": conv_bard,
+                },
+            },
+            timeout=30,
+        )
+    except Exception as e:
+        logger.error(f"==== error ====\n{e}")
+        yield {
+            "text": f"**API REQUEST ERROR** Reason: {e}.",
+            "error_code": 1,
+        }
+
+    if res.status_code != 200:
+        logger.error(f"==== error ==== ({res.status_code}): {res.text}")
+        yield {
+            "text": f"**API REQUEST ERROR** Reason: status code {res.status_code}.",
+            "error_code": 1,
+        }
+
+    response_json = res.json()
+    if "candidates" not in response_json:
+        logger.error(f"==== error ==== response blocked: {response_json}")
+        reason = response_json["filters"][0]["reason"]
+        yield {
+            "text": f"**API REQUEST ERROR** Reason: {reason}.",
+            "error_code": 1,
+        }
+
+    response = response_json["candidates"][0]["content"]
+    pos = 0
+    while pos < len(response):
+        # simulate token streaming
+        pos += random.randint(3, 6)
+        time.sleep(0.002)
+        data = {
+            "text": response[:pos],
+            "error_code": 0,
+        }
+        yield data
+
+
+def ai2_api_stream_iter(
+    model_name,
+    model_id,
+    messages,
+    temperature,
+    top_p,
+    max_new_tokens,
+    api_key=None,
+    api_base=None,
+):
+    # get keys and needed values
+    ai2_key = api_key or os.environ.get("AI2_API_KEY")
+    api_base = api_base or "https://inferd.allen.ai/api/v1/infer"
+
+    # Make requests
+    gen_params = {
+        "model": model_name,
+        "prompt": messages,
+        "temperature": temperature,
+        "top_p": top_p,
+        "max_new_tokens": max_new_tokens,
+    }
+    logger.info(f"==== request ====\n{gen_params}")
+
+    # AI2 uses vLLM, which requires that `top_p` be 1.0 for greedy sampling:
+    # https://github.com/vllm-project/vllm/blob/v0.1.7/vllm/sampling_params.py#L156-L157
+    if temperature == 0.0 and top_p < 1.0:
+        raise ValueError("top_p must be 1 when temperature is 0.0")
+
+    res = requests.post(
+        api_base,
+        stream=True,
+        headers={"Authorization": f"Bearer {ai2_key}"},
+        json={
+            "model_id": model_id,
+            # This input format is specific to the Tulu2 model. Other models
+            # may require different input formats. See the model's schema
+            # documentation on InferD for more information.
+            "input": {
+                "messages": messages,
+                "opts": {
+                    "max_tokens": max_new_tokens,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "logprobs": 1,  # increase for more choices
+                },
+            },
+        },
+        timeout=5,
+    )
+
+    if res.status_code != 200:
+        logger.error(f"unexpected response ({res.status_code}): {res.text}")
+        raise ValueError("unexpected response from InferD", res)
+
+    text = ""
+    for line in res.iter_lines():
+        if line:
+            part = json.loads(line)
+            if "result" in part and "output" in part["result"]:
+                for t in part["result"]["output"]["text"]:
+                    text += t
+            else:
+                logger.error(f"unexpected part: {part}")
+                raise ValueError("empty result in InferD response")
+
+            data = {
+                "text": text,
+                "error_code": 0,
+            }
+            yield data
+
+
+def mistral_api_stream_iter(model_name, messages, temperature, top_p, max_new_tokens):
+    from mistralai.client import MistralClient
+    from mistralai.models.chat_completion import ChatMessage
+
+    api_key = os.environ["MISTRAL_API_KEY"]
+
+    client = MistralClient(api_key=api_key)
+
+    # Make requests
+    gen_params = {
+        "model": model_name,
+        "prompt": messages,
+        "temperature": temperature,
+        "top_p": top_p,
+        "max_new_tokens": max_new_tokens,
+    }
+    logger.info(f"==== request ====\n{gen_params}")
+
+    new_messages = [
+        ChatMessage(role=message["role"], content=message["content"])
+        for message in messages
+    ]
+
+    res = client.chat_stream(
+        model=model_name,
+        temperature=temperature,
+        messages=new_messages,
+        max_tokens=max_new_tokens,
+        top_p=top_p,
+    )
+
+    text = ""
+    for chunk in res:
+        if chunk.choices[0].delta.content is not None:
+            text += chunk.choices[0].delta.content
+            data = {
+                "text": text,
+                "error_code": 0,
+            }
+            yield data
+
+
+def nvidia_api_stream_iter(model_name, messages, temp, top_p, max_tokens, api_base):
+    assert model_name in ["llama2-70b-steerlm-chat", "yi-34b-chat"]
+
+    api_key = os.environ["NVIDIA_API_KEY"]
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "accept": "text/event-stream",
+        "content-type": "application/json",
+    }
+    # nvidia api does not accept 0 temperature
+    if temp == 0.0:
+        temp = 0.0001
+
+    payload = {
+        "messages": messages,
+        "temperature": temp,
+        "top_p": top_p,
+        "max_tokens": max_tokens,
+        "seed": 42,
+        "stream": True,
+    }
+    logger.info(f"==== request ====\n{payload}")
+
+    response = requests.post(
+        api_base, headers=headers, json=payload, stream=True, timeout=1
+    )
+    text = ""
+    for line in response.iter_lines():
+        if line:
+            data = line.decode("utf-8")
+            if data.endswith("[DONE]"):
+                break
+            data = json.loads(data[6:])["choices"][0]["delta"]["content"]
+            text += data
+            yield {"text": text, "error_code": 0}

base_model_worker.py

@@ -0,0 +1,241 @@
+import asyncio
+import threading
+import time
+from typing import List
+
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse, JSONResponse
+import requests
+
+from fastchat.constants import WORKER_HEART_BEAT_INTERVAL
+from fastchat.conversation import Conversation
+from fastchat.utils import pretty_print_semaphore, build_logger
+
+
+worker = None
+logger = None
+
+app = FastAPI()
+
+
+def heart_beat_worker(obj):
+    while True:
+        time.sleep(WORKER_HEART_BEAT_INTERVAL)
+        obj.send_heart_beat()
+
+
+class BaseModelWorker:
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        conv_template: str = None,
+        multimodal: bool = False,
+    ):
+        global logger, worker
+
+        self.controller_addr = controller_addr
+        self.worker_addr = worker_addr
+        self.worker_id = worker_id
+        if model_path.endswith("/"):
+            model_path = model_path[:-1]
+        self.model_names = model_names or [model_path.split("/")[-1]]
+        self.limit_worker_concurrency = limit_worker_concurrency
+        self.conv = self.make_conv_template(conv_template, model_path)
+        self.conv.sep_style = int(self.conv.sep_style)
+        self.multimodal = multimodal
+        self.tokenizer = None
+        self.context_len = None
+        self.call_ct = 0
+        self.semaphore = None
+
+        self.heart_beat_thread = None
+
+        if logger is None:
+            logger = build_logger("model_worker", f"model_worker_{self.worker_id}.log")
+        if worker is None:
+            worker = self
+
+    def make_conv_template(
+        self,
+        conv_template: str = None,
+        model_path: str = None,
+    ) -> Conversation:
+        """
+        can be overrided to costomize the conversation template for different model workers.
+        """
+        from fastchat.conversation import get_conv_template
+        from fastchat.model.model_adapter import get_conversation_template
+
+        if conv_template:
+            conv = get_conv_template(conv_template)
+        else:
+            conv = get_conversation_template(model_path)
+        return conv
+
+    def init_heart_beat(self):
+        self.register_to_controller()
+        self.heart_beat_thread = threading.Thread(
+            target=heart_beat_worker,
+            args=(self,),
+            daemon=True,
+        )
+        self.heart_beat_thread.start()
+
+    def register_to_controller(self):
+        logger.info("Register to controller")
+
+        url = self.controller_addr + "/register_worker"
+        data = {
+            "worker_name": self.worker_addr,
+            "check_heart_beat": True,
+            "worker_status": self.get_status(),
+            "multimodal": self.multimodal,
+        }
+        r = requests.post(url, json=data)
+        assert r.status_code == 200
+
+    def send_heart_beat(self):
+        logger.info(
+            f"Send heart beat. Models: {self.model_names}. "
+            f"Semaphore: {pretty_print_semaphore(self.semaphore)}. "
+            f"call_ct: {self.call_ct}. "
+            f"worker_id: {self.worker_id}. "
+        )
+
+        url = self.controller_addr + "/receive_heart_beat"
+
+        while True:
+            try:
+                ret = requests.post(
+                    url,
+                    json={
+                        "worker_name": self.worker_addr,
+                        "queue_length": self.get_queue_length(),
+                    },
+                    timeout=5,
+                )
+                exist = ret.json()["exist"]
+                break
+            except (requests.exceptions.RequestException, KeyError) as e:
+                logger.error(f"heart beat error: {e}")
+            time.sleep(5)
+
+        if not exist:
+            self.register_to_controller()
+
+    def get_queue_length(self):
+        if self.semaphore is None:
+            return 0
+        else:
+            sempahore_value = (
+                self.semaphore._value
+                if self.semaphore._value is not None
+                else self.limit_worker_concurrency
+            )
+            waiter_count = (
+                0 if self.semaphore._waiters is None else len(self.semaphore._waiters)
+            )
+            return self.limit_worker_concurrency - sempahore_value + waiter_count
+
+    def get_status(self):
+        return {
+            "model_names": self.model_names,
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
+        }
+
+    def count_token(self, params):
+        prompt = params["prompt"]
+
+        try:
+            input_ids = self.tokenizer(prompt).input_ids
+            input_echo_len = len(input_ids)
+        except TypeError:
+            input_echo_len = self.tokenizer.num_tokens(prompt)
+
+        ret = {
+            "count": input_echo_len,
+            "error_code": 0,
+        }
+        return ret
+
+    def get_conv_template(self):
+        return {"conv": self.conv}
+
+    def generate_stream_gate(self, params):
+        raise NotImplementedError
+
+    def generate_gate(self, params):
+        raise NotImplementedError
+
+    def get_embeddings(self, params):
+        raise NotImplementedError
+
+
+def release_worker_semaphore():
+    worker.semaphore.release()
+
+
+def acquire_worker_semaphore():
+    if worker.semaphore is None:
+        worker.semaphore = asyncio.Semaphore(worker.limit_worker_concurrency)
+    return worker.semaphore.acquire()
+
+
+def create_background_tasks():
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(release_worker_semaphore)
+    return background_tasks
+
+
+@app.post("/worker_generate_stream")
+async def api_generate_stream(request: Request):
+    params = await request.json()
+    await acquire_worker_semaphore()
+    generator = worker.generate_stream_gate(params)
+    background_tasks = create_background_tasks()
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_generate")
+async def api_generate(request: Request):
+    params = await request.json()
+    await acquire_worker_semaphore()
+    output = await asyncio.to_thread(worker.generate_gate, params)
+    release_worker_semaphore()
+    return JSONResponse(output)
+
+
+@app.post("/worker_get_embeddings")
+async def api_get_embeddings(request: Request):
+    params = await request.json()
+    await acquire_worker_semaphore()
+    embedding = worker.get_embeddings(params)
+    release_worker_semaphore()
+    return JSONResponse(content=embedding)
+
+
+@app.post("/worker_get_status")
+async def api_get_status(request: Request):
+    return worker.get_status()
+
+
+@app.post("/count_token")
+async def api_count_token(request: Request):
+    params = await request.json()
+    return worker.count_token(params)
+
+
+@app.post("/worker_get_conv_template")
+async def api_get_conv(request: Request):
+    return worker.get_conv_template()
+
+
+@app.post("/model_details")
+async def api_model_details(request: Request):
+    return {"context_length": worker.context_len}

call_monitor.py

@@ -0,0 +1,219 @@
+import json
+import os
+import glob
+import time
+
+from fastapi import FastAPI
+import hashlib
+import asyncio
+
+REFRESH_INTERVAL_SEC = 60
+LOG_DIR = "/home/vicuna/fastchat_logs/server0"
+# LOG_DIR = "/home/vicuna/tmp/test_env"
+
+
+class Monitor:
+    """Monitor the number of calls to each model."""
+
+    def __init__(self, log_dir: str):
+        self.log_dir = log_dir
+        self.model_call = {}
+        self.user_call = {}
+        self.model_call_limit_global = {
+            "gpt-4-1106-preview": 300,
+            "gpt-4-0125-preview": 300,
+        }
+        self.model_call_day_limit_per_user = {"gpt-4-1106-preview": 10}
+
+    async def update_stats(self, num_file=1) -> None:
+        while True:
+            # find the latest num_file log under log_dir
+            json_files = glob.glob(os.path.join(self.log_dir, "*.json"))
+            json_files.sort(key=os.path.getctime, reverse=True)
+            json_files = json_files[:num_file]
+
+            model_call = {}
+            user_call = {}
+            for json_file in json_files:
+                for line in open(json_file, "r", encoding="utf-8"):
+                    obj = json.loads(line)
+                    if obj["type"] != "chat":
+                        continue
+                    if obj["model"] not in model_call:
+                        model_call[obj["model"]] = []
+                    model_call[obj["model"]].append(
+                        {"tstamp": obj["tstamp"], "user_id": obj["ip"]}
+                    )
+                    if obj["ip"] not in user_call:
+                        user_call[obj["ip"]] = []
+                    user_call[obj["ip"]].append(
+                        {"tstamp": obj["tstamp"], "model": obj["model"]}
+                    )
+
+            self.model_call = model_call
+            self.model_call_stats_hour = self.get_model_call_stats(top_k=None)
+            self.model_call_stats_day = self.get_model_call_stats(
+                top_k=None, most_recent_min=24 * 60
+            )
+
+            self.user_call = user_call
+            self.user_call_stats_hour = self.get_user_call_stats(top_k=None)
+            self.user_call_stats_day = self.get_user_call_stats(
+                top_k=None, most_recent_min=24 * 60
+            )
+            await asyncio.sleep(REFRESH_INTERVAL_SEC)
+
+    def get_model_call_limit(self, model: str) -> int:
+        if model not in self.model_call_limit_global:
+            return -1
+        return self.model_call_limit_global[model]
+
+    def update_model_call_limit(self, model: str, limit: int) -> bool:
+        if model not in self.model_call_limit_global:
+            return False
+        self.model_call_limit_global[model] = limit
+        return True
+
+    def is_model_limit_reached(self, model: str) -> bool:
+        if model not in self.model_call_limit_global:
+            return False
+        if model not in self.model_call_stats_hour:
+            return False
+        # check if the model call limit is reached
+        if self.model_call_stats_hour[model] >= self.model_call_limit_global[model]:
+            return True
+        return False
+
+    def is_user_limit_reached(self, model: str, user_id: str) -> bool:
+        if model not in self.model_call_day_limit_per_user:
+            return False
+        if user_id not in self.user_call_stats_day:
+            return False
+        if model not in self.user_call_stats_day[user_id]["call_dict"]:
+            return False
+        # check if the user call limit is reached
+        if (
+            self.user_call_stats_day[user_id]["call_dict"][model]
+            >= self.model_call_day_limit_per_user[model]
+        ):
+            return True
+        return False
+
+    def get_model_call_stats(
+        self, target_model=None, most_recent_min: int = 60, top_k: int = 20
+    ) -> dict:
+        model_call_stats = {}
+        for model, reqs in self.model_call.items():
+            if target_model is not None and model != target_model:
+                continue
+            model_call = []
+            for req in reqs:
+                if req["tstamp"] < time.time() - most_recent_min * 60:
+                    continue
+                model_call.append(req["tstamp"])
+            model_call_stats[model] = len(model_call)
+        if top_k is not None:
+            top_k_model = sorted(
+                model_call_stats, key=lambda x: model_call_stats[x], reverse=True
+            )[:top_k]
+            model_call_stats = {model: model_call_stats[model] for model in top_k_model}
+        return model_call_stats
+
+    def get_user_call_stats(
+        self, target_model=None, most_recent_min: int = 60, top_k: int = 20
+    ) -> dict:
+        user_call_stats = {}
+        for user_id, reqs in self.user_call.items():
+            user_model_call = {"call_dict": {}}
+            for req in reqs:
+                if req["tstamp"] < time.time() - most_recent_min * 60:
+                    continue
+                if target_model is not None and req["model"] != target_model:
+                    continue
+                if req["model"] not in user_model_call["call_dict"]:
+                    user_model_call["call_dict"][req["model"]] = 0
+                user_model_call["call_dict"][req["model"]] += 1
+
+            user_model_call["total_calls"] = sum(user_model_call["call_dict"].values())
+            if user_model_call["total_calls"] > 0:
+                user_call_stats[user_id] = user_model_call
+        if top_k is not None:
+            top_k_user = sorted(
+                user_call_stats,
+                key=lambda x: user_call_stats[x]["total_calls"],
+                reverse=True,
+            )[:top_k]
+            user_call_stats = {
+                user_id: user_call_stats[user_id] for user_id in top_k_user
+            }
+        return user_call_stats
+
+    def get_num_users(self, most_recent_min: int = 60) -> int:
+        user_call_stats = self.get_user_call_stats(
+            most_recent_min=most_recent_min, top_k=None
+        )
+        return len(user_call_stats)
+
+
+monitor = Monitor(log_dir=LOG_DIR)
+app = FastAPI()
+
+
+@app.on_event("startup")
+async def app_startup():
+    asyncio.create_task(monitor.update_stats(2))
+
+
+@app.get("/get_model_call_limit/{model}")
+async def get_model_call_limit(model: str):
+    return {"model_call_limit": {model: monitor.get_model_call_limit(model)}}
+
+
+@app.get("/update_model_call_limit/{model}/{limit}")
+async def update_model_call_limit(model: str, limit: int):
+    if not monitor.update_model_call_limit(model, limit):
+        return {"success": False}
+    return {"success": True}
+
+
+@app.get("/is_limit_reached")
+async def is_limit_reached(model: str, user_id: str):
+    if monitor.is_model_limit_reached(model):
+        return {
+            "is_limit_reached": True,
+            "reason": f"MODEL_HOURLY_LIMIT ({model}): {monitor.get_model_call_limit(model)}",
+        }
+    if monitor.is_user_limit_reached(model, user_id):
+        return {
+            "is_limit_reached": True,
+            "reason": f"USER_DAILY_LIMIT ({model}): {monitor.model_call_day_limit_per_user[model]}",
+        }
+    return {"is_limit_reached": False}
+
+
+@app.get("/get_num_users_hr")
+async def get_num_users():
+    return {"num_users": len(monitor.user_call_stats_hour)}
+
+
+@app.get("/get_num_users_day")
+async def get_num_users_day():
+    return {"num_users": len(monitor.user_call_stats_day)}
+
+
+@app.get("/get_user_call_stats")
+async def get_user_call_stats(
+    model: str = None, most_recent_min: int = 60, top_k: int = None
+):
+    return {
+        "user_call_stats": monitor.get_user_call_stats(model, most_recent_min, top_k)
+    }
+
+
+@app.get("/get_model_call_stats")
+async def get_model_call_stats(
+    model: str = None, most_recent_min: int = 60, top_k: int = None
+):
+    return {
+        "model_call_stats": monitor.get_model_call_stats(model, most_recent_min, top_k)
+    }

cli.py

@@ -0,0 +1,304 @@
+"""
+Chat with a model with command line interface.
+
+Usage:
+python3 -m fastchat.serve.cli --model lmsys/vicuna-7b-v1.5
+python3 -m fastchat.serve.cli --model lmsys/fastchat-t5-3b-v1.0
+
+Other commands:
+- Type "!!exit" or an empty line to exit.
+- Type "!!reset" to start a new conversation.
+- Type "!!remove" to remove the last prompt.
+- Type "!!regen" to regenerate the last message.
+- Type "!!save <filename>" to save the conversation history to a json file.
+- Type "!!load <filename>" to load a conversation history from a json file.
+"""
+import argparse
+import os
+import re
+import sys
+
+from prompt_toolkit import PromptSession
+from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
+from prompt_toolkit.completion import WordCompleter
+from prompt_toolkit.history import InMemoryHistory
+from prompt_toolkit.key_binding import KeyBindings
+from rich.console import Console
+from rich.live import Live
+from rich.markdown import Markdown
+import torch
+
+from fastchat.model.model_adapter import add_model_args
+from fastchat.modules.awq import AWQConfig
+from fastchat.modules.exllama import ExllamaConfig
+from fastchat.modules.xfastertransformer import XftConfig
+from fastchat.modules.gptq import GptqConfig
+from fastchat.serve.inference import ChatIO, chat_loop
+from fastchat.utils import str_to_torch_dtype
+
+
+class SimpleChatIO(ChatIO):
+    def __init__(self, multiline: bool = False):
+        self._multiline = multiline
+
+    def prompt_for_input(self, role) -> str:
+        if not self._multiline:
+            return input(f"{role}: ")
+
+        prompt_data = []
+        line = input(f"{role} [ctrl-d/z on empty line to end]: ")
+        while True:
+            prompt_data.append(line.strip())
+            try:
+                line = input()
+            except EOFError as e:
+                break
+        return "\n".join(prompt_data)
+
+    def prompt_for_output(self, role: str):
+        print(f"{role}: ", end="", flush=True)
+
+    def stream_output(self, output_stream):
+        pre = 0
+        for outputs in output_stream:
+            output_text = outputs["text"]
+            output_text = output_text.strip().split(" ")
+            now = len(output_text) - 1
+            if now > pre:
+                print(" ".join(output_text[pre:now]), end=" ", flush=True)
+                pre = now
+        print(" ".join(output_text[pre:]), flush=True)
+        return " ".join(output_text)
+
+    def print_output(self, text: str):
+        print(text)
+
+
+class RichChatIO(ChatIO):
+    bindings = KeyBindings()
+
+    @bindings.add("escape", "enter")
+    def _(event):
+        event.app.current_buffer.newline()
+
+    def __init__(self, multiline: bool = False, mouse: bool = False):
+        self._prompt_session = PromptSession(history=InMemoryHistory())
+        self._completer = WordCompleter(
+            words=["!!exit", "!!reset", "!!remove", "!!regen", "!!save", "!!load"],
+            pattern=re.compile("$"),
+        )
+        self._console = Console()
+        self._multiline = multiline
+        self._mouse = mouse
+
+    def prompt_for_input(self, role) -> str:
+        self._console.print(f"[bold]{role}:")
+        # TODO(suquark): multiline input has some issues. fix it later.
+        prompt_input = self._prompt_session.prompt(
+            completer=self._completer,
+            multiline=False,
+            mouse_support=self._mouse,
+            auto_suggest=AutoSuggestFromHistory(),
+            key_bindings=self.bindings if self._multiline else None,
+        )
+        self._console.print()
+        return prompt_input
+
+    def prompt_for_output(self, role: str):
+        self._console.print(f"[bold]{role.replace('/', '|')}:")
+
+    def stream_output(self, output_stream):
+        """Stream output from a role."""
+        # TODO(suquark): the console flickers when there is a code block
+        #  above it. We need to cut off "live" when a code block is done.
+
+        # Create a Live context for updating the console output
+        with Live(console=self._console, refresh_per_second=4) as live:
+            # Read lines from the stream
+            for outputs in output_stream:
+                if not outputs:
+                    continue
+                text = outputs["text"]
+                # Render the accumulated text as Markdown
+                # NOTE: this is a workaround for the rendering "unstandard markdown"
+                #  in rich. The chatbots output treat "\n" as a new line for
+                #  better compatibility with real-world text. However, rendering
+                #  in markdown would break the format. It is because standard markdown
+                #  treat a single "\n" in normal text as a space.
+                #  Our workaround is adding two spaces at the end of each line.
+                #  This is not a perfect solution, as it would
+                #  introduce trailing spaces (only) in code block, but it works well
+                #  especially for console output, because in general the console does not
+                #  care about trailing spaces.
+                lines = []
+                for line in text.splitlines():
+                    lines.append(line)
+                    if line.startswith("```"):
+                        # Code block marker - do not add trailing spaces, as it would
+                        #  break the syntax highlighting
+                        lines.append("\n")
+                    else:
+                        lines.append("  \n")
+                markdown = Markdown("".join(lines))
+                # Update the Live console output
+                live.update(markdown)
+        self._console.print()
+        return text
+
+    def print_output(self, text: str):
+        self.stream_output([{"text": text}])
+
+
+class ProgrammaticChatIO(ChatIO):
+    def prompt_for_input(self, role) -> str:
+        contents = ""
+        # `end_sequence` signals the end of a message. It is unlikely to occur in
+        #  message content.
+        end_sequence = " __END_OF_A_MESSAGE_47582648__\n"
+        len_end = len(end_sequence)
+        while True:
+            if len(contents) >= len_end:
+                last_chars = contents[-len_end:]
+                if last_chars == end_sequence:
+                    break
+            try:
+                char = sys.stdin.read(1)
+                contents = contents + char
+            except EOFError:
+                continue
+        contents = contents[:-len_end]
+        print(f"[!OP:{role}]: {contents}", flush=True)
+        return contents
+
+    def prompt_for_output(self, role: str):
+        print(f"[!OP:{role}]: ", end="", flush=True)
+
+    def stream_output(self, output_stream):
+        pre = 0
+        for outputs in output_stream:
+            output_text = outputs["text"]
+            output_text = output_text.strip().split(" ")
+            now = len(output_text) - 1
+            if now > pre:
+                print(" ".join(output_text[pre:now]), end=" ", flush=True)
+                pre = now
+        print(" ".join(output_text[pre:]), flush=True)
+        return " ".join(output_text)
+
+    def print_output(self, text: str):
+        print(text)
+
+
+def main(args):
+    if args.gpus:
+        if len(args.gpus.split(",")) < args.num_gpus:
+            raise ValueError(
+                f"Larger --num-gpus ({args.num_gpus}) than --gpus {args.gpus}!"
+            )
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
+        os.environ["XPU_VISIBLE_DEVICES"] = args.gpus
+    if args.enable_exllama:
+        exllama_config = ExllamaConfig(
+            max_seq_len=args.exllama_max_seq_len,
+            gpu_split=args.exllama_gpu_split,
+            cache_8bit=args.exllama_cache_8bit,
+        )
+    else:
+        exllama_config = None
+    if args.enable_xft:
+        xft_config = XftConfig(
+            max_seq_len=args.xft_max_seq_len,
+            data_type=args.xft_dtype,
+        )
+        if args.device != "cpu":
+            print("xFasterTransformer now is only support CPUs. Reset device to CPU")
+            args.device = "cpu"
+    else:
+        xft_config = None
+    if args.style == "simple":
+        chatio = SimpleChatIO(args.multiline)
+    elif args.style == "rich":
+        chatio = RichChatIO(args.multiline, args.mouse)
+    elif args.style == "programmatic":
+        chatio = ProgrammaticChatIO()
+    else:
+        raise ValueError(f"Invalid style for console: {args.style}")
+    try:
+        chat_loop(
+            args.model_path,
+            args.device,
+            args.num_gpus,
+            args.max_gpu_memory,
+            str_to_torch_dtype(args.dtype),
+            args.load_8bit,
+            args.cpu_offloading,
+            args.conv_template,
+            args.conv_system_msg,
+            args.temperature,
+            args.repetition_penalty,
+            args.max_new_tokens,
+            chatio,
+            gptq_config=GptqConfig(
+                ckpt=args.gptq_ckpt or args.model_path,
+                wbits=args.gptq_wbits,
+                groupsize=args.gptq_groupsize,
+                act_order=args.gptq_act_order,
+            ),
+            awq_config=AWQConfig(
+                ckpt=args.awq_ckpt or args.model_path,
+                wbits=args.awq_wbits,
+                groupsize=args.awq_groupsize,
+            ),
+            exllama_config=exllama_config,
+            xft_config=xft_config,
+            revision=args.revision,
+            judge_sent_end=args.judge_sent_end,
+            debug=args.debug,
+            history=not args.no_history,
+        )
+    except KeyboardInterrupt:
+        print("exit...")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    add_model_args(parser)
+    parser.add_argument(
+        "--conv-template", type=str, default=None, help="Conversation prompt template."
+    )
+    parser.add_argument(
+        "--conv-system-msg", type=str, default=None, help="Conversation system message."
+    )
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--repetition_penalty", type=float, default=1.0)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--no-history", action="store_true")
+    parser.add_argument(
+        "--style",
+        type=str,
+        default="simple",
+        choices=["simple", "rich", "programmatic"],
+        help="Display style.",
+    )
+    parser.add_argument(
+        "--multiline",
+        action="store_true",
+        help="Enable multiline input. Use ESC+Enter for newline.",
+    )
+    parser.add_argument(
+        "--mouse",
+        action="store_true",
+        help="[Rich Style]: Enable mouse support for cursor positioning.",
+    )
+    parser.add_argument(
+        "--judge-sent-end",
+        action="store_true",
+        help="Whether enable the correction logic that interrupts the output of sentences due to EOS.",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Print useful debug information (e.g., prompts)",
+    )
+    args = parser.parse_args()
+    main(args)

controller.py

@@ -0,0 +1,389 @@
+"""
+A controller manages distributed workers.
+It sends worker addresses to clients.
+"""
+import argparse
+import asyncio
+import dataclasses
+from enum import Enum, auto
+import json
+import logging
+import os
+import time
+from typing import List, Union
+import threading
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+import numpy as np
+import requests
+import uvicorn
+
+from fastchat.constants import (
+    CONTROLLER_HEART_BEAT_EXPIRATION,
+    WORKER_API_TIMEOUT,
+    ErrorCode,
+    SERVER_ERROR_MSG,
+)
+from fastchat.utils import build_logger
+
+
+logger = build_logger("controller", "controller.log")
+
+
+class DispatchMethod(Enum):
+    LOTTERY = auto()
+    SHORTEST_QUEUE = auto()
+
+    @classmethod
+    def from_str(cls, name):
+        if name == "lottery":
+            return cls.LOTTERY
+        elif name == "shortest_queue":
+            return cls.SHORTEST_QUEUE
+        else:
+            raise ValueError(f"Invalid dispatch method")
+
+
+@dataclasses.dataclass
+class WorkerInfo:
+    model_names: List[str]
+    speed: int
+    queue_length: int
+    check_heart_beat: bool
+    last_heart_beat: str
+    multimodal: bool
+
+
+def heart_beat_controller(controller):
+    while True:
+        time.sleep(CONTROLLER_HEART_BEAT_EXPIRATION)
+        controller.remove_stale_workers_by_expiration()
+
+
+class Controller:
+    def __init__(self, dispatch_method: str):
+        # Dict[str -> WorkerInfo]
+        self.worker_info = {}
+        self.dispatch_method = DispatchMethod.from_str(dispatch_method)
+
+        self.heart_beat_thread = threading.Thread(
+            target=heart_beat_controller, args=(self,)
+        )
+        self.heart_beat_thread.start()
+
+    def register_worker(
+        self,
+        worker_name: str,
+        check_heart_beat: bool,
+        worker_status: dict,
+        multimodal: bool,
+    ):
+        if worker_name not in self.worker_info:
+            logger.info(f"Register a new worker: {worker_name}")
+        else:
+            logger.info(f"Register an existing worker: {worker_name}")
+
+        if not worker_status:
+            worker_status = self.get_worker_status(worker_name)
+        if not worker_status:
+            return False
+
+        self.worker_info[worker_name] = WorkerInfo(
+            worker_status["model_names"],
+            worker_status["speed"],
+            worker_status["queue_length"],
+            check_heart_beat,
+            time.time(),
+            multimodal,
+        )
+
+        logger.info(f"Register done: {worker_name}, {worker_status}")
+        return True
+
+    def get_worker_status(self, worker_name: str):
+        try:
+            r = requests.post(worker_name + "/worker_get_status", timeout=5)
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Get status fails: {worker_name}, {e}")
+            return None
+
+        if r.status_code != 200:
+            logger.error(f"Get status fails: {worker_name}, {r}")
+            return None
+
+        return r.json()
+
+    def remove_worker(self, worker_name: str):
+        del self.worker_info[worker_name]
+
+    def refresh_all_workers(self):
+        old_info = dict(self.worker_info)
+        self.worker_info = {}
+
+        for w_name, w_info in old_info.items():
+            if not self.register_worker(
+                w_name, w_info.check_heart_beat, None, w_info.multimodal
+            ):
+                logger.info(f"Remove stale worker: {w_name}")
+
+    def list_models(self):
+        model_names = set()
+
+        for w_name, w_info in self.worker_info.items():
+            model_names.update(w_info.model_names)
+
+        return list(model_names)
+
+    def list_multimodal_models(self):
+        model_names = set()
+
+        for w_name, w_info in self.worker_info.items():
+            if w_info.multimodal:
+                model_names.update(w_info.model_names)
+
+        return list(model_names)
+
+    def list_language_models(self):
+        model_names = set()
+
+        for w_name, w_info in self.worker_info.items():
+            if not w_info.multimodal:
+                model_names.update(w_info.model_names)
+
+        return list(model_names)
+
+    def get_worker_address(self, model_name: str):
+        if self.dispatch_method == DispatchMethod.LOTTERY:
+            worker_names = []
+            worker_speeds = []
+            for w_name, w_info in self.worker_info.items():
+                if model_name in w_info.model_names:
+                    worker_names.append(w_name)
+                    worker_speeds.append(w_info.speed)
+            worker_speeds = np.array(worker_speeds, dtype=np.float32)
+            norm = np.sum(worker_speeds)
+            if norm < 1e-4:
+                return ""
+            worker_speeds = worker_speeds / norm
+            if True:  # Directly return address
+                pt = np.random.choice(np.arange(len(worker_names)), p=worker_speeds)
+                worker_name = worker_names[pt]
+                return worker_name
+
+            # Check status before returning
+            while True:
+                pt = np.random.choice(np.arange(len(worker_names)), p=worker_speeds)
+                worker_name = worker_names[pt]
+
+                if self.get_worker_status(worker_name):
+                    break
+                else:
+                    self.remove_worker(worker_name)
+                    worker_speeds[pt] = 0
+                    norm = np.sum(worker_speeds)
+                    if norm < 1e-4:
+                        return ""
+                    worker_speeds = worker_speeds / norm
+                    continue
+            return worker_name
+        elif self.dispatch_method == DispatchMethod.SHORTEST_QUEUE:
+            worker_names = []
+            worker_qlen = []
+            for w_name, w_info in self.worker_info.items():
+                if model_name in w_info.model_names:
+                    worker_names.append(w_name)
+                    worker_qlen.append(w_info.queue_length / w_info.speed)
+            if len(worker_names) == 0:
+                return ""
+            min_index = np.argmin(worker_qlen)
+            w_name = worker_names[min_index]
+            self.worker_info[w_name].queue_length += 1
+            logger.info(
+                f"names: {worker_names}, queue_lens: {worker_qlen}, ret: {w_name}"
+            )
+            return w_name
+        else:
+            raise ValueError(f"Invalid dispatch method: {self.dispatch_method}")
+
+    def receive_heart_beat(self, worker_name: str, queue_length: int):
+        if worker_name not in self.worker_info:
+            logger.info(f"Receive unknown heart beat. {worker_name}")
+            return False
+
+        self.worker_info[worker_name].queue_length = queue_length
+        self.worker_info[worker_name].last_heart_beat = time.time()
+        logger.info(f"Receive heart beat. {worker_name}")
+        return True
+
+    def remove_stale_workers_by_expiration(self):
+        expire = time.time() - CONTROLLER_HEART_BEAT_EXPIRATION
+        to_delete = []
+        for worker_name, w_info in self.worker_info.items():
+            if w_info.check_heart_beat and w_info.last_heart_beat < expire:
+                to_delete.append(worker_name)
+
+        for worker_name in to_delete:
+            self.remove_worker(worker_name)
+
+    def handle_no_worker(self, params):
+        logger.info(f"no worker: {params['model']}")
+        ret = {
+            "text": SERVER_ERROR_MSG,
+            "error_code": ErrorCode.CONTROLLER_NO_WORKER,
+        }
+        return json.dumps(ret).encode() + b"\0"
+
+    def handle_worker_timeout(self, worker_address):
+        logger.info(f"worker timeout: {worker_address}")
+        ret = {
+            "text": SERVER_ERROR_MSG,
+            "error_code": ErrorCode.CONTROLLER_WORKER_TIMEOUT,
+        }
+        return json.dumps(ret).encode() + b"\0"
+
+    # Let the controller act as a worker to achieve hierarchical
+    # management. This can be used to connect isolated sub networks.
+    def worker_api_get_status(self):
+        model_names = set()
+        speed = 0
+        queue_length = 0
+
+        for w_name in self.worker_info:
+            worker_status = self.get_worker_status(w_name)
+            if worker_status is not None:
+                model_names.update(worker_status["model_names"])
+                speed += worker_status["speed"]
+                queue_length += worker_status["queue_length"]
+
+        model_names = sorted(list(model_names))
+        return {
+            "model_names": model_names,
+            "speed": speed,
+            "queue_length": queue_length,
+        }
+
+    def worker_api_generate_stream(self, params):
+        worker_addr = self.get_worker_address(params["model"])
+        if not worker_addr:
+            yield self.handle_no_worker(params)
+
+        try:
+            response = requests.post(
+                worker_addr + "/worker_generate_stream",
+                json=params,
+                stream=True,
+                timeout=WORKER_API_TIMEOUT,
+            )
+            for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+                if chunk:
+                    yield chunk + b"\0"
+        except requests.exceptions.RequestException as e:
+            yield self.handle_worker_timeout(worker_addr)
+
+
+app = FastAPI()
+
+
+@app.post("/register_worker")
+async def register_worker(request: Request):
+    data = await request.json()
+    controller.register_worker(
+        data["worker_name"],
+        data["check_heart_beat"],
+        data.get("worker_status", None),
+        data.get("multimodal", False),
+    )
+
+
+@app.post("/refresh_all_workers")
+async def refresh_all_workers():
+    models = controller.refresh_all_workers()
+
+
+@app.post("/list_models")
+async def list_models():
+    models = controller.list_models()
+    return {"models": models}
+
+
+@app.post("/list_multimodal_models")
+async def list_multimodal_models():
+    models = controller.list_multimodal_models()
+    return {"models": models}
+
+
+@app.post("/list_language_models")
+async def list_language_models():
+    models = controller.list_language_models()
+    return {"models": models}
+
+
+@app.post("/get_worker_address")
+async def get_worker_address(request: Request):
+    data = await request.json()
+    addr = controller.get_worker_address(data["model"])
+    return {"address": addr}
+
+
+@app.post("/receive_heart_beat")
+async def receive_heart_beat(request: Request):
+    data = await request.json()
+    exist = controller.receive_heart_beat(data["worker_name"], data["queue_length"])
+    return {"exist": exist}
+
+
+@app.post("/worker_generate_stream")
+async def worker_api_generate_stream(request: Request):
+    params = await request.json()
+    generator = controller.worker_api_generate_stream(params)
+    return StreamingResponse(generator)
+
+
+@app.post("/worker_get_status")
+async def worker_api_get_status(request: Request):
+    return controller.worker_api_get_status()
+
+
+@app.get("/test_connection")
+async def worker_api_get_status(request: Request):
+    return "success"
+
+
+def create_controller():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21001)
+    parser.add_argument(
+        "--dispatch-method",
+        type=str,
+        choices=["lottery", "shortest_queue"],
+        default="shortest_queue",
+    )
+    parser.add_argument(
+        "--ssl",
+        action="store_true",
+        required=False,
+        default=False,
+        help="Enable SSL. Requires OS Environment variables 'SSL_KEYFILE' and 'SSL_CERTFILE'.",
+    )
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    controller = Controller(args.dispatch_method)
+    return args, controller
+
+
+if __name__ == "__main__":
+    args, controller = create_controller()
+    if args.ssl:
+        uvicorn.run(
+            app,
+            host=args.host,
+            port=args.port,
+            log_level="info",
+            ssl_keyfile=os.environ["SSL_KEYFILE"],
+            ssl_certfile=os.environ["SSL_CERTFILE"],
+        )
+    else:
+        uvicorn.run(app, host=args.host, port=args.port, log_level="info")

gradio_block_arena_anony.py

@@ -0,0 +1,811 @@
+"""
+Chatbot Arena (battle) tab.
+Users chat with two anonymous models.
+"""
+
+import json
+import time
+
+import gradio as gr
+import numpy as np
+
+from fastchat.constants import (
+    MODERATION_MSG,
+    CONVERSATION_LIMIT_MSG,
+    SLOW_MODEL_MSG,
+    INPUT_CHAR_LEN_LIMIT,
+    CONVERSATION_TURN_LIMIT,
+)
+from fastchat.model.model_adapter import get_conversation_template
+from fastchat.serve.gradio_block_arena_named import flash_buttons
+from fastchat.serve.gradio_web_server import (
+    State,
+    bot_response,
+    get_conv_log_filename,
+    no_change_btn,
+    enable_btn,
+    disable_btn,
+    invisible_btn,
+    acknowledgment_md,
+    get_ip,
+    get_model_description_md,
+)
+from fastchat.utils import (
+    build_logger,
+    moderation_filter,
+)
+
+logger = build_logger("gradio_web_server_multi", "gradio_web_server_multi.log")
+
+num_sides = 2
+enable_moderation = False
+anony_names = ["", ""]
+models = []
+
+
+def set_global_vars_anony(enable_moderation_):
+    global enable_moderation
+    enable_moderation = enable_moderation_
+
+
+def load_demo_side_by_side_anony(models_, url_params):
+    global models
+    models = models_
+
+    states = (None,) * num_sides
+    selector_updates = (
+        gr.Markdown(visible=True),
+        gr.Markdown(visible=True),
+    )
+
+    return states + selector_updates
+
+
+def vote_last_response(states, vote_type, model_selectors, request: gr.Request):
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(time.time(), 4),
+            "type": vote_type,
+            "models": [x for x in model_selectors],
+            "states": [x.dict() for x in states],
+            "ip": get_ip(request),
+        }
+        fout.write(json.dumps(data) + "\n")
+
+    if ":" not in model_selectors[0]:
+        for i in range(5):
+            names = (
+                "### Model A: " + states[0].model_name,
+                "### Model B: " + states[1].model_name,
+            )
+            yield names + ("",) + (disable_btn,) * 4
+            time.sleep(0.1)
+    else:
+        names = (
+            "### Model A: " + states[0].model_name,
+            "### Model B: " + states[1].model_name,
+        )
+        yield names + ("",) + (disable_btn,) * 4
+
+
+def leftvote_last_response(
+    state0, state1, model_selector0, model_selector1, request: gr.Request
+):
+    logger.info(f"leftvote (anony). ip: {get_ip(request)}")
+    for x in vote_last_response(
+        [state0, state1], "leftvote", [model_selector0, model_selector1], request
+    ):
+        yield x
+
+
+def rightvote_last_response(
+    state0, state1, model_selector0, model_selector1, request: gr.Request
+):
+    logger.info(f"rightvote (anony). ip: {get_ip(request)}")
+    for x in vote_last_response(
+        [state0, state1], "rightvote", [model_selector0, model_selector1], request
+    ):
+        yield x
+
+
+def tievote_last_response(
+    state0, state1, model_selector0, model_selector1, request: gr.Request
+):
+    logger.info(f"tievote (anony). ip: {get_ip(request)}")
+    for x in vote_last_response(
+        [state0, state1], "tievote", [model_selector0, model_selector1], request
+    ):
+        yield x
+
+
+def bothbad_vote_last_response(
+    state0, state1, model_selector0, model_selector1, request: gr.Request
+):
+    logger.info(f"bothbad_vote (anony). ip: {get_ip(request)}")
+    for x in vote_last_response(
+        [state0, state1], "bothbad_vote", [model_selector0, model_selector1], request
+    ):
+        yield x
+
+
+def regenerate(state0, state1, request: gr.Request):
+    logger.info(f"regenerate (anony). ip: {get_ip(request)}")
+    states = [state0, state1]
+    for i in range(num_sides):
+        states[i].conv.update_last_message(None)
+    return states + [x.to_gradio_chatbot() for x in states] + [""] + [disable_btn] * 6
+
+
+def clear_history(request: gr.Request):
+    logger.info(f"clear_history (anony). ip: {get_ip(request)}")
+    return (
+        [None] * num_sides
+        + [None] * num_sides
+        + anony_names
+        + [""]
+        + [invisible_btn] * 4
+        + [disable_btn] * 2
+        + [""]
+    )
+
+
+def share_click(state0, state1, model_selector0, model_selector1, request: gr.Request):
+    logger.info(f"share (anony). ip: {get_ip(request)}")
+    if state0 is not None and state1 is not None:
+        vote_last_response(
+            [state0, state1], "share", [model_selector0, model_selector1], request
+        )
+
+
+SAMPLING_WEIGHTS = {
+    # tier 0
+    "gpt-4": 4,
+    "gpt-4-0314": 4,
+    "gpt-4-0613": 4,
+    "gpt-4-turbo": 4,
+    "gpt-4-1106-preview": 4,
+    "gpt-4-0125-preview": 4,
+    "gpt-3.5-turbo-0613": 2,
+    "gpt-3.5-turbo-1106": 2,
+    "gpt-3.5-turbo-0125": 4,
+    "claude-2.1": 4,
+    "claude-2.0": 2,
+    "claude-1": 2,
+    "claude-instant-1": 2,
+    "gemini-pro": 4,
+    "gemini-pro-dev-api": 4,
+    "bard-jan-24-gemini-pro": 4,
+    "bard-feb-2024": 4,
+    "mixtral-8x7b-instruct-v0.1": 4,
+    "mistral-medium": 4,
+    "qwen1.5-72b-chat": 4,
+    "qwen1.5-7b-chat": 2,
+    "qwen1.5-4b-chat": 2,
+    "nous-hermes-2-mixtral-8x7b-dpo": 2,
+    "deepseek-llm-67b-chat": 2,
+    "stripedhyena-nous-7b": 2,
+    "openchat-3.5-0106": 2,
+    "mistral-7b-instruct-v0.2": 2,
+    "solar-10.7b-instruct-v1.0": 2,
+    "dolphin-2.2.1-mistral-7b": 2,
+    "starling-lm-7b-alpha": 2,
+    "tulu-2-dpo-70b": 2,
+    "yi-34b-chat": 2,
+    "zephyr-7b-beta": 2,
+    # tier 1
+    "deluxe-chat-v1.2": 4,
+    "llama-2-70b-chat": 4,
+    "llama-2-13b-chat": 2,
+    "llama-2-7b-chat": 2,
+    "mistral-7b-instruct": 2,
+    "codellama-34b-instruct": 1.5,
+    "vicuna-33b": 2,
+    "vicuna-13b": 1.5,
+    "wizardlm-13b": 1.5,
+    "qwen-14b-chat": 1.5,
+    # tier 2
+    "pplx-7b-online": 1,
+    "pplx-70b-online": 1,
+    "openhermes-2.5-mistral-7b": 1.0,
+    "llama2-70b-steerlm-chat": 1.0,
+    "chatglm3-6b": 1.0,
+    "openchat-3.5": 1.0,
+    "wizardlm-70b": 1.0,
+    "vicuna-7b": 1.0,
+    "chatglm2-6b": 1.0,
+    # deprecated
+    "zephyr-7b-alpha": 1.5,
+    "codellama-13b-instruct": 1.0,
+    "mpt-30b-chat": 1.5,
+    "guanaco-33b": 1.0,
+    "fastchat-t5-3b": 0.5,
+    "alpaca-13b": 0.5,
+    "mpt-7b-chat": 0.1,
+    "oasst-pythia-12b": 0.1,
+    "RWKV-4-Raven-14B": 0.1,
+    "gpt4all-13b-snoozy": 0.1,
+    "koala-13b": 0.1,
+    "stablelm-tuned-alpha-7b": 0.1,
+    "dolly-v2-12b": 0.1,
+    "llama-13b": 0.1,
+    "chatglm-6b": 0.5,
+    "deluxe-chat-v1": 4,
+    "palm-2": 1.5,
+}
+
+# target model sampling weights will be boosted.
+BATTLE_TARGETS = {
+    "gpt-4": {"gpt-4-0314", "claude-2.1", "gpt-4-1106-preview"},
+    "gpt-4-0613": {"gpt-4-0314", "claude-2.1", "gpt-4-1106-preview"},
+    "gpt-4-0314": {
+        "gpt-4-1106-preview",
+        "gpt-4-0613",
+        "claude-2.1",
+        "gpt-3.5-turbo-0613",
+    },
+    "gpt-4-1106-preview": {
+        "gpt-4-0613",
+        "gpt-3.5-turbo-0613",
+        "gpt-3.5-turbo-1106",
+        "claude-2.1",
+        "bard-feb-2024",
+    },
+    "gpt-4-0125-preview": {
+        "gpt-4-1106-preview",
+        "gpt-4-0613",
+        "gpt-3.5-turbo-0613",
+        "claude-2.1",
+        "mistral-medium",
+        "bard-feb-2024",
+    },
+    "gpt-3.5-turbo-0613": {"claude-instant-1", "gpt-4-0613", "claude-2.1"},
+    "gpt-3.5-turbo-1106": {"gpt-4-0613", "claude-instant-1", "gpt-3.5-turbo-0613"},
+    "gpt-3.5-turbo-0125": {
+        "gpt-4-0613",
+        "gpt-4-1106-preview",
+        "gpt-3.5-turbo-0613",
+        "gpt-3.5-turbo-1106",
+        "mixtral-8x7b-instruct-v0.1",
+    },
+    "qwen1.5-72b-chat": {
+        "gpt-3.5-turbo-0125",
+        "gpt-4-0613",
+        "gpt-4-1106-preview",
+        "llama-2-70b-chat",
+        "mixtral-8x7b-instruct-v0.1",
+        "mistral-medium",
+        "yi-34b-chat",
+    },
+    "qwen1.5-7b-chat": {
+        "gpt-3.5-turbo-0125",
+        "starling-lm-7b-alpha",
+        "llama-2-70b-chat",
+        "openchat-3.5",
+        "mixtral-8x7b-instruct-v0.1",
+    },
+    "qwen1.5-4b-chat": {
+        "llama-2-70b-chat",
+        "llama-2-13b-chat",
+        "llama-2-7b-chat",
+        "openchat-3.5",
+    },
+    "openchat-3.5-0106": {
+        "gpt-3.5-turbo-0125",
+        "gpt-3.5-turbo-0613",
+        "llama-2-70b-chat",
+        "openchat-3.5",
+        "mixtral-8x7b-instruct-v0.1",
+    },
+    "nous-hermes-2-mixtral-8x7b-dpo": {
+        "gpt-4-1106-preview",
+        "claude-2.1",
+        "mistral-medium",
+        "gpt-3.5-turbo-0613",
+        "mixtral-8x7b-instruct-v0.1",
+    },
+    "mistral-7b-instruct-v0.2": {
+        "llama-2-70b-chat",
+        "mixtral-8x7b-instruct-v0.1",
+        "starling-lm-7b-alpha",
+        "openhermes-2.5-mistral-7b",
+    },
+    "solar-10.7b-instruct-v1.0": {
+        "mixtral-8x7b-instruct-v0.1",
+        "gpt-3.5-turbo-0613",
+        "llama-2-70b-chat",
+    },
+    "mistral-medium": {
+        "gpt-3.5-turbo-0125",
+        "gpt-3.5-turbo-0613",
+        "gpt-4-1106-preview",
+        "mixtral-8x7b-instruct-v0.1",
+        "bard-feb-2024",
+    },
+    "mixtral-8x7b-instruct-v0.1": {
+        "gpt-3.5-turbo-0125",
+        "gpt-3.5-turbo-0613",
+        "gpt-4-1106-preview",
+        "llama-2-70b-chat",
+    },
+    "claude-2.1": {"gpt-4-1106-preview", "gpt-4-0613", "claude-1"},
+    "claude-2.0": {"gpt-4-1106-preview", "gpt-4-0613", "claude-1"},
+    "claude-1": {"claude-2.1", "gpt-4-0613", "gpt-3.5-turbo-0613"},
+    "claude-instant-1": {"gpt-3.5-turbo-0125", "claude-2.1"},
+    "gemini-pro": {"gpt-4-1106-preview", "gpt-4-0613", "gpt-3.5-turbo-0613"},
+    "gemini-pro-dev-api": {
+        "gpt-4-1106-preview",
+        "gpt-4-0613",
+        "gpt-3.5-turbo-0613",
+        "bard-feb-2024",
+    },
+    "bard-jan-24-gemini-pro": {
+        "gpt-4-1106-preview",
+        "gpt-4-0613",
+        "gpt-3.5-turbo-0613",
+        "gemini-pro-dev-api",
+    },
+    "bard-feb-2024": {
+        "gpt-4-1106-preview",
+        "gpt-4-0613",
+        "gpt-3.5-turbo-0613",
+        "bard-jan-24-gemini-pro",
+    },
+    "deepseek-llm-67b-chat": {
+        "gpt-4-1106-preview",
+        "gpt-4-turbo",
+        "gpt-3.5-turbo-0613",
+    },
+    "llama2-70b-steerlm-chat": {
+        "llama-2-70b-chat",
+        "tulu-2-dpo-70b",
+        "yi-34b-chat",
+    },
+    "stripedhyena-nous-7b": {
+        "starling-lm-7b-alpha",
+        "openhermes-2.5-mistral-7b",
+        "mistral-7b-instruct",
+        "llama-2-7b-chat",
+    },
+    "deluxe-chat-v1.1": {"gpt-4-0613", "gpt-4-1106-preview"},
+    "deluxe-chat-v1.2": {"gpt-4-0613", "gpt-4-1106-preview"},
+    "pplx-7b-online": {"gpt-3.5-turbo-0125", "llama-2-70b-chat"},
+    "pplx-70b-online": {"gpt-3.5-turbo-0125", "llama-2-70b-chat"},
+    "openhermes-2.5-mistral-7b": {
+        "gpt-3.5-turbo-0613",
+        "openchat-3.5",
+        "zephyr-7b-beta",
+    },
+    "dolphin-2.2.1-mistral-7b": {
+        "gpt-3.5-turbo-0613",
+        "vicuna-33b",
+        "starling-lm-7b-alpha",
+        "openhermes-2.5-mistral-7b",
+    },
+    "starling-lm-7b-alpha": {"gpt-3.5-turbo-0613", "openchat-3.5", "zephyr-7b-beta"},
+    "tulu-2-dpo-70b": {"gpt-3.5-turbo-0613", "vicuna-33b", "claude-instant-1"},
+    "yi-34b-chat": {"gpt-3.5-turbo-0613", "vicuna-33b", "claude-instant-1"},
+    "openchat-3.5": {"gpt-3.5-turbo-0613", "llama-2-70b-chat", "zephyr-7b-beta"},
+    "chatglm3-6b": {"yi-34b-chat", "qwen-14b-chat"},
+    "qwen-14b-chat": {"vicuna-13b", "llama-2-13b-chat", "llama-2-70b-chat"},
+    "zephyr-7b-alpha": {"mistral-7b-instruct", "llama-2-13b-chat"},
+    "zephyr-7b-beta": {
+        "mistral-7b-instruct",
+        "llama-2-13b-chat",
+        "llama-2-7b-chat",
+        "wizardlm-13b",
+    },
+    "llama-2-70b-chat": {"gpt-3.5-turbo-0125", "claude-instant-1"},
+    "llama-2-13b-chat": {"mistral-7b-instruct", "vicuna-13b", "llama-2-70b-chat"},
+    "llama-2-7b-chat": {"mistral-7b-instruct", "vicuna-7b", "llama-2-13b-chat"},
+    "mistral-7b-instruct": {
+        "llama-2-7b-chat",
+        "llama-2-13b-chat",
+        "llama-2-70b-chat",
+    },
+    "vicuna-33b": {"llama-2-70b-chat", "gpt-3.5-turbo-0613", "claude-instant-1"},
+    "vicuna-13b": {"llama-2-13b-chat", "llama-2-70b-chat"},
+    "vicuna-7b": {"llama-2-7b-chat", "mistral-7b-instruct", "llama-2-13b-chat"},
+    "wizardlm-70b": {"gpt-3.5-turbo-0613", "vicuna-33b", "claude-instant-1"},
+}
+
+SAMPLING_BOOST_MODELS = [
+    # "claude-2.1",
+    # "gpt-4-0613",
+    # "gpt-4-0314",
+    # "gpt-4-1106-preview",
+    # "gpt-4-0125-preview",
+    "gpt-3.5-turbo-0125",
+    # "mistral-medium",
+    "nous-hermes-2-mixtral-8x7b-dpo",
+    "openchat-3.5-0106",
+    "qwen1.5-72b-chat",
+    "qwen1.5-7b-chat",
+    "qwen1.5-4b-chat",
+    # "mistral-7b-instruct-v0.2",
+]
+
+# outage models won't be sampled.
+OUTAGE_MODELS = []
+
+
+def get_sample_weight(model):
+    if model in OUTAGE_MODELS:
+        return 0
+    weight = SAMPLING_WEIGHTS.get(model, 1.0)
+    if model in SAMPLING_BOOST_MODELS:
+        weight *= 5
+    return weight
+
+
+def get_battle_pair():
+    if len(models) == 1:
+        return models[0], models[0]
+
+    model_weights = []
+    for model in models:
+        weight = get_sample_weight(model)
+        model_weights.append(weight)
+    total_weight = np.sum(model_weights)
+    model_weights = model_weights / total_weight
+    chosen_idx = np.random.choice(len(models), p=model_weights)
+    chosen_model = models[chosen_idx]
+    # for p, w in zip(models, model_weights):
+    #     print(p, w)
+
+    rival_models = []
+    rival_weights = []
+    for model in models:
+        if model == chosen_model:
+            continue
+        weight = get_sample_weight(model)
+        if (
+            weight != 0
+            and chosen_model in BATTLE_TARGETS
+            and model in BATTLE_TARGETS[chosen_model]
+        ):
+            # boost to 50% chance
+            weight = total_weight / len(BATTLE_TARGETS[chosen_model])
+        rival_models.append(model)
+        rival_weights.append(weight)
+    # for p, w in zip(rival_models, rival_weights):
+    #     print(p, w)
+    rival_weights = rival_weights / np.sum(rival_weights)
+    rival_idx = np.random.choice(len(rival_models), p=rival_weights)
+    rival_model = rival_models[rival_idx]
+
+    swap = np.random.randint(2)
+    if swap == 0:
+        return chosen_model, rival_model
+    else:
+        return rival_model, chosen_model
+
+
+def add_text(
+    state0, state1, model_selector0, model_selector1, text, request: gr.Request
+):
+    ip = get_ip(request)
+    logger.info(f"add_text (anony). ip: {ip}. len: {len(text)}")
+    states = [state0, state1]
+    model_selectors = [model_selector0, model_selector1]
+
+    # Init states if necessary
+    if states[0] is None:
+        assert states[1] is None
+
+        model_left, model_right = get_battle_pair()
+        states = [
+            State(model_left),
+            State(model_right),
+        ]
+
+    if len(text) <= 0:
+        for i in range(num_sides):
+            states[i].skip_next = True
+        return (
+            states
+            + [x.to_gradio_chatbot() for x in states]
+            + [""]
+            + [
+                no_change_btn,
+            ]
+            * 6
+            + [""]
+        )
+
+    model_list = [states[i].model_name for i in range(num_sides)]
+    flagged = moderation_filter(text, model_list)
+    if flagged:
+        logger.info(f"violate moderation (anony). ip: {ip}. text: {text}")
+        # overwrite the original text
+        text = MODERATION_MSG
+
+    conv = states[0].conv
+    if (len(conv.messages) - conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
+        logger.info(f"conversation turn limit. ip: {get_ip(request)}. text: {text}")
+        for i in range(num_sides):
+            states[i].skip_next = True
+        return (
+            states
+            + [x.to_gradio_chatbot() for x in states]
+            + [CONVERSATION_LIMIT_MSG]
+            + [
+                no_change_btn,
+            ]
+            * 6
+            + [""]
+        )
+
+    text = text[:INPUT_CHAR_LEN_LIMIT]  # Hard cut-off
+    for i in range(num_sides):
+        states[i].conv.append_message(states[i].conv.roles[0], text)
+        states[i].conv.append_message(states[i].conv.roles[1], None)
+        states[i].skip_next = False
+
+    hint_msg = ""
+    for i in range(num_sides):
+        if "deluxe" in states[i].model_name:
+            hint_msg = SLOW_MODEL_MSG
+    return (
+        states
+        + [x.to_gradio_chatbot() for x in states]
+        + [""]
+        + [
+            disable_btn,
+        ]
+        * 6
+        + [hint_msg]
+    )
+
+
+def bot_response_multi(
+    state0,
+    state1,
+    temperature,
+    top_p,
+    max_new_tokens,
+    request: gr.Request,
+):
+    logger.info(f"bot_response_multi (anony). ip: {get_ip(request)}")
+
+    if state0 is None or state0.skip_next:
+        # This generate call is skipped due to invalid inputs
+        yield (
+            state0,
+            state1,
+            state0.to_gradio_chatbot(),
+            state1.to_gradio_chatbot(),
+        ) + (no_change_btn,) * 6
+        return
+
+    states = [state0, state1]
+    gen = []
+    for i in range(num_sides):
+        gen.append(
+            bot_response(
+                states[i],
+                temperature,
+                top_p,
+                max_new_tokens,
+                request,
+                apply_rate_limit=False,
+            )
+        )
+
+    is_gemini = []
+    for i in range(num_sides):
+        is_gemini.append(states[i].model_name in ["gemini-pro", "gemini-pro-dev-api"])
+    chatbots = [None] * num_sides
+    iters = 0
+    while True:
+        stop = True
+        iters += 1
+        for i in range(num_sides):
+            try:
+                # yield gemini fewer times as its chunk size is larger
+                # otherwise, gemini will stream too fast
+                if not is_gemini[i] or (iters % 30 == 1 or iters < 3):
+                    ret = next(gen[i])
+                    states[i], chatbots[i] = ret[0], ret[1]
+                stop = False
+            except StopIteration:
+                pass
+        yield states + chatbots + [disable_btn] * 6
+        if stop:
+            break
+
+
+def build_side_by_side_ui_anony(models):
+    notice_markdown = """
+# ⚔️  Chatbot Arena: Benchmarking LLMs in the Wild
+| [Blog](https://lmsys.org/blog/2023-05-03-arena/) | [GitHub](https://github.com/lm-sys/FastChat) | [Paper](https://arxiv.org/abs/2306.05685) | [Dataset](https://github.com/lm-sys/FastChat/blob/main/docs/dataset_release.md) | [Twitter](https://twitter.com/lmsysorg) | [Discord](https://discord.gg/HSWAKCrnFx) |
+
+## 📜 Rules
+- Ask any question to two anonymous models (e.g., ChatGPT, Claude, Llama) and vote for the better one!
+- You can continue chatting until you identify a winner.
+- Vote won't be counted if model identity is revealed during conversation.
+
+## 🏆 Arena Elo&nbsp;[Leaderboard](https://huggingface.co/spaces/lmsys/chatbot-arena-leaderboard)
+We collect **200K+** human votes to compute an Elo-based LLM leaderboard.
+Find out who is the 🥇LLM Champion!
+
+## 👇 Chat now!
+"""
+
+    states = [gr.State() for _ in range(num_sides)]
+    model_selectors = [None] * num_sides
+    chatbots = [None] * num_sides
+
+    gr.Markdown(notice_markdown, elem_id="notice_markdown")
+
+    with gr.Group(elem_id="share-region-anony"):
+        with gr.Accordion(
+            f"🔍 Expand to see the descriptions of {len(models)} models", open=False
+        ):
+            model_description_md = get_model_description_md(models)
+            gr.Markdown(model_description_md, elem_id="model_description_markdown")
+        with gr.Row():
+            for i in range(num_sides):
+                label = "Model A" if i == 0 else "Model B"
+                with gr.Column():
+                    chatbots[i] = gr.Chatbot(
+                        label=label,
+                        elem_id="chatbot",
+                        height=550,
+                        show_copy_button=True,
+                    )
+
+        with gr.Row():
+            for i in range(num_sides):
+                with gr.Column():
+                    model_selectors[i] = gr.Markdown(
+                        anony_names[i], elem_id="model_selector_md"
+                    )
+        with gr.Row():
+            slow_warning = gr.Markdown("", elem_id="notice_markdown")
+
+    with gr.Row():
+        leftvote_btn = gr.Button(
+            value="👈  A is better", visible=False, interactive=False
+        )
+        rightvote_btn = gr.Button(
+            value="👉  B is better", visible=False, interactive=False
+        )
+        tie_btn = gr.Button(value="🤝  Tie", visible=False, interactive=False)
+        bothbad_btn = gr.Button(
+            value="👎  Both are bad", visible=False, interactive=False
+        )
+
+    with gr.Row():
+        textbox = gr.Textbox(
+            show_label=False,
+            placeholder="👉 Enter your prompt and press ENTER",
+            elem_id="input_box",
+        )
+        send_btn = gr.Button(value="Send", variant="primary", scale=0)
+
+    with gr.Row() as button_row:
+        clear_btn = gr.Button(value="🎲 New Round", interactive=False)
+        regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+        share_btn = gr.Button(value="📷  Share")
+
+    with gr.Accordion("Parameters", open=False) as parameter_row:
+        temperature = gr.Slider(
+            minimum=0.0,
+            maximum=1.0,
+            value=0.7,
+            step=0.1,
+            interactive=True,
+            label="Temperature",
+        )
+        top_p = gr.Slider(
+            minimum=0.0,
+            maximum=1.0,
+            value=1.0,
+            step=0.1,
+            interactive=True,
+            label="Top P",
+        )
+        max_output_tokens = gr.Slider(
+            minimum=16,
+            maximum=2048,
+            value=1024,
+            step=64,
+            interactive=True,
+            label="Max output tokens",
+        )
+
+    gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
+
+    # Register listeners
+    btn_list = [
+        leftvote_btn,
+        rightvote_btn,
+        tie_btn,
+        bothbad_btn,
+        regenerate_btn,
+        clear_btn,
+    ]
+    leftvote_btn.click(
+        leftvote_last_response,
+        states + model_selectors,
+        model_selectors + [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
+    )
+    rightvote_btn.click(
+        rightvote_last_response,
+        states + model_selectors,
+        model_selectors + [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
+    )
+    tie_btn.click(
+        tievote_last_response,
+        states + model_selectors,
+        model_selectors + [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
+    )
+    bothbad_btn.click(
+        bothbad_vote_last_response,
+        states + model_selectors,
+        model_selectors + [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
+    )
+    regenerate_btn.click(
+        regenerate, states, states + chatbots + [textbox] + btn_list
+    ).then(
+        bot_response_multi,
+        states + [temperature, top_p, max_output_tokens],
+        states + chatbots + btn_list,
+    ).then(
+        flash_buttons, [], btn_list
+    )
+    clear_btn.click(
+        clear_history,
+        None,
+        states + chatbots + model_selectors + [textbox] + btn_list + [slow_warning],
+    )
+
+    share_js = """
+function (a, b, c, d) {
+    const captureElement = document.querySelector('#share-region-anony');
+    html2canvas(captureElement)
+        .then(canvas => {
+            canvas.style.display = 'none'
+            document.body.appendChild(canvas)
+            return canvas
+        })
+        .then(canvas => {
+            const image = canvas.toDataURL('image/png')
+            const a = document.createElement('a')
+            a.setAttribute('download', 'chatbot-arena.png')
+            a.setAttribute('href', image)
+            a.click()
+            canvas.remove()
+        });
+    return [a, b, c, d];
+}
+"""
+    share_btn.click(share_click, states + model_selectors, [], js=share_js)
+
+    textbox.submit(
+        add_text,
+        states + model_selectors + [textbox],
+        states + chatbots + [textbox] + btn_list + [slow_warning],
+    ).then(
+        bot_response_multi,
+        states + [temperature, top_p, max_output_tokens],
+        states + chatbots + btn_list,
+    ).then(
+        flash_buttons,
+        [],
+        btn_list,
+    )
+
+    send_btn.click(
+        add_text,
+        states + model_selectors + [textbox],
+        states + chatbots + [textbox] + btn_list,
+    ).then(
+        bot_response_multi,
+        states + [temperature, top_p, max_output_tokens],
+        states + chatbots + btn_list,
+    ).then(
+        flash_buttons, [], btn_list
+    )
+
+    return states + model_selectors

gradio_block_arena_named.py

@@ -0,0 +1,469 @@
+"""
+Chatbot Arena (side-by-side) tab.
+Users chat with two chosen models.
+"""
+
+import json
+import time
+
+import gradio as gr
+import numpy as np
+
+from fastchat.constants import (
+    MODERATION_MSG,
+    CONVERSATION_LIMIT_MSG,
+    INPUT_CHAR_LEN_LIMIT,
+    CONVERSATION_TURN_LIMIT,
+)
+from fastchat.model.model_adapter import get_conversation_template
+from fastchat.serve.gradio_web_server import (
+    State,
+    bot_response,
+    get_conv_log_filename,
+    no_change_btn,
+    enable_btn,
+    disable_btn,
+    invisible_btn,
+    acknowledgment_md,
+    get_ip,
+    get_model_description_md,
+)
+from fastchat.utils import (
+    build_logger,
+    moderation_filter,
+)
+
+
+logger = build_logger("gradio_web_server_multi", "gradio_web_server_multi.log")
+
+num_sides = 2
+enable_moderation = False
+
+
+def set_global_vars_named(enable_moderation_):
+    global enable_moderation
+    enable_moderation = enable_moderation_
+
+
+def load_demo_side_by_side_named(models, url_params):
+    states = (None,) * num_sides
+
+    model_left = models[0] if len(models) > 0 else ""
+    if len(models) > 1:
+        weights = ([8] * 4 + [4] * 8 + [1] * 32)[: len(models) - 1]
+        weights = weights / np.sum(weights)
+        model_right = np.random.choice(models[1:], p=weights)
+    else:
+        model_right = model_left
+
+    selector_updates = (
+        gr.Dropdown(choices=models, value=model_left, visible=True),
+        gr.Dropdown(choices=models, value=model_right, visible=True),
+    )
+
+    return states + selector_updates
+
+
+def vote_last_response(states, vote_type, model_selectors, request: gr.Request):
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(time.time(), 4),
+            "type": vote_type,
+            "models": [x for x in model_selectors],
+            "states": [x.dict() for x in states],
+            "ip": get_ip(request),
+        }
+        fout.write(json.dumps(data) + "\n")
+
+
+def leftvote_last_response(
+    state0, state1, model_selector0, model_selector1, request: gr.Request
+):
+    logger.info(f"leftvote (named). ip: {get_ip(request)}")
+    vote_last_response(
+        [state0, state1], "leftvote", [model_selector0, model_selector1], request
+    )
+    return ("",) + (disable_btn,) * 4
+
+
+def rightvote_last_response(
+    state0, state1, model_selector0, model_selector1, request: gr.Request
+):
+    logger.info(f"rightvote (named). ip: {get_ip(request)}")
+    vote_last_response(
+        [state0, state1], "rightvote", [model_selector0, model_selector1], request
+    )
+    return ("",) + (disable_btn,) * 4
+
+
+def tievote_last_response(
+    state0, state1, model_selector0, model_selector1, request: gr.Request
+):
+    logger.info(f"tievote (named). ip: {get_ip(request)}")
+    vote_last_response(
+        [state0, state1], "tievote", [model_selector0, model_selector1], request
+    )
+    return ("",) + (disable_btn,) * 4
+
+
+def bothbad_vote_last_response(
+    state0, state1, model_selector0, model_selector1, request: gr.Request
+):
+    logger.info(f"bothbad_vote (named). ip: {get_ip(request)}")
+    vote_last_response(
+        [state0, state1], "bothbad_vote", [model_selector0, model_selector1], request
+    )
+    return ("",) + (disable_btn,) * 4
+
+
+def regenerate(state0, state1, request: gr.Request):
+    logger.info(f"regenerate (named). ip: {get_ip(request)}")
+    states = [state0, state1]
+    for i in range(num_sides):
+        states[i].conv.update_last_message(None)
+    return states + [x.to_gradio_chatbot() for x in states] + [""] + [disable_btn] * 6
+
+
+def clear_history(request: gr.Request):
+    logger.info(f"clear_history (named). ip: {get_ip(request)}")
+    return (
+        [None] * num_sides
+        + [None] * num_sides
+        + [""]
+        + [invisible_btn] * 4
+        + [disable_btn] * 2
+    )
+
+
+def share_click(state0, state1, model_selector0, model_selector1, request: gr.Request):
+    logger.info(f"share (named). ip: {get_ip(request)}")
+    if state0 is not None and state1 is not None:
+        vote_last_response(
+            [state0, state1], "share", [model_selector0, model_selector1], request
+        )
+
+
+def add_text(
+    state0, state1, model_selector0, model_selector1, text, request: gr.Request
+):
+    ip = get_ip(request)
+    logger.info(f"add_text (named). ip: {ip}. len: {len(text)}")
+    states = [state0, state1]
+    model_selectors = [model_selector0, model_selector1]
+
+    # Init states if necessary
+    for i in range(num_sides):
+        if states[i] is None:
+            states[i] = State(model_selectors[i])
+
+    if len(text) <= 0:
+        for i in range(num_sides):
+            states[i].skip_next = True
+        return (
+            states
+            + [x.to_gradio_chatbot() for x in states]
+            + [""]
+            + [
+                no_change_btn,
+            ]
+            * 6
+        )
+
+    model_list = [states[i].model_name for i in range(num_sides)]
+    flagged = moderation_filter(text, model_list)
+    if flagged:
+        logger.info(f"violate moderation (named). ip: {ip}. text: {text}")
+        # overwrite the original text
+        text = MODERATION_MSG
+
+    conv = states[0].conv
+    if (len(conv.messages) - conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
+        logger.info(f"conversation turn limit. ip: {ip}. text: {text}")
+        for i in range(num_sides):
+            states[i].skip_next = True
+        return (
+            states
+            + [x.to_gradio_chatbot() for x in states]
+            + [CONVERSATION_LIMIT_MSG]
+            + [
+                no_change_btn,
+            ]
+            * 6
+        )
+
+    text = text[:INPUT_CHAR_LEN_LIMIT]  # Hard cut-off
+    for i in range(num_sides):
+        states[i].conv.append_message(states[i].conv.roles[0], text)
+        states[i].conv.append_message(states[i].conv.roles[1], None)
+        states[i].skip_next = False
+
+    return (
+        states
+        + [x.to_gradio_chatbot() for x in states]
+        + [""]
+        + [
+            disable_btn,
+        ]
+        * 6
+    )
+
+
+def bot_response_multi(
+    state0,
+    state1,
+    temperature,
+    top_p,
+    max_new_tokens,
+    request: gr.Request,
+):
+    logger.info(f"bot_response_multi (named). ip: {get_ip(request)}")
+
+    if state0.skip_next:
+        # This generate call is skipped due to invalid inputs
+        yield (
+            state0,
+            state1,
+            state0.to_gradio_chatbot(),
+            state1.to_gradio_chatbot(),
+        ) + (no_change_btn,) * 6
+        return
+
+    states = [state0, state1]
+    gen = []
+    for i in range(num_sides):
+        gen.append(
+            bot_response(
+                states[i],
+                temperature,
+                top_p,
+                max_new_tokens,
+                request,
+            )
+        )
+
+    is_gemini = []
+    for i in range(num_sides):
+        is_gemini.append(states[i].model_name in ["gemini-pro", "gemini-pro-dev-api"])
+
+    chatbots = [None] * num_sides
+    iters = 0
+    while True:
+        stop = True
+        iters += 1
+        for i in range(num_sides):
+            try:
+                # yield gemini fewer times as its chunk size is larger
+                # otherwise, gemini will stream too fast
+                if not is_gemini[i] or (iters % 30 == 1 or iters < 3):
+                    ret = next(gen[i])
+                    states[i], chatbots[i] = ret[0], ret[1]
+                stop = False
+            except StopIteration:
+                pass
+        yield states + chatbots + [disable_btn] * 6
+        if stop:
+            break
+
+
+def flash_buttons():
+    btn_updates = [
+        [disable_btn] * 4 + [enable_btn] * 2,
+        [enable_btn] * 6,
+    ]
+    for i in range(4):
+        yield btn_updates[i % 2]
+        time.sleep(0.3)
+
+
+def build_side_by_side_ui_named(models):
+    notice_markdown = """
+# ⚔️  Chatbot Arena: Benchmarking LLMs in the Wild
+| [Blog](https://lmsys.org/blog/2023-05-03-arena/) | [GitHub](https://github.com/lm-sys/FastChat) | [Paper](https://arxiv.org/abs/2306.05685) | [Dataset](https://github.com/lm-sys/FastChat/blob/main/docs/dataset_release.md) | [Twitter](https://twitter.com/lmsysorg) | [Discord](https://discord.gg/HSWAKCrnFx) |
+
+## 📜 Rules
+- Chat with any two models side-by-side and vote!
+- You can continue chatting for multiple rounds.
+- Click "Clear history" to start a new round.
+
+## 🤖 Choose two models to compare
+"""
+
+    states = [gr.State() for _ in range(num_sides)]
+    model_selectors = [None] * num_sides
+    chatbots = [None] * num_sides
+
+    notice = gr.Markdown(notice_markdown, elem_id="notice_markdown")
+
+    with gr.Group(elem_id="share-region-named"):
+        with gr.Row():
+            for i in range(num_sides):
+                with gr.Column():
+                    model_selectors[i] = gr.Dropdown(
+                        choices=models,
+                        value=models[i] if len(models) > i else "",
+                        interactive=True,
+                        show_label=False,
+                        container=False,
+                    )
+        with gr.Row():
+            with gr.Accordion(
+                f"🔍 Expand to see the descriptions of {len(models)} models", open=False
+            ):
+                model_description_md = get_model_description_md(models)
+                gr.Markdown(model_description_md, elem_id="model_description_markdown")
+
+        with gr.Row():
+            for i in range(num_sides):
+                label = "Model A" if i == 0 else "Model B"
+                with gr.Column():
+                    chatbots[i] = gr.Chatbot(
+                        label=label,
+                        elem_id=f"chatbot",
+                        height=550,
+                        show_copy_button=True,
+                    )
+
+    with gr.Row():
+        leftvote_btn = gr.Button(
+            value="👈  A is better", visible=False, interactive=False
+        )
+        rightvote_btn = gr.Button(
+            value="👉  B is better", visible=False, interactive=False
+        )
+        tie_btn = gr.Button(value="🤝  Tie", visible=False, interactive=False)
+        bothbad_btn = gr.Button(
+            value="👎  Both are bad", visible=False, interactive=False
+        )
+
+    with gr.Row():
+        textbox = gr.Textbox(
+            show_label=False,
+            placeholder="👉 Enter your prompt and press ENTER",
+            elem_id="input_box",
+        )
+        send_btn = gr.Button(value="Send", variant="primary", scale=0)
+
+    with gr.Row() as button_row:
+        clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)
+        regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+        share_btn = gr.Button(value="📷  Share")
+
+    with gr.Accordion("Parameters", open=False) as parameter_row:
+        temperature = gr.Slider(
+            minimum=0.0,
+            maximum=1.0,
+            value=0.7,
+            step=0.1,
+            interactive=True,
+            label="Temperature",
+        )
+        top_p = gr.Slider(
+            minimum=0.0,
+            maximum=1.0,
+            value=1.0,
+            step=0.1,
+            interactive=True,
+            label="Top P",
+        )
+        max_output_tokens = gr.Slider(
+            minimum=16,
+            maximum=2048,
+            value=1024,
+            step=64,
+            interactive=True,
+            label="Max output tokens",
+        )
+
+    gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
+
+    # Register listeners
+    btn_list = [
+        leftvote_btn,
+        rightvote_btn,
+        tie_btn,
+        bothbad_btn,
+        regenerate_btn,
+        clear_btn,
+    ]
+    leftvote_btn.click(
+        leftvote_last_response,
+        states + model_selectors,
+        [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
+    )
+    rightvote_btn.click(
+        rightvote_last_response,
+        states + model_selectors,
+        [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
+    )
+    tie_btn.click(
+        tievote_last_response,
+        states + model_selectors,
+        [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
+    )
+    bothbad_btn.click(
+        bothbad_vote_last_response,
+        states + model_selectors,
+        [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
+    )
+    regenerate_btn.click(
+        regenerate, states, states + chatbots + [textbox] + btn_list
+    ).then(
+        bot_response_multi,
+        states + [temperature, top_p, max_output_tokens],
+        states + chatbots + btn_list,
+    ).then(
+        flash_buttons, [], btn_list
+    )
+    clear_btn.click(clear_history, None, states + chatbots + [textbox] + btn_list)
+
+    share_js = """
+function (a, b, c, d) {
+    const captureElement = document.querySelector('#share-region-named');
+    html2canvas(captureElement)
+        .then(canvas => {
+            canvas.style.display = 'none'
+            document.body.appendChild(canvas)
+            return canvas
+        })
+        .then(canvas => {
+            const image = canvas.toDataURL('image/png')
+            const a = document.createElement('a')
+            a.setAttribute('download', 'chatbot-arena.png')
+            a.setAttribute('href', image)
+            a.click()
+            canvas.remove()
+        });
+    return [a, b, c, d];
+}
+"""
+    share_btn.click(share_click, states + model_selectors, [], js=share_js)
+
+    for i in range(num_sides):
+        model_selectors[i].change(
+            clear_history, None, states + chatbots + [textbox] + btn_list
+        )
+
+    textbox.submit(
+        add_text,
+        states + model_selectors + [textbox],
+        states + chatbots + [textbox] + btn_list,
+    ).then(
+        bot_response_multi,
+        states + [temperature, top_p, max_output_tokens],
+        states + chatbots + btn_list,
+    ).then(
+        flash_buttons, [], btn_list
+    )
+    send_btn.click(
+        add_text,
+        states + model_selectors + [textbox],
+        states + chatbots + [textbox] + btn_list,
+    ).then(
+        bot_response_multi,
+        states + [temperature, top_p, max_output_tokens],
+        states + chatbots + btn_list,
+    ).then(
+        flash_buttons, [], btn_list
+    )
+
+    return states + model_selectors

gradio_block_arena_vision.py

@@ -0,0 +1,187 @@
+"""
+The gradio demo server for chatting with a large multimodal model.
+
+Usage:
+python3 -m fastchat.serve.controller
+python3 -m fastchat.serve.sglang_worker --model-path liuhaotian/llava-v1.5-7b --tokenizer-path llava-hf/llava-1.5-7b-hf
+python3 -m fastchat.serve.gradio_web_server_multi --share --multimodal
+"""
+
+import os
+
+import gradio as gr
+
+from fastchat.serve.gradio_web_server import (
+    upvote_last_response,
+    downvote_last_response,
+    flag_last_response,
+    get_model_description_md,
+    acknowledgment_md,
+    bot_response,
+    add_text,
+    clear_history,
+    regenerate,
+)
+from fastchat.utils import (
+    build_logger,
+)
+
+logger = build_logger("gradio_web_server_multi", "gradio_web_server_multi.log")
+
+
+def build_single_vision_language_model_ui(models, add_promotion_links=False):
+    promotion = (
+        """
+| [GitHub](https://github.com/lm-sys/FastChat) | [Dataset](https://github.com/lm-sys/FastChat/blob/main/docs/dataset_release.md) | [Twitter](https://twitter.com/lmsysorg) | [Discord](https://discord.gg/HSWAKCrnFx) |
+"""
+        if add_promotion_links
+        else ""
+    )
+
+    notice_markdown = f"""
+# 🏔️ Chat with Open Large Vision-Language Models
+{promotion}
+"""
+
+    state = gr.State()
+    gr.Markdown(notice_markdown, elem_id="notice_markdown")
+
+    with gr.Group():
+        with gr.Row(elem_id="model_selector_row"):
+            model_selector = gr.Dropdown(
+                choices=models,
+                value=models[0] if len(models) > 0 else "",
+                interactive=True,
+                show_label=False,
+                container=False,
+            )
+
+        with gr.Accordion(
+            f"🔍 Expand to see the descriptions of {len(models)} models", open=False
+        ):
+            model_description_md = get_model_description_md(models)
+            gr.Markdown(model_description_md, elem_id="model_description_markdown")
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            textbox = gr.Textbox(
+                show_label=False,
+                placeholder="👉 Enter your prompt and press ENTER",
+                container=False,
+                render=False,
+                elem_id="input_box",
+            )
+            imagebox = gr.Image(type="pil")
+
+            cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+            with gr.Accordion("Parameters", open=False) as parameter_row:
+                temperature = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.2,
+                    step=0.1,
+                    interactive=True,
+                    label="Temperature",
+                )
+                top_p = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.7,
+                    step=0.1,
+                    interactive=True,
+                    label="Top P",
+                )
+                max_output_tokens = gr.Slider(
+                    minimum=0,
+                    maximum=1024,
+                    value=512,
+                    step=64,
+                    interactive=True,
+                    label="Max output tokens",
+                )
+
+            gr.Examples(
+                examples=[
+                    [
+                        f"{cur_dir}/example_images/city.jpeg",
+                        "What is unusual about this image?",
+                    ],
+                    [
+                        f"{cur_dir}/example_images/fridge.jpeg",
+                        "What is in this fridge?",
+                    ],
+                ],
+                inputs=[imagebox, textbox],
+            )
+
+        with gr.Column(scale=8):
+            chatbot = gr.Chatbot(
+                elem_id="chatbot", label="Scroll down and start chatting", height=550
+            )
+
+            with gr.Row():
+                with gr.Column(scale=8):
+                    textbox.render()
+                with gr.Column(scale=1, min_width=50):
+                    send_btn = gr.Button(value="Send", variant="primary")
+            with gr.Row(elem_id="buttons"):
+                upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
+                downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
+                flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
+                regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+                clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
+
+    if add_promotion_links:
+        gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
+
+    # Register listeners
+    btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
+    upvote_btn.click(
+        upvote_last_response,
+        [state, model_selector],
+        [textbox, upvote_btn, downvote_btn, flag_btn],
+    )
+    downvote_btn.click(
+        downvote_last_response,
+        [state, model_selector],
+        [textbox, upvote_btn, downvote_btn, flag_btn],
+    )
+    flag_btn.click(
+        flag_last_response,
+        [state, model_selector],
+        [textbox, upvote_btn, downvote_btn, flag_btn],
+    )
+    regenerate_btn.click(
+        regenerate, state, [state, chatbot, textbox, imagebox] + btn_list
+    ).then(
+        bot_response,
+        [state, temperature, top_p, max_output_tokens],
+        [state, chatbot] + btn_list,
+    )
+    clear_btn.click(clear_history, None, [state, chatbot, textbox, imagebox] + btn_list)
+
+    model_selector.change(
+        clear_history, None, [state, chatbot, textbox, imagebox] + btn_list
+    )
+
+    textbox.submit(
+        add_text,
+        [state, model_selector, textbox, imagebox],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    ).then(
+        bot_response,
+        [state, temperature, top_p, max_output_tokens],
+        [state, chatbot] + btn_list,
+    )
+    send_btn.click(
+        add_text,
+        [state, model_selector, textbox, imagebox],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    ).then(
+        bot_response,
+        [state, temperature, top_p, max_output_tokens],
+        [state, chatbot] + btn_list,
+    )
+
+    return [state, model_selector]

gradio_web_server.py

@@ -0,0 +1,887 @@
+"""
+The gradio demo server for chatting with a single model.
+"""
+
+import argparse
+from collections import defaultdict
+import datetime
+import hashlib
+import json
+import os
+import random
+import time
+import uuid
+
+import gradio as gr
+import requests
+
+from fastchat.constants import (
+    LOGDIR,
+    WORKER_API_TIMEOUT,
+    ErrorCode,
+    MODERATION_MSG,
+    CONVERSATION_LIMIT_MSG,
+    RATE_LIMIT_MSG,
+    SERVER_ERROR_MSG,
+    INPUT_CHAR_LEN_LIMIT,
+    CONVERSATION_TURN_LIMIT,
+    SESSION_EXPIRATION_TIME,
+)
+from fastchat.model.model_adapter import (
+    get_conversation_template,
+)
+from fastchat.model.model_registry import get_model_info, model_info
+from fastchat.serve.api_provider import get_api_provider_stream_iter
+from fastchat.utils import (
+    build_logger,
+    get_window_url_params_js,
+    get_window_url_params_with_tos_js,
+    moderation_filter,
+    parse_gradio_auth_creds,
+    load_image,
+)
+
+
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
+
+headers = {"User-Agent": "FastChat Client"}
+
+no_change_btn = gr.Button()
+enable_btn = gr.Button(interactive=True, visible=True)
+disable_btn = gr.Button(interactive=False)
+invisible_btn = gr.Button(interactive=False, visible=False)
+
+controller_url = None
+enable_moderation = False
+
+acknowledgment_md = """
+### Terms of Service
+
+Users are required to agree to the following terms before using the service:
+
+The service is a research preview. It only provides limited safety measures and may generate offensive content.
+It must not be used for any illegal, harmful, violent, racist, or sexual purposes.
+The service collects user dialogue data and reserves the right to distribute it under a Creative Commons Attribution (CC-BY) or a similar license.
+Additionally, Bard is offered on LMSys for research purposes only. To access the Bard product, please visit its [website](http://bard.google.com).
+
+### Acknowledgment
+We thank [Kaggle](https://www.kaggle.com/), [MBZUAI](https://mbzuai.ac.ae/), [a16z](https://www.a16z.com/), [Together AI](https://www.together.ai/), [Anyscale](https://www.anyscale.com/), [HuggingFace](https://huggingface.co/) for their generous [sponsorship](https://lmsys.org/donations/).
+
+<div class="sponsor-image-about">
+    <img src="https://storage.googleapis.com/public-arena-asset/kaggle.png" alt="Kaggle">
+    <img src="https://storage.googleapis.com/public-arena-asset/mbzuai.jpeg" alt="MBZUAI">
+    <img src="https://storage.googleapis.com/public-arena-asset/a16z.jpeg" alt="a16z">
+    <img src="https://storage.googleapis.com/public-arena-asset/together.png" alt="Together AI">
+    <img src="https://storage.googleapis.com/public-arena-asset/anyscale.png" alt="AnyScale">
+    <img src="https://storage.googleapis.com/public-arena-asset/huggingface.png" alt="HuggingFace">
+</div>
+"""
+
+# JSON file format of API-based models:
+# {
+#   "gpt-3.5-turbo-0613": {
+#     "model_name": "gpt-3.5-turbo-0613",
+#     "api_type": "openai",
+#     "api_base": "https://api.openai.com/v1",
+#     "api_key": "sk-******",
+#     "anony_only": false
+#   }
+# }
+# "api_type" can be one of the following: openai, anthropic, gemini, mistral.
+# "anony_only" means whether to show this model in anonymous mode only.
+api_endpoint_info = {}
+
+
+class State:
+    def __init__(self, model_name):
+        self.conv = get_conversation_template(model_name)
+        self.conv_id = uuid.uuid4().hex
+        self.skip_next = False
+        self.model_name = model_name
+
+    def to_gradio_chatbot(self):
+        return self.conv.to_gradio_chatbot()
+
+    def dict(self):
+        base = self.conv.dict()
+        base.update(
+            {
+                "conv_id": self.conv_id,
+                "model_name": self.model_name,
+            }
+        )
+        return base
+
+
+def set_global_vars(controller_url_, enable_moderation_):
+    global controller_url, enable_moderation
+    controller_url = controller_url_
+    enable_moderation = enable_moderation_
+
+
+def get_conv_log_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
+    return name
+
+
+def get_model_list(controller_url, register_api_endpoint_file, multimodal):
+    global api_endpoint_info
+
+    # Add models from the controller
+    if controller_url:
+        ret = requests.post(controller_url + "/refresh_all_workers")
+        assert ret.status_code == 200
+
+        if multimodal:
+            ret = requests.post(controller_url + "/list_multimodal_models")
+            models = ret.json()["models"]
+        else:
+            ret = requests.post(controller_url + "/list_language_models")
+            models = ret.json()["models"]
+    else:
+        models = []
+
+    # Add models from the API providers
+    if register_api_endpoint_file:
+        api_endpoint_info = json.load(open(register_api_endpoint_file))
+        for mdl, mdl_dict in api_endpoint_info.items():
+            mdl_multimodal = mdl_dict.get("multimodal", False)
+            if multimodal and mdl_multimodal:
+                models += [mdl]
+            elif not multimodal and not mdl_multimodal:
+                models += [mdl]
+
+    # Remove anonymous models
+    models = list(set(models))
+    visible_models = models.copy()
+    for mdl in visible_models:
+        if mdl not in api_endpoint_info:
+            continue
+        mdl_dict = api_endpoint_info[mdl]
+        if mdl_dict["anony_only"]:
+            visible_models.remove(mdl)
+
+    # Sort models and add descriptions
+    priority = {k: f"___{i:03d}" for i, k in enumerate(model_info)}
+    models.sort(key=lambda x: priority.get(x, x))
+    visible_models.sort(key=lambda x: priority.get(x, x))
+    logger.info(f"All models: {models}")
+    logger.info(f"Visible models: {visible_models}")
+    return visible_models, models
+
+
+def load_demo_single(models, url_params):
+    selected_model = models[0] if len(models) > 0 else ""
+    if "model" in url_params:
+        model = url_params["model"]
+        if model in models:
+            selected_model = model
+
+    dropdown_update = gr.Dropdown(choices=models, value=selected_model, visible=True)
+    state = None
+    return state, dropdown_update
+
+
+def load_demo(url_params, request: gr.Request):
+    global models
+
+    ip = get_ip(request)
+    logger.info(f"load_demo. ip: {ip}. params: {url_params}")
+
+    if args.model_list_mode == "reload":
+        models, all_models = get_model_list(
+            controller_url, args.register_api_endpoint_file, False
+        )
+
+    return load_demo_single(models, url_params)
+
+
+def vote_last_response(state, vote_type, model_selector, request: gr.Request):
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(time.time(), 4),
+            "type": vote_type,
+            "model": model_selector,
+            "state": state.dict(),
+            "ip": get_ip(request),
+        }
+        fout.write(json.dumps(data) + "\n")
+
+
+def upvote_last_response(state, model_selector, request: gr.Request):
+    ip = get_ip(request)
+    logger.info(f"upvote. ip: {ip}")
+    vote_last_response(state, "upvote", model_selector, request)
+    return ("",) + (disable_btn,) * 3
+
+
+def downvote_last_response(state, model_selector, request: gr.Request):
+    ip = get_ip(request)
+    logger.info(f"downvote. ip: {ip}")
+    vote_last_response(state, "downvote", model_selector, request)
+    return ("",) + (disable_btn,) * 3
+
+
+def flag_last_response(state, model_selector, request: gr.Request):
+    ip = get_ip(request)
+    logger.info(f"flag. ip: {ip}")
+    vote_last_response(state, "flag", model_selector, request)
+    return ("",) + (disable_btn,) * 3
+
+
+def regenerate(state, request: gr.Request):
+    ip = get_ip(request)
+    logger.info(f"regenerate. ip: {ip}")
+    state.conv.update_last_message(None)
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+
+def clear_history(request: gr.Request):
+    ip = get_ip(request)
+    logger.info(f"clear_history. ip: {ip}")
+    state = None
+    return (state, [], "", None) + (disable_btn,) * 5
+
+
+def get_ip(request: gr.Request):
+    if "cf-connecting-ip" in request.headers:
+        ip = request.headers["cf-connecting-ip"]
+    else:
+        ip = request.client.host
+    return ip
+
+
+def _prepare_text_with_image(state, text, image):
+    if image is not None:
+        if len(state.conv.get_images()) > 0:
+            # reset convo with new image
+            state.conv = get_conversation_template(state.model_name)
+
+        image = state.conv.convert_image_to_base64(
+            image
+        )  # PIL type is not JSON serializable
+
+        text = text, [image]
+
+    return text
+
+
+def add_text(state, model_selector, text, image, request: gr.Request):
+    ip = get_ip(request)
+    logger.info(f"add_text. ip: {ip}. len: {len(text)}")
+
+    if state is None:
+        state = State(model_selector)
+
+    if len(text) <= 0:
+        state.skip_next = True
+        return (state, state.to_gradio_chatbot(), "") + (no_change_btn,) * 5
+
+    flagged = moderation_filter(text, [state.model_name])
+    if flagged:
+        logger.info(f"violate moderation. ip: {ip}. text: {text}")
+        # overwrite the original text
+        text = MODERATION_MSG
+
+    if (len(state.conv.messages) - state.conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
+        logger.info(f"conversation turn limit. ip: {ip}. text: {text}")
+        state.skip_next = True
+        return (state, state.to_gradio_chatbot(), CONVERSATION_LIMIT_MSG) + (
+            no_change_btn,
+        ) * 5
+
+    text = text[:INPUT_CHAR_LEN_LIMIT]  # Hard cut-off
+    text = _prepare_text_with_image(state, text, image)
+    state.conv.append_message(state.conv.roles[0], text)
+    state.conv.append_message(state.conv.roles[1], None)
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+
+def model_worker_stream_iter(
+    conv,
+    model_name,
+    worker_addr,
+    prompt,
+    temperature,
+    repetition_penalty,
+    top_p,
+    max_new_tokens,
+    images,
+):
+    # Make requests
+    gen_params = {
+        "model": model_name,
+        "prompt": prompt,
+        "temperature": temperature,
+        "repetition_penalty": repetition_penalty,
+        "top_p": top_p,
+        "max_new_tokens": max_new_tokens,
+        "stop": conv.stop_str,
+        "stop_token_ids": conv.stop_token_ids,
+        "echo": False,
+    }
+
+    logger.info(f"==== request ====\n{gen_params}")
+
+    if len(images) > 0:
+        gen_params["images"] = images
+
+    # Stream output
+    response = requests.post(
+        worker_addr + "/worker_generate_stream",
+        headers=headers,
+        json=gen_params,
+        stream=True,
+        timeout=WORKER_API_TIMEOUT,
+    )
+    for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+        if chunk:
+            data = json.loads(chunk.decode())
+            yield data
+
+
+def is_limit_reached(model_name, ip):
+    monitor_url = "http://localhost:9090"
+    try:
+        ret = requests.get(
+            f"{monitor_url}/is_limit_reached?model={model_name}&user_id={ip}", timeout=1
+        )
+        obj = ret.json()
+        return obj
+    except Exception as e:
+        logger.info(f"monitor error: {e}")
+        return None
+
+
+def bot_response(
+    state,
+    temperature,
+    top_p,
+    max_new_tokens,
+    request: gr.Request,
+    apply_rate_limit=True,
+):
+    ip = get_ip(request)
+    logger.info(f"bot_response. ip: {ip}")
+    start_tstamp = time.time()
+    temperature = float(temperature)
+    top_p = float(top_p)
+    max_new_tokens = int(max_new_tokens)
+
+    if state.skip_next:
+        # This generate call is skipped due to invalid inputs
+        state.skip_next = False
+        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
+        return
+
+    if apply_rate_limit:
+        ret = is_limit_reached(state.model_name, ip)
+        if ret is not None and ret["is_limit_reached"]:
+            error_msg = RATE_LIMIT_MSG + "\n\n" + ret["reason"]
+            logger.info(f"rate limit reached. ip: {ip}. error_msg: {ret['reason']}")
+            state.conv.update_last_message(error_msg)
+            yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
+            return
+
+    conv, model_name = state.conv, state.model_name
+    model_api_dict = (
+        api_endpoint_info[model_name] if model_name in api_endpoint_info else None
+    )
+    images = conv.get_images()
+
+    if model_api_dict is None:
+        # Query worker address
+        ret = requests.post(
+            controller_url + "/get_worker_address", json={"model": model_name}
+        )
+        worker_addr = ret.json()["address"]
+        logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
+
+        # No available worker
+        if worker_addr == "":
+            conv.update_last_message(SERVER_ERROR_MSG)
+            yield (
+                state,
+                state.to_gradio_chatbot(),
+                disable_btn,
+                disable_btn,
+                disable_btn,
+                enable_btn,
+                enable_btn,
+            )
+            return
+
+        # Construct prompt.
+        # We need to call it here, so it will not be affected by "▌".
+        prompt = conv.get_prompt()
+
+        # Set repetition_penalty
+        if "t5" in model_name:
+            repetition_penalty = 1.2
+        else:
+            repetition_penalty = 1.0
+
+        stream_iter = model_worker_stream_iter(
+            conv,
+            model_name,
+            worker_addr,
+            prompt,
+            temperature,
+            repetition_penalty,
+            top_p,
+            max_new_tokens,
+            images,
+        )
+    else:
+        stream_iter = get_api_provider_stream_iter(
+            conv,
+            model_name,
+            model_api_dict,
+            temperature,
+            top_p,
+            max_new_tokens,
+        )
+
+    conv.update_last_message("▌")
+    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+
+    try:
+        for i, data in enumerate(stream_iter):
+            if data["error_code"] == 0:
+                output = data["text"].strip()
+                conv.update_last_message(output + "▌")
+                yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+            else:
+                output = data["text"] + f"\n\n(error_code: {data['error_code']})"
+                conv.update_last_message(output)
+                yield (state, state.to_gradio_chatbot()) + (
+                    disable_btn,
+                    disable_btn,
+                    disable_btn,
+                    enable_btn,
+                    enable_btn,
+                )
+                return
+        output = data["text"].strip()
+        conv.update_last_message(output)
+        yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5
+    except requests.exceptions.RequestException as e:
+        conv.update_last_message(
+            f"{SERVER_ERROR_MSG}\n\n"
+            f"(error_code: {ErrorCode.GRADIO_REQUEST_ERROR}, {e})"
+        )
+        yield (state, state.to_gradio_chatbot()) + (
+            disable_btn,
+            disable_btn,
+            disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+    except Exception as e:
+        conv.update_last_message(
+            f"{SERVER_ERROR_MSG}\n\n"
+            f"(error_code: {ErrorCode.GRADIO_STREAM_UNKNOWN_ERROR}, {e})"
+        )
+        yield (state, state.to_gradio_chatbot()) + (
+            disable_btn,
+            disable_btn,
+            disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    finish_tstamp = time.time()
+    logger.info(f"{output}")
+
+    # We load the image because gradio accepts base64 but that increases file size by ~1.33x
+    loaded_images = [load_image(image) for image in images]
+    images_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in loaded_images]
+    for image, hash_str in zip(loaded_images, images_hash):
+        t = datetime.datetime.now()
+        filename = os.path.join(
+            LOGDIR,
+            "serve_images",
+            f"{hash_str}.jpg",
+        )
+        if not os.path.isfile(filename):
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            image.save(filename)
+
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(finish_tstamp, 4),
+            "type": "chat",
+            "model": model_name,
+            "gen_params": {
+                "temperature": temperature,
+                "top_p": top_p,
+                "max_new_tokens": max_new_tokens,
+            },
+            "start": round(start_tstamp, 4),
+            "finish": round(finish_tstamp, 4),
+            "state": state.dict(),
+            "ip": get_ip(request),
+            "images": images_hash,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+
+block_css = """
+#notice_markdown .prose {
+    font-size: 120% !important;
+}
+#notice_markdown th {
+    display: none;
+}
+#notice_markdown td {
+    padding-top: 6px;
+    padding-bottom: 6px;
+}
+#model_description_markdown {
+    font-size: 120% !important;
+}
+#leaderboard_markdown .prose {
+    font-size: 120% !important;
+}
+#leaderboard_markdown td {
+    padding-top: 6px;
+    padding-bottom: 6px;
+}
+#leaderboard_dataframe td {
+    line-height: 0.1em;
+}
+#about_markdown .prose {
+    font-size: 120% !important;
+}
+#ack_markdown .prose {
+    font-size: 120% !important;
+}
+footer {
+    display:none !important;
+}
+.sponsor-image-about img {
+    margin: 0 20px;
+    margin-top: 20px;
+    height: 40px;
+    max-height: 100%;
+    width: auto;
+    float: left;
+}
+"""
+
+
+def get_model_description_md(models):
+    model_description_md = """
+| | | |
+| ---- | ---- | ---- |
+"""
+    ct = 0
+    visited = set()
+    for i, name in enumerate(models):
+        minfo = get_model_info(name)
+        if minfo.simple_name in visited:
+            continue
+        visited.add(minfo.simple_name)
+        one_model_md = f"[{minfo.simple_name}]({minfo.link}): {minfo.description}"
+
+        if ct % 3 == 0:
+            model_description_md += "|"
+        model_description_md += f" {one_model_md} |"
+        if ct % 3 == 2:
+            model_description_md += "\n"
+        ct += 1
+    return model_description_md
+
+
+def build_about():
+    about_markdown = """
+# About Us
+Chatbot Arena is an open-source research project developed by members from [LMSYS](https://lmsys.org/about/) and UC Berkeley [SkyLab](https://sky.cs.berkeley.edu/).  Our mission is to build an open crowdsourced platform to collect human feedback and evaluate LLMs under real-world scenarios. We open-source our [FastChat](https://github.com/lm-sys/FastChat) project at GitHub and release chat and human feedback datasets [here](https://github.com/lm-sys/FastChat/blob/main/docs/dataset_release.md). We invite everyone to join us in this journey!
+
+## Read More
+- Chatbot Arena [launch post](https://lmsys.org/blog/2023-05-03-arena/), [data release](https://lmsys.org/blog/2023-07-20-dataset/)
+- LMSYS-Chat-1M [report](https://arxiv.org/abs/2309.11998)
+
+## Core Members
+[Lianmin Zheng](https://lmzheng.net/), [Wei-Lin Chiang](https://infwinston.github.io/), [Ying Sheng](https://sites.google.com/view/yingsheng/home), [Siyuan Zhuang](https://scholar.google.com/citations?user=KSZmI5EAAAAJ)
+
+## Advisors
+[Ion Stoica](http://people.eecs.berkeley.edu/~istoica/), [Joseph E. Gonzalez](https://people.eecs.berkeley.edu/~jegonzal/), [Hao Zhang](https://cseweb.ucsd.edu/~haozhang/)
+
+## Contact Us
+- Follow our [Twitter](https://twitter.com/lmsysorg), [Discord](https://discord.gg/HSWAKCrnFx) or email us at [email protected]
+- File issues on [GitHub](https://github.com/lm-sys/FastChat)
+- Download our datasets and models on [HuggingFace](https://huggingface.co/lmsys)
+
+## Acknowledgment
+We thank [SkyPilot](https://github.com/skypilot-org/skypilot) and [Gradio](https://github.com/gradio-app/gradio) team for their system support.
+We also thank [Kaggle](https://www.kaggle.com/), [MBZUAI](https://mbzuai.ac.ae/), [a16z](https://www.a16z.com/), [Together AI](https://www.together.ai/), [Anyscale](https://www.anyscale.com/), [HuggingFace](https://huggingface.co/) for their generous sponsorship. Learn more about partnership [here](https://lmsys.org/donations/).
+
+<div class="sponsor-image-about">
+    <img src="https://storage.googleapis.com/public-arena-asset/kaggle.png" alt="Kaggle">
+    <img src="https://storage.googleapis.com/public-arena-asset/mbzuai.jpeg" alt="MBZUAI">
+    <img src="https://storage.googleapis.com/public-arena-asset/a16z.jpeg" alt="a16z">
+    <img src="https://storage.googleapis.com/public-arena-asset/together.png" alt="Together AI">
+    <img src="https://storage.googleapis.com/public-arena-asset/anyscale.png" alt="AnyScale">
+    <img src="https://storage.googleapis.com/public-arena-asset/huggingface.png" alt="HuggingFace">
+</div>
+"""
+    gr.Markdown(about_markdown, elem_id="about_markdown")
+
+
+def build_single_model_ui(models, add_promotion_links=False):
+    promotion = (
+        """
+- | [GitHub](https://github.com/lm-sys/FastChat) | [Dataset](https://github.com/lm-sys/FastChat/blob/main/docs/dataset_release.md) | [Twitter](https://twitter.com/lmsysorg) | [Discord](https://discord.gg/HSWAKCrnFx) |
+- Introducing Llama 2: The Next Generation Open Source Large Language Model. [[Website]](https://ai.meta.com/llama/)
+- Vicuna: An Open-Source Chatbot Impressing GPT-4 with 90% ChatGPT Quality. [[Blog]](https://lmsys.org/blog/2023-03-30-vicuna/)
+
+## 🤖 Choose any model to chat
+"""
+        if add_promotion_links
+        else ""
+    )
+
+    notice_markdown = f"""
+# 🏔️ Chat with Open Large Language Models
+{promotion}
+"""
+
+    state = gr.State()
+    gr.Markdown(notice_markdown, elem_id="notice_markdown")
+
+    with gr.Group(elem_id="share-region-named"):
+        with gr.Row(elem_id="model_selector_row"):
+            model_selector = gr.Dropdown(
+                choices=models,
+                value=models[0] if len(models) > 0 else "",
+                interactive=True,
+                show_label=False,
+                container=False,
+            )
+        with gr.Row():
+            with gr.Accordion(
+                f"🔍 Expand to see the descriptions of {len(models)} models",
+                open=False,
+            ):
+                model_description_md = get_model_description_md(models)
+                gr.Markdown(model_description_md, elem_id="model_description_markdown")
+
+        chatbot = gr.Chatbot(
+            elem_id="chatbot",
+            label="Scroll down and start chatting",
+            height=550,
+            show_copy_button=True,
+        )
+    with gr.Row():
+        textbox = gr.Textbox(
+            show_label=False,
+            placeholder="👉 Enter your prompt and press ENTER",
+            elem_id="input_box",
+        )
+        send_btn = gr.Button(value="Send", variant="primary", scale=0)
+
+    with gr.Row() as button_row:
+        upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
+        downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
+        flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
+        regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+        clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)
+
+    with gr.Accordion("Parameters", open=False) as parameter_row:
+        temperature = gr.Slider(
+            minimum=0.0,
+            maximum=1.0,
+            value=0.7,
+            step=0.1,
+            interactive=True,
+            label="Temperature",
+        )
+        top_p = gr.Slider(
+            minimum=0.0,
+            maximum=1.0,
+            value=1.0,
+            step=0.1,
+            interactive=True,
+            label="Top P",
+        )
+        max_output_tokens = gr.Slider(
+            minimum=16,
+            maximum=2048,
+            value=1024,
+            step=64,
+            interactive=True,
+            label="Max output tokens",
+        )
+
+    if add_promotion_links:
+        gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
+
+    # Register listeners
+    imagebox = gr.State(None)
+    btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
+    upvote_btn.click(
+        upvote_last_response,
+        [state, model_selector],
+        [textbox, upvote_btn, downvote_btn, flag_btn],
+    )
+    downvote_btn.click(
+        downvote_last_response,
+        [state, model_selector],
+        [textbox, upvote_btn, downvote_btn, flag_btn],
+    )
+    flag_btn.click(
+        flag_last_response,
+        [state, model_selector],
+        [textbox, upvote_btn, downvote_btn, flag_btn],
+    )
+    regenerate_btn.click(
+        regenerate, state, [state, chatbot, textbox, imagebox] + btn_list
+    ).then(
+        bot_response,
+        [state, temperature, top_p, max_output_tokens],
+        [state, chatbot] + btn_list,
+    )
+    clear_btn.click(clear_history, None, [state, chatbot, textbox, imagebox] + btn_list)
+
+    model_selector.change(
+        clear_history, None, [state, chatbot, textbox, imagebox] + btn_list
+    )
+
+    textbox.submit(
+        add_text,
+        [state, model_selector, textbox, imagebox],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    ).then(
+        bot_response,
+        [state, temperature, top_p, max_output_tokens],
+        [state, chatbot] + btn_list,
+    )
+    send_btn.click(
+        add_text,
+        [state, model_selector, textbox, imagebox],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    ).then(
+        bot_response,
+        [state, temperature, top_p, max_output_tokens],
+        [state, chatbot] + btn_list,
+    )
+
+    return [state, model_selector]
+
+
+def build_demo(models):
+    with gr.Blocks(
+        title="Chat with Open Large Language Models",
+        theme=gr.themes.Default(),
+        css=block_css,
+    ) as demo:
+        url_params = gr.JSON(visible=False)
+
+        state, model_selector = build_single_model_ui(models)
+
+        if args.model_list_mode not in ["once", "reload"]:
+            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
+
+        if args.show_terms_of_use:
+            load_js = get_window_url_params_with_tos_js
+        else:
+            load_js = get_window_url_params_js
+
+        demo.load(
+            load_demo,
+            [url_params],
+            [
+                state,
+                model_selector,
+            ],
+            js=load_js,
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int)
+    parser.add_argument(
+        "--share",
+        action="store_true",
+        help="Whether to generate a public, shareable link",
+    )
+    parser.add_argument(
+        "--controller-url",
+        type=str,
+        default="http://localhost:21001",
+        help="The address of the controller",
+    )
+    parser.add_argument(
+        "--concurrency-count",
+        type=int,
+        default=10,
+        help="The concurrency count of the gradio queue",
+    )
+    parser.add_argument(
+        "--model-list-mode",
+        type=str,
+        default="once",
+        choices=["once", "reload"],
+        help="Whether to load the model list once or reload the model list every time",
+    )
+    parser.add_argument(
+        "--moderate",
+        action="store_true",
+        help="Enable content moderation to block unsafe inputs",
+    )
+    parser.add_argument(
+        "--show-terms-of-use",
+        action="store_true",
+        help="Shows term of use before loading the demo",
+    )
+    parser.add_argument(
+        "--register-api-endpoint-file",
+        type=str,
+        help="Register API-based model endpoints from a JSON file",
+    )
+    parser.add_argument(
+        "--gradio-auth-path",
+        type=str,
+        help='Set the gradio authentication file path. The file should contain one or more user:password pairs in this format: "u1:p1,u2:p2,u3:p3"',
+    )
+    parser.add_argument(
+        "--gradio-root-path",
+        type=str,
+        help="Sets the gradio root path, eg /abc/def. Useful when running behind a reverse-proxy or at a custom URL path prefix",
+    )
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    # Set global variables
+    set_global_vars(args.controller_url, args.moderate)
+    models, all_models = get_model_list(
+        args.controller_url, args.register_api_endpoint_file, False
+    )
+
+    # Set authorization credentials
+    auth = None
+    if args.gradio_auth_path is not None:
+        auth = parse_gradio_auth_creds(args.gradio_auth_path)
+
+    # Launch the demo
+    demo = build_demo(models)
+    demo.queue(
+        default_concurrency_limit=args.concurrency_count,
+        status_update_rate=10,
+        api_open=False,
+    ).launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share,
+        max_threads=200,
+        auth=auth,
+        root_path=args.gradio_root_path,
+    )

gradio_web_server_multi.py

@@ -0,0 +1,277 @@
+"""
+The gradio demo server with multiple tabs.
+It supports chatting with a single model or chatting with two models side-by-side.
+"""
+
+import argparse
+import pickle
+import time
+
+import gradio as gr
+
+from fastchat.serve.gradio_block_arena_anony import (
+    build_side_by_side_ui_anony,
+    load_demo_side_by_side_anony,
+    set_global_vars_anony,
+)
+from fastchat.serve.gradio_block_arena_named import (
+    build_side_by_side_ui_named,
+    load_demo_side_by_side_named,
+    set_global_vars_named,
+)
+from fastchat.serve.gradio_block_arena_vision import (
+    build_single_vision_language_model_ui,
+)
+from fastchat.serve.gradio_web_server import (
+    set_global_vars,
+    block_css,
+    build_single_model_ui,
+    build_about,
+    get_model_list,
+    load_demo_single,
+    get_ip,
+)
+from fastchat.serve.monitor.monitor import build_leaderboard_tab
+from fastchat.utils import (
+    build_logger,
+    get_window_url_params_js,
+    get_window_url_params_with_tos_js,
+    parse_gradio_auth_creds,
+)
+
+logger = build_logger("gradio_web_server_multi", "gradio_web_server_multi.log")
+
+
+def load_demo(url_params, request: gr.Request):
+    global models, all_models, vl_models
+
+    ip = get_ip(request)
+    logger.info(f"load_demo. ip: {ip}. params: {url_params}")
+
+    selected = 0
+    if "arena" in url_params:
+        selected = 0
+    elif "compare" in url_params:
+        selected = 1
+    elif "direct" in url_params or "model" in url_params:
+        selected = 2
+    elif "vision" in url_params:
+        selected = 3
+    elif "leaderboard" in url_params:
+        selected = 4
+
+    if args.model_list_mode == "reload":
+        models, all_models = get_model_list(
+            args.controller_url,
+            args.register_api_endpoint_file,
+            False,
+        )
+
+        vl_models, all_vl_models = get_model_list(
+            args.controller_url,
+            args.register_api_endpoint_file,
+            True,
+        )
+
+    single_updates = load_demo_single(models, url_params)
+    side_by_side_anony_updates = load_demo_side_by_side_anony(all_models, url_params)
+    side_by_side_named_updates = load_demo_side_by_side_named(models, url_params)
+    vision_language_updates = load_demo_single(vl_models, url_params)
+
+    return (
+        (gr.Tabs(selected=selected),)
+        + single_updates
+        + side_by_side_anony_updates
+        + side_by_side_named_updates
+        + vision_language_updates
+    )
+
+
+def build_demo(models, vl_models, elo_results_file, leaderboard_table_file):
+    text_size = gr.themes.sizes.text_md
+    if args.show_terms_of_use:
+        load_js = get_window_url_params_with_tos_js
+    else:
+        load_js = get_window_url_params_js
+
+    head_js = """
+<script src="https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.min.js"></script>
+"""
+    if args.ga_id is not None:
+        head_js += f"""
+<script async src="https://www.googletagmanager.com/gtag/js?id={args.ga_id}"></script>
+<script>
+window.dataLayer = window.dataLayer || [];
+function gtag(){{dataLayer.push(arguments);}}
+gtag('js', new Date());
+
+gtag('config', '{args.ga_id}');
+window.__gradio_mode__ = "app";
+</script>
+        """
+
+    with gr.Blocks(
+        title="Chat with Open Large Language Models",
+        theme=gr.themes.Default(text_size=text_size),
+        css=block_css,
+        head=head_js,
+    ) as demo:
+        with gr.Tabs() as tabs:
+            with gr.Tab("Arena (battle)", id=0):
+                side_by_side_anony_list = build_side_by_side_ui_anony(models)
+
+            with gr.Tab("Arena (side-by-side)", id=1):
+                side_by_side_named_list = build_side_by_side_ui_named(models)
+
+            with gr.Tab("Direct Chat", id=2):
+                single_model_list = build_single_model_ui(
+                    models, add_promotion_links=True
+                )
+
+            with gr.Tab(
+                "Vision-Language Model Direct Chat", id=3, visible=args.multimodal
+            ):
+                single_vision_language_model_list = (
+                    build_single_vision_language_model_ui(
+                        vl_models, add_promotion_links=True
+                    )
+                )
+
+            if elo_results_file:
+                with gr.Tab("Leaderboard", id=4):
+                    build_leaderboard_tab(elo_results_file, leaderboard_table_file)
+
+            with gr.Tab("About Us", id=5):
+                about = build_about()
+
+        url_params = gr.JSON(visible=False)
+
+        if args.model_list_mode not in ["once", "reload"]:
+            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
+
+        demo.load(
+            load_demo,
+            [url_params],
+            [tabs]
+            + single_model_list
+            + side_by_side_anony_list
+            + side_by_side_named_list
+            + single_vision_language_model_list,
+            js=load_js,
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int)
+    parser.add_argument(
+        "--share",
+        action="store_true",
+        help="Whether to generate a public, shareable link",
+    )
+    parser.add_argument(
+        "--controller-url",
+        type=str,
+        default="http://localhost:21001",
+        help="The address of the controller",
+    )
+    parser.add_argument(
+        "--concurrency-count",
+        type=int,
+        default=10,
+        help="The concurrency count of the gradio queue",
+    )
+    parser.add_argument(
+        "--model-list-mode",
+        type=str,
+        default="once",
+        choices=["once", "reload"],
+        help="Whether to load the model list once or reload the model list every time.",
+    )
+    parser.add_argument(
+        "--moderate",
+        action="store_true",
+        help="Enable content moderation to block unsafe inputs",
+    )
+    parser.add_argument(
+        "--show-terms-of-use",
+        action="store_true",
+        help="Shows term of use before loading the demo",
+    )
+    parser.add_argument(
+        "--multimodal", action="store_true", help="Show multi modal tabs."
+    )
+    parser.add_argument(
+        "--register-api-endpoint-file",
+        type=str,
+        help="Register API-based model endpoints from a JSON file",
+    )
+    parser.add_argument(
+        "--gradio-auth-path",
+        type=str,
+        help='Set the gradio authentication file path. The file should contain one or more user:password pairs in this format: "u1:p1,u2:p2,u3:p3"',
+        default=None,
+    )
+    parser.add_argument(
+        "--elo-results-file", type=str, help="Load leaderboard results and plots"
+    )
+    parser.add_argument(
+        "--leaderboard-table-file", type=str, help="Load leaderboard results and plots"
+    )
+    parser.add_argument(
+        "--gradio-root-path",
+        type=str,
+        help="Sets the gradio root path, eg /abc/def. Useful when running behind a reverse-proxy or at a custom URL path prefix",
+    )
+    parser.add_argument(
+        "--ga-id",
+        type=str,
+        help="the Google Analytics ID",
+        default=None,
+    )
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    # Set global variables
+    set_global_vars(args.controller_url, args.moderate)
+    set_global_vars_named(args.moderate)
+    set_global_vars_anony(args.moderate)
+    models, all_models = get_model_list(
+        args.controller_url,
+        args.register_api_endpoint_file,
+        False,
+    )
+
+    vl_models, all_vl_models = get_model_list(
+        args.controller_url,
+        args.register_api_endpoint_file,
+        True,
+    )
+
+    # Set authorization credentials
+    auth = None
+    if args.gradio_auth_path is not None:
+        auth = parse_gradio_auth_creds(args.gradio_auth_path)
+
+    # Launch the demo
+    demo = build_demo(
+        models,
+        vl_models,
+        args.elo_results_file,
+        args.leaderboard_table_file,
+    )
+    demo.queue(
+        default_concurrency_limit=args.concurrency_count,
+        status_update_rate=10,
+        api_open=False,
+    ).launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share,
+        max_threads=200,
+        auth=auth,
+        root_path=args.gradio_root_path,
+    )

huggingface_api.py

@@ -0,0 +1,73 @@
+"""
+Use FastChat with Hugging Face generation APIs.
+
+Usage:
+python3 -m fastchat.serve.huggingface_api --model lmsys/vicuna-7b-v1.5
+python3 -m fastchat.serve.huggingface_api --model lmsys/fastchat-t5-3b-v1.0
+"""
+import argparse
+
+import torch
+
+from fastchat.model import load_model, get_conversation_template, add_model_args
+
+
+@torch.inference_mode()
+def main(args):
+    # Load model
+    model, tokenizer = load_model(
+        args.model_path,
+        device=args.device,
+        num_gpus=args.num_gpus,
+        max_gpu_memory=args.max_gpu_memory,
+        load_8bit=args.load_8bit,
+        cpu_offloading=args.cpu_offloading,
+        revision=args.revision,
+        debug=args.debug,
+    )
+
+    # Build the prompt with a conversation template
+    msg = args.message
+    conv = get_conversation_template(args.model_path)
+    conv.append_message(conv.roles[0], msg)
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt()
+
+    # Run inference
+    inputs = tokenizer([prompt], return_tensors="pt").to(args.device)
+    output_ids = model.generate(
+        **inputs,
+        do_sample=True if args.temperature > 1e-5 else False,
+        temperature=args.temperature,
+        repetition_penalty=args.repetition_penalty,
+        max_new_tokens=args.max_new_tokens,
+    )
+
+    if model.config.is_encoder_decoder:
+        output_ids = output_ids[0]
+    else:
+        output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
+    outputs = tokenizer.decode(
+        output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
+    )
+
+    # Print results
+    print(f"{conv.roles[0]}: {msg}")
+    print(f"{conv.roles[1]}: {outputs}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    add_model_args(parser)
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--repetition_penalty", type=float, default=1.0)
+    parser.add_argument("--max-new-tokens", type=int, default=1024)
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--message", type=str, default="Hello! Who are you?")
+    args = parser.parse_args()
+
+    # Reset default repetition penalty for T5 models.
+    if "t5" in args.model_path and args.repetition_penalty == 1.0:
+        args.repetition_penalty = 1.2
+
+    main(args)

huggingface_api_worker.py

@@ -0,0 +1,415 @@
+"""
+A model worker that calls huggingface inference endpoint.
+
+Register models in a JSON file with the following format:
+{
+    "falcon-180b-chat": {
+        "model_name": "falcon-180B-chat",
+        "api_base": "https://api-inference.huggingface.co/models",
+        "model_path": "tiiuae/falcon-180B-chat",
+        "token": "hf_XXX",
+        "context_length": 2048
+    },
+    "zephyr-7b-beta": {
+        "model_name": "zephyr-7b-beta",
+        "model_path": "",
+        "api_base": "xxx",
+        "token": "hf_XXX",
+        "context_length": 4096
+    }
+}
+
+"model_path", "api_base", "token", and "context_length" are necessary, while others are optional.
+"""
+import argparse
+import asyncio
+import json
+import uuid
+import os
+from typing import List, Optional
+
+import requests
+import uvicorn
+from fastapi import BackgroundTasks, FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+from huggingface_hub import InferenceClient
+
+from fastchat.constants import SERVER_ERROR_MSG, ErrorCode
+from fastchat.serve.base_model_worker import BaseModelWorker
+from fastchat.utils import build_logger
+
+worker_id = str(uuid.uuid4())[:8]
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
+
+workers = []
+worker_map = {}
+app = FastAPI()
+
+
+# reference to
+# https://github.com/philschmid/easyllm/blob/cbd908b3b3f44a97a22cb0fc2c93df3660bacdad/easyllm/clients/huggingface.py#L374-L392
+def get_gen_kwargs(
+    params,
+    seed: Optional[int] = None,
+):
+    stop = params.get("stop", None)
+    if isinstance(stop, list):
+        stop_sequences = stop
+    elif isinstance(stop, str):
+        stop_sequences = [stop]
+    else:
+        stop_sequences = []
+    gen_kwargs = {
+        "do_sample": True,
+        "return_full_text": bool(params.get("echo", False)),
+        "max_new_tokens": int(params.get("max_new_tokens", 256)),
+        "top_p": float(params.get("top_p", 1.0)),
+        "temperature": float(params.get("temperature", 1.0)),
+        "stop_sequences": stop_sequences,
+        "repetition_penalty": float(params.get("repetition_penalty", 1.0)),
+        "top_k": params.get("top_k", None),
+        "seed": seed,
+    }
+    if gen_kwargs["top_p"] == 1:
+        gen_kwargs["top_p"] = 0.9999999
+    if gen_kwargs["top_p"] == 0:
+        gen_kwargs.pop("top_p")
+    if gen_kwargs["temperature"] == 0:
+        gen_kwargs.pop("temperature")
+        gen_kwargs["do_sample"] = False
+    return gen_kwargs
+
+
+def could_be_stop(text, stop):
+    for s in stop:
+        if any(text.endswith(s[:i]) for i in range(1, len(s) + 1)):
+            return True
+    return False
+
+
+class HuggingfaceApiWorker(BaseModelWorker):
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        api_base: str,
+        token: str,
+        context_length: int,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        no_register: bool,
+        conv_template: Optional[str] = None,
+        seed: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            controller_addr,
+            worker_addr,
+            worker_id,
+            model_path,
+            model_names,
+            limit_worker_concurrency,
+            conv_template=conv_template,
+        )
+
+        self.model_path = model_path
+        self.api_base = api_base
+        self.token = token
+        self.context_len = context_length
+        self.seed = seed
+
+        logger.info(
+            f"Connecting with huggingface api {self.model_path} as {self.model_names} on worker {worker_id} ..."
+        )
+
+        if not no_register:
+            self.init_heart_beat()
+
+    def count_token(self, params):
+        # No tokenizer here
+        ret = {
+            "count": 0,
+            "error_code": 0,
+        }
+        return ret
+
+    def generate_stream_gate(self, params):
+        self.call_ct += 1
+
+        prompt = params["prompt"]
+        gen_kwargs = get_gen_kwargs(params, seed=self.seed)
+        stop = gen_kwargs["stop_sequences"]
+        if "falcon" in self.model_path and "chat" in self.model_path:
+            stop.extend(["\nUser:", "<|endoftext|>", " User:", "###"])
+            stop = list(set(stop))
+            gen_kwargs["stop_sequences"] = stop
+
+        logger.info(f"prompt: {prompt}")
+        logger.info(f"gen_kwargs: {gen_kwargs}")
+
+        try:
+            if self.model_path == "":
+                url = f"{self.api_base}"
+            else:
+                url = f"{self.api_base}/{self.model_path}"
+            client = InferenceClient(url, token=self.token)
+            res = client.text_generation(
+                prompt, stream=True, details=True, **gen_kwargs
+            )
+
+            reason = None
+            text = ""
+            for chunk in res:
+                if chunk.token.special:
+                    continue
+                text += chunk.token.text
+
+                s = next((x for x in stop if text.endswith(x)), None)
+                if s is not None:
+                    text = text[: -len(s)]
+                    reason = "stop"
+                    break
+                if could_be_stop(text, stop):
+                    continue
+                if (
+                    chunk.details is not None
+                    and chunk.details.finish_reason is not None
+                ):
+                    reason = chunk.details.finish_reason
+                if reason not in ["stop", "length"]:
+                    reason = None
+                ret = {
+                    "text": text,
+                    "error_code": 0,
+                    "finish_reason": reason,
+                }
+                yield json.dumps(ret).encode() + b"\0"
+        except Exception as e:
+            ret = {
+                "text": f"{SERVER_ERROR_MSG}\n\n({e})",
+                "error_code": ErrorCode.INTERNAL_ERROR,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+    def generate_gate(self, params):
+        for x in self.generate_stream_gate(params):
+            pass
+        return json.loads(x[:-1].decode())
+
+    def get_embeddings(self, params):
+        raise NotImplementedError()
+
+
+def release_worker_semaphore(worker):
+    worker.semaphore.release()
+
+
+def acquire_worker_semaphore(worker):
+    if worker.semaphore is None:
+        worker.semaphore = asyncio.Semaphore(worker.limit_worker_concurrency)
+    return worker.semaphore.acquire()
+
+
+def create_background_tasks(worker):
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(lambda: release_worker_semaphore(worker))
+    return background_tasks
+
+
+@app.post("/worker_generate_stream")
+async def api_generate_stream(request: Request):
+    params = await request.json()
+    worker = worker_map[params["model"]]
+    await acquire_worker_semaphore(worker)
+    generator = worker.generate_stream_gate(params)
+    background_tasks = create_background_tasks(worker)
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_generate")
+async def api_generate(request: Request):
+    params = await request.json()
+    worker = worker_map[params["model"]]
+    await acquire_worker_semaphore(worker)
+    output = worker.generate_gate(params)
+    release_worker_semaphore(worker)
+    return JSONResponse(output)
+
+
+@app.post("/worker_get_embeddings")
+async def api_get_embeddings(request: Request):
+    params = await request.json()
+    worker = worker_map[params["model"]]
+    await acquire_worker_semaphore(worker)
+    embedding = worker.get_embeddings(params)
+    release_worker_semaphore(worker)
+    return JSONResponse(content=embedding)
+
+
+@app.post("/worker_get_status")
+async def api_get_status(request: Request):
+    return {
+        "model_names": [m for w in workers for m in w.model_names],
+        "speed": 1,
+        "queue_length": sum([w.get_queue_length() for w in workers]),
+    }
+
+
+@app.post("/count_token")
+async def api_count_token(request: Request):
+    params = await request.json()
+    worker = worker_map[params["model"]]
+    return worker.count_token(params)
+
+
+@app.post("/worker_get_conv_template")
+async def api_get_conv(request: Request):
+    params = await request.json()
+    worker = worker_map[params["model"]]
+    return worker.get_conv_template()
+
+
+@app.post("/model_details")
+async def api_model_details(request: Request):
+    params = await request.json()
+    worker = worker_map[params["model"]]
+    return {"context_length": worker.context_len}
+
+
+def create_huggingface_api_worker():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+    parser.add_argument(
+        "--controller-address", type=str, default="http://localhost:21001"
+    )
+    # all model-related parameters are listed in --model-info-file
+    parser.add_argument(
+        "--model-info-file",
+        type=str,
+        required=True,
+        help="Huggingface API model's info file path",
+    )
+
+    parser.add_argument(
+        "--limit-worker-concurrency",
+        type=int,
+        default=5,
+        help="Limit the model concurrency to prevent OOM.",
+    )
+    parser.add_argument("--no-register", action="store_true")
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Overwrite the random seed for each generation.",
+    )
+    parser.add_argument(
+        "--ssl",
+        action="store_true",
+        required=False,
+        default=False,
+        help="Enable SSL. Requires OS Environment variables 'SSL_KEYFILE' and 'SSL_CERTFILE'.",
+    )
+    args = parser.parse_args()
+
+    with open(args.model_info_file, "r", encoding="UTF-8") as f:
+        model_info = json.load(f)
+
+    logger.info(f"args: {args}")
+
+    model_path_list = []
+    api_base_list = []
+    token_list = []
+    context_length_list = []
+    model_names_list = []
+    conv_template_list = []
+
+    for m in model_info:
+        model_path_list.append(model_info[m]["model_path"])
+        api_base_list.append(model_info[m]["api_base"])
+        token_list.append(model_info[m]["token"])
+
+        context_length = model_info[m]["context_length"]
+        model_names = model_info[m].get("model_names", [m.split("/")[-1]])
+        if isinstance(model_names, str):
+            model_names = [model_names]
+        conv_template = model_info[m].get("conv_template", None)
+
+        context_length_list.append(context_length)
+        model_names_list.append(model_names)
+        conv_template_list.append(conv_template)
+
+    logger.info(f"Model paths: {model_path_list}")
+    logger.info(f"API bases: {api_base_list}")
+    logger.info(f"Tokens: {token_list}")
+    logger.info(f"Context lengths: {context_length_list}")
+    logger.info(f"Model names: {model_names_list}")
+    logger.info(f"Conv templates: {conv_template_list}")
+
+    for (
+        model_names,
+        conv_template,
+        model_path,
+        api_base,
+        token,
+        context_length,
+    ) in zip(
+        model_names_list,
+        conv_template_list,
+        model_path_list,
+        api_base_list,
+        token_list,
+        context_length_list,
+    ):
+        m = HuggingfaceApiWorker(
+            args.controller_address,
+            args.worker_address,
+            worker_id,
+            model_path,
+            api_base,
+            token,
+            context_length,
+            model_names,
+            args.limit_worker_concurrency,
+            no_register=args.no_register,
+            conv_template=conv_template,
+            seed=args.seed,
+        )
+        workers.append(m)
+        for name in model_names:
+            worker_map[name] = m
+
+    # register all the models
+    url = args.controller_address + "/register_worker"
+    data = {
+        "worker_name": workers[0].worker_addr,
+        "check_heart_beat": not args.no_register,
+        "worker_status": {
+            "model_names": [m for w in workers for m in w.model_names],
+            "speed": 1,
+            "queue_length": sum([w.get_queue_length() for w in workers]),
+        },
+    }
+    r = requests.post(url, json=data)
+    assert r.status_code == 200
+
+    return args, workers
+
+
+if __name__ == "__main__":
+    args, workers = create_huggingface_api_worker()
+    if args.ssl:
+        uvicorn.run(
+            app,
+            host=args.host,
+            port=args.port,
+            log_level="info",
+            ssl_keyfile=os.environ["SSL_KEYFILE"],
+            ssl_certfile=os.environ["SSL_CERTFILE"],
+        )
+    else:
+        uvicorn.run(app, host=args.host, port=args.port, log_level="info")

inference.py

@@ -0,0 +1,555 @@
+"""Inference for FastChat models."""
+import abc
+import gc
+import json
+import math
+import os
+import sys
+import time
+from typing import Iterable, Optional, Dict
+import warnings
+
+import psutil
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    AutoModel,
+    AutoModelForSeq2SeqLM,
+    T5Tokenizer,
+    AutoConfig,
+)
+from transformers.generation.logits_process import (
+    LogitsProcessorList,
+    RepetitionPenaltyLogitsProcessor,
+    TemperatureLogitsWarper,
+    TopKLogitsWarper,
+    TopPLogitsWarper,
+)
+
+from fastchat.conversation import get_conv_template, SeparatorStyle
+from fastchat.model.model_adapter import (
+    load_model,
+    get_conversation_template,
+    get_generate_stream_function,
+)
+from fastchat.modules.awq import AWQConfig
+from fastchat.modules.gptq import GptqConfig
+from fastchat.modules.exllama import ExllamaConfig
+from fastchat.modules.xfastertransformer import XftConfig
+from fastchat.utils import is_partial_stop, is_sentence_complete, get_context_length
+
+
+def prepare_logits_processor(
+    temperature: float, repetition_penalty: float, top_p: float, top_k: int
+) -> LogitsProcessorList:
+    processor_list = LogitsProcessorList()
+    # TemperatureLogitsWarper doesn't accept 0.0, 1.0 makes it a no-op so we skip two cases.
+    if temperature >= 1e-5 and temperature != 1.0:
+        processor_list.append(TemperatureLogitsWarper(temperature))
+    if repetition_penalty > 1.0:
+        processor_list.append(RepetitionPenaltyLogitsProcessor(repetition_penalty))
+    if 1e-8 <= top_p < 1.0:
+        processor_list.append(TopPLogitsWarper(top_p))
+    if top_k > 0:
+        processor_list.append(TopKLogitsWarper(top_k))
+    return processor_list
+
+
+@torch.inference_mode()
+def generate_stream(
+    model,
+    tokenizer,
+    params: Dict,
+    device: str,
+    context_len: int,
+    stream_interval: int = 2,
+    judge_sent_end: bool = False,
+):
+    if hasattr(model, "device"):
+        device = model.device
+
+    # Read parameters
+    prompt = params["prompt"]
+    len_prompt = len(prompt)
+    temperature = float(params.get("temperature", 1.0))
+    repetition_penalty = float(params.get("repetition_penalty", 1.0))
+    top_p = float(params.get("top_p", 1.0))
+    top_k = int(params.get("top_k", -1))  # -1 means disable
+    max_new_tokens = int(params.get("max_new_tokens", 256))
+    logprobs = params.get("logprobs", None)  # FIXME: Support logprobs>1.
+    echo = bool(params.get("echo", True))
+    stop_str = params.get("stop", None)
+    stop_token_ids = params.get("stop_token_ids", None) or []
+    if tokenizer.eos_token_id not in stop_token_ids:
+        stop_token_ids.append(tokenizer.eos_token_id)
+
+    logits_processor = prepare_logits_processor(
+        temperature, repetition_penalty, top_p, top_k
+    )
+    input_ids = tokenizer(prompt).input_ids
+
+    if model.config.is_encoder_decoder:
+        max_src_len = context_len
+    else:  # truncate
+        max_src_len = context_len - max_new_tokens - 1
+
+    input_ids = input_ids[-max_src_len:]
+    output_ids = list(input_ids)
+    input_echo_len = len(input_ids)
+
+    if model.config.is_encoder_decoder:
+        if logprobs is not None:  # FIXME: Support logprobs for encoder-decoder models.
+            raise NotImplementedError
+        encoder_output = model.encoder(
+            input_ids=torch.as_tensor([input_ids], device=device)
+        )[0]
+        start_ids = torch.as_tensor(
+            [[model.generation_config.decoder_start_token_id]],
+            dtype=torch.int64,
+            device=device,
+        )
+    else:
+        start_ids = torch.as_tensor([input_ids], device=device)
+
+    past_key_values = out = None
+    token_logprobs = [None]  # The first token has no logprobs.
+    sent_interrupt = False
+    finish_reason = None
+    stopped = False
+    for i in range(max_new_tokens):
+        if i == 0:  # prefill
+            if model.config.is_encoder_decoder:
+                out = model.decoder(
+                    input_ids=start_ids,
+                    encoder_hidden_states=encoder_output,
+                    use_cache=True,
+                )
+                logits = model.lm_head(out[0])
+            else:
+                out = model(input_ids=start_ids, use_cache=True)
+                logits = out.logits
+            past_key_values = out.past_key_values
+
+            if logprobs is not None:
+                # Prefull logprobs for the prompt.
+                shift_input_ids = start_ids[..., 1:].contiguous()
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_logits = torch.log_softmax(shift_logits, dim=-1).tolist()
+                for label_id, logit in zip(
+                    shift_input_ids[0].tolist(), shift_logits[0]
+                ):
+                    token_logprobs.append(logit[label_id])
+        else:  # decoding
+            if model.config.is_encoder_decoder:
+                out = model.decoder(
+                    input_ids=torch.as_tensor(
+                        [[token] if not sent_interrupt else output_ids],
+                        device=device,
+                    ),
+                    encoder_hidden_states=encoder_output,
+                    use_cache=True,
+                    past_key_values=past_key_values if not sent_interrupt else None,
+                )
+                sent_interrupt = False
+
+                logits = model.lm_head(out[0])
+            else:
+                out = model(
+                    input_ids=torch.as_tensor(
+                        [[token] if not sent_interrupt else output_ids],
+                        device=device,
+                    ),
+                    use_cache=True,
+                    past_key_values=past_key_values if not sent_interrupt else None,
+                )
+                sent_interrupt = False
+                logits = out.logits
+            past_key_values = out.past_key_values
+
+        if logits_processor:
+            if repetition_penalty > 1.0:
+                tmp_output_ids = torch.as_tensor([output_ids], device=logits.device)
+            else:
+                tmp_output_ids = None
+            last_token_logits = logits_processor(tmp_output_ids, logits[:, -1, :])[0]
+        else:
+            last_token_logits = logits[0, -1, :]
+
+        if device == "mps":
+            # Switch to CPU by avoiding some bugs in mps backend.
+            last_token_logits = last_token_logits.float().to("cpu")
+
+        if temperature < 1e-5 or top_p < 1e-8:  # greedy
+            _, indices = torch.topk(last_token_logits, 2)
+            tokens = [int(index) for index in indices.tolist()]
+        else:
+            probs = torch.softmax(last_token_logits, dim=-1)
+            indices = torch.multinomial(probs, num_samples=2)
+            tokens = [int(token) for token in indices.tolist()]
+        token = tokens[0]
+        output_ids.append(token)
+        if logprobs is not None:
+            # Cannot use last_token_logits because logprobs is based on raw logits.
+            token_logprobs.append(
+                torch.log_softmax(logits[0, -1, :], dim=-1)[token].tolist()
+            )
+
+        if token in stop_token_ids:
+            stopped = True
+        else:
+            stopped = False
+
+        # Yield the output tokens
+        if i % stream_interval == 0 or i == max_new_tokens - 1 or stopped:
+            if echo:
+                tmp_output_ids = output_ids
+                rfind_start = len_prompt
+            else:
+                tmp_output_ids = output_ids[input_echo_len:]
+                rfind_start = 0
+
+            output = tokenizer.decode(
+                tmp_output_ids,
+                skip_special_tokens=True,
+                spaces_between_special_tokens=False,
+                clean_up_tokenization_spaces=True,
+            )
+            ret_logprobs = None
+            if logprobs is not None:
+                ret_logprobs = {
+                    "text_offset": [],
+                    "tokens": [
+                        tokenizer.decode(token)
+                        for token in (
+                            output_ids if echo else output_ids[input_echo_len:]
+                        )
+                    ],
+                    "token_logprobs": token_logprobs
+                    if echo
+                    else token_logprobs[input_echo_len:],
+                    "top_logprobs": [{}]
+                    * len(token_logprobs if echo else token_logprobs[input_echo_len:]),
+                }
+                # Compute text_offset
+                curr_pos = 0
+                for text in ret_logprobs["tokens"]:
+                    ret_logprobs["text_offset"].append(curr_pos)
+                    curr_pos += len(text)
+
+            # TODO: For the issue of incomplete sentences interrupting output, apply a patch and others can also modify it to a more elegant way
+            if judge_sent_end and stopped and not is_sentence_complete(output):
+                if len(tokens) > 1:
+                    token = tokens[1]
+                    output_ids[-1] = token
+                else:
+                    output_ids.pop()
+                stopped = False
+                sent_interrupt = True
+
+            partially_stopped = False
+            if stop_str:
+                if isinstance(stop_str, str):
+                    pos = output.rfind(stop_str, rfind_start)
+                    if pos != -1:
+                        output = output[:pos]
+                        stopped = True
+                    else:
+                        partially_stopped = is_partial_stop(output, stop_str)
+                elif isinstance(stop_str, Iterable):
+                    for each_stop in stop_str:
+                        pos = output.rfind(each_stop, rfind_start)
+                        if pos != -1:
+                            output = output[:pos]
+                            stopped = True
+                            break
+                        else:
+                            partially_stopped = is_partial_stop(output, each_stop)
+                            if partially_stopped:
+                                break
+                else:
+                    raise ValueError("Invalid stop field type.")
+
+            # Prevent yielding partial stop sequence
+            if not partially_stopped:
+                yield {
+                    "text": output,
+                    "logprobs": ret_logprobs,
+                    "usage": {
+                        "prompt_tokens": input_echo_len,
+                        "completion_tokens": i,
+                        "total_tokens": input_echo_len + i,
+                    },
+                    "finish_reason": None,
+                }
+
+        if stopped:
+            break
+
+    # Finish stream event, which contains finish reason
+    else:
+        finish_reason = "length"
+
+    if stopped:
+        finish_reason = "stop"
+
+    yield {
+        "text": output,
+        "logprobs": ret_logprobs,
+        "usage": {
+            "prompt_tokens": input_echo_len,
+            "completion_tokens": i,
+            "total_tokens": input_echo_len + i,
+        },
+        "finish_reason": finish_reason,
+    }
+
+    # Clean
+    del past_key_values, out
+    gc.collect()
+    torch.cuda.empty_cache()
+    if device == "xpu":
+        torch.xpu.empty_cache()
+    if device == "npu":
+        torch.npu.empty_cache()
+
+
+class ChatIO(abc.ABC):
+    @abc.abstractmethod
+    def prompt_for_input(self, role: str) -> str:
+        """Prompt for input from a role."""
+
+    @abc.abstractmethod
+    def prompt_for_output(self, role: str):
+        """Prompt for output from a role."""
+
+    @abc.abstractmethod
+    def stream_output(self, output_stream):
+        """Stream output."""
+
+    @abc.abstractmethod
+    def print_output(self, text: str):
+        """Print output."""
+
+
+def chat_loop(
+    model_path: str,
+    device: str,
+    num_gpus: int,
+    max_gpu_memory: str,
+    dtype: Optional[torch.dtype],
+    load_8bit: bool,
+    cpu_offloading: bool,
+    conv_template: Optional[str],
+    conv_system_msg: Optional[str],
+    temperature: float,
+    repetition_penalty: float,
+    max_new_tokens: int,
+    chatio: ChatIO,
+    gptq_config: Optional[GptqConfig] = None,
+    awq_config: Optional[AWQConfig] = None,
+    exllama_config: Optional[ExllamaConfig] = None,
+    xft_config: Optional[XftConfig] = None,
+    revision: str = "main",
+    judge_sent_end: bool = True,
+    debug: bool = True,
+    history: bool = True,
+):
+    # Model
+    model, tokenizer = load_model(
+        model_path,
+        device=device,
+        num_gpus=num_gpus,
+        max_gpu_memory=max_gpu_memory,
+        dtype=dtype,
+        load_8bit=load_8bit,
+        cpu_offloading=cpu_offloading,
+        gptq_config=gptq_config,
+        awq_config=awq_config,
+        exllama_config=exllama_config,
+        xft_config=xft_config,
+        revision=revision,
+        debug=debug,
+    )
+    generate_stream_func = get_generate_stream_function(model, model_path)
+
+    model_type = str(type(model)).lower()
+    is_t5 = "t5" in model_type
+    is_codet5p = "codet5p" in model_type
+    is_xft = "xft" in model_type
+
+    # Hardcode T5's default repetition penalty to be 1.2
+    if is_t5 and repetition_penalty == 1.0:
+        repetition_penalty = 1.2
+
+    # Set context length
+    context_len = get_context_length(model.config)
+
+    # Chat
+    def new_chat():
+        if conv_template:
+            conv = get_conv_template(conv_template)
+        else:
+            conv = get_conversation_template(model_path)
+        if conv_system_msg is not None:
+            conv.set_system_message(conv_system_msg)
+        return conv
+
+    def reload_conv(conv):
+        """
+        Reprints the conversation from the start.
+        """
+        for message in conv.messages[conv.offset :]:
+            chatio.prompt_for_output(message[0])
+            chatio.print_output(message[1])
+
+    conv = None
+
+    while True:
+        if not history or not conv:
+            conv = new_chat()
+
+        try:
+            inp = chatio.prompt_for_input(conv.roles[0])
+        except EOFError:
+            inp = ""
+
+        if inp == "!!exit" or not inp:
+            print("exit...")
+            break
+        elif inp == "!!reset":
+            print("resetting...")
+            conv = new_chat()
+            continue
+        elif inp == "!!remove":
+            print("removing last message...")
+            if len(conv.messages) > conv.offset:
+                # Assistant
+                if conv.messages[-1][0] == conv.roles[1]:
+                    conv.messages.pop()
+                # User
+                if conv.messages[-1][0] == conv.roles[0]:
+                    conv.messages.pop()
+                reload_conv(conv)
+            else:
+                print("No messages to remove.")
+            continue
+        elif inp == "!!regen":
+            print("regenerating last message...")
+            if len(conv.messages) > conv.offset:
+                # Assistant
+                if conv.messages[-1][0] == conv.roles[1]:
+                    conv.messages.pop()
+                # User
+                if conv.messages[-1][0] == conv.roles[0]:
+                    reload_conv(conv)
+                    # Set inp to previous message
+                    inp = conv.messages.pop()[1]
+                else:
+                    # Shouldn't happen in normal circumstances
+                    print("No user message to regenerate from.")
+                    continue
+            else:
+                print("No messages to regenerate.")
+                continue
+        elif inp.startswith("!!save"):
+            args = inp.split(" ", 1)
+
+            if len(args) != 2:
+                print("usage: !!save <filename>")
+                continue
+            else:
+                filename = args[1]
+
+            # Add .json if extension not present
+            if not "." in filename:
+                filename += ".json"
+
+            print("saving...", filename)
+            with open(filename, "w") as outfile:
+                json.dump(conv.dict(), outfile)
+            continue
+        elif inp.startswith("!!load"):
+            args = inp.split(" ", 1)
+
+            if len(args) != 2:
+                print("usage: !!load <filename>")
+                continue
+            else:
+                filename = args[1]
+
+            # Check if file exists and add .json if needed
+            if not os.path.exists(filename):
+                if (not filename.endswith(".json")) and os.path.exists(
+                    filename + ".json"
+                ):
+                    filename += ".json"
+                else:
+                    print("file not found:", filename)
+                    continue
+
+            print("loading...", filename)
+            with open(filename, "r") as infile:
+                new_conv = json.load(infile)
+
+            conv = get_conv_template(new_conv["template_name"])
+            conv.set_system_message(new_conv["system_message"])
+            conv.messages = new_conv["messages"]
+            reload_conv(conv)
+            continue
+
+        conv.append_message(conv.roles[0], inp)
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+
+        if is_codet5p:  # codet5p is a code completion model.
+            prompt = inp
+
+        gen_params = {
+            "model": model_path,
+            "prompt": prompt,
+            "temperature": temperature,
+            "repetition_penalty": repetition_penalty,
+            "max_new_tokens": max_new_tokens,
+            "stop": conv.stop_str,
+            "stop_token_ids": conv.stop_token_ids,
+            "echo": False,
+        }
+
+        try:
+            chatio.prompt_for_output(conv.roles[1])
+            output_stream = generate_stream_func(
+                model,
+                tokenizer,
+                gen_params,
+                device,
+                context_len=context_len,
+                judge_sent_end=judge_sent_end,
+            )
+            t = time.time()
+            outputs = chatio.stream_output(output_stream)
+            duration = time.time() - t
+            conv.update_last_message(outputs.strip())
+
+            if debug:
+                num_tokens = len(tokenizer.encode(outputs))
+                msg = {
+                    "conv_template": conv.name,
+                    "prompt": prompt,
+                    "outputs": outputs,
+                    "speed (token/s)": round(num_tokens / duration, 2),
+                }
+                print(f"\n{msg}\n")
+
+        except KeyboardInterrupt:
+            print("stopped generation.")
+            # If generation didn't finish
+            if conv.messages[-1][1] is None:
+                conv.messages.pop()
+                # Remove last user message, so there isn't a double up
+                if conv.messages[-1][0] == conv.roles[0]:
+                    conv.messages.pop()
+
+                reload_conv(conv)

launch_all_serve.py

@@ -0,0 +1,284 @@
+"""
+Usage: python launch_all_serve_by_shell.py --model-path-address "THUDM/chatglm2-6b@localhost@2021" "huggyllama/llama-7b@localhost@2022" 
+
+Workers are listed in format of `model-path`@`host`@`port` 
+
+The key mechanism behind this scripts is: 
+    1, execute shell cmd to launch the controller/worker/openai-api-server;
+    2, check the log of controller/worker/openai-api-server to ensure that the serve is launched properly.
+Note that a few of non-critical `fastchat.serve` cmd options are not supported currently.
+"""
+import sys
+import os
+
+sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+
+import subprocess
+import re
+import argparse
+
+LOGDIR = "./logs/"
+
+if not os.path.exists(LOGDIR):
+    os.makedirs(LOGDIR)
+
+parser = argparse.ArgumentParser()
+# ------multi worker-----------------
+parser.add_argument(
+    "--model-path-address",
+    default="THUDM/chatglm2-6b@localhost@20002",
+    nargs="+",
+    type=str,
+    help="model path, host, and port, formatted as model-path@host@port",
+)
+# ---------------controller-------------------------
+
+parser.add_argument("--controller-host", type=str, default="localhost")
+parser.add_argument("--controller-port", type=int, default=21001)
+parser.add_argument(
+    "--dispatch-method",
+    type=str,
+    choices=["lottery", "shortest_queue"],
+    default="shortest_queue",
+)
+controller_args = ["controller-host", "controller-port", "dispatch-method"]
+
+# ----------------------worker------------------------------------------
+
+parser.add_argument("--worker-host", type=str, default="localhost")
+parser.add_argument("--worker-port", type=int, default=21002)
+# parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+# parser.add_argument(
+#     "--controller-address", type=str, default="http://localhost:21001"
+# )
+parser.add_argument(
+    "--model-path",
+    type=str,
+    default="lmsys/vicuna-7b-v1.5",
+    help="The path to the weights. This can be a local folder or a Hugging Face repo ID.",
+)
+parser.add_argument(
+    "--revision",
+    type=str,
+    default="main",
+    help="Hugging Face Hub model revision identifier",
+)
+parser.add_argument(
+    "--device",
+    type=str,
+    choices=["cpu", "cuda", "mps", "xpu", "npu"],
+    default="cuda",
+    help="The device type",
+)
+parser.add_argument(
+    "--gpus",
+    type=str,
+    default="0",
+    help="A single GPU like 1 or multiple GPUs like 0,2",
+)
+parser.add_argument("--num-gpus", type=int, default=1)
+parser.add_argument(
+    "--max-gpu-memory",
+    type=str,
+    help="The maximum memory per gpu. Use a string like '13Gib'",
+)
+parser.add_argument("--load-8bit", action="store_true", help="Use 8-bit quantization")
+parser.add_argument(
+    "--cpu-offloading",
+    action="store_true",
+    help="Only when using 8-bit quantization: Offload excess weights to the CPU that don't fit on the GPU",
+)
+parser.add_argument(
+    "--gptq-ckpt",
+    type=str,
+    default=None,
+    help="Load quantized model. The path to the local GPTQ checkpoint.",
+)
+parser.add_argument(
+    "--gptq-wbits",
+    type=int,
+    default=16,
+    choices=[2, 3, 4, 8, 16],
+    help="#bits to use for quantization",
+)
+parser.add_argument(
+    "--gptq-groupsize",
+    type=int,
+    default=-1,
+    help="Groupsize to use for quantization; default uses full row.",
+)
+parser.add_argument(
+    "--gptq-act-order",
+    action="store_true",
+    help="Whether to apply the activation order GPTQ heuristic",
+)
+parser.add_argument(
+    "--model-names",
+    type=lambda s: s.split(","),
+    help="Optional display comma separated names",
+)
+parser.add_argument(
+    "--limit-worker-concurrency",
+    type=int,
+    default=5,
+    help="Limit the model concurrency to prevent OOM.",
+)
+parser.add_argument("--stream-interval", type=int, default=2)
+parser.add_argument("--no-register", action="store_true")
+
+worker_args = [
+    "worker-host",
+    "worker-port",
+    "model-path",
+    "revision",
+    "device",
+    "gpus",
+    "num-gpus",
+    "max-gpu-memory",
+    "load-8bit",
+    "cpu-offloading",
+    "gptq-ckpt",
+    "gptq-wbits",
+    "gptq-groupsize",
+    "gptq-act-order",
+    "model-names",
+    "limit-worker-concurrency",
+    "stream-interval",
+    "no-register",
+    "controller-address",
+]
+# -----------------openai server---------------------------
+
+parser.add_argument("--server-host", type=str, default="localhost", help="host name")
+parser.add_argument("--server-port", type=int, default=8001, help="port number")
+parser.add_argument(
+    "--allow-credentials", action="store_true", help="allow credentials"
+)
+# parser.add_argument(
+#     "--allowed-origins", type=json.loads, default=["*"], help="allowed origins"
+# )
+# parser.add_argument(
+#     "--allowed-methods", type=json.loads, default=["*"], help="allowed methods"
+# )
+# parser.add_argument(
+#     "--allowed-headers", type=json.loads, default=["*"], help="allowed headers"
+# )
+parser.add_argument(
+    "--api-keys",
+    type=lambda s: s.split(","),
+    help="Optional list of comma separated API keys",
+)
+server_args = [
+    "server-host",
+    "server-port",
+    "allow-credentials",
+    "api-keys",
+    "controller-address",
+]
+
+args = parser.parse_args()
+
+args = argparse.Namespace(
+    **vars(args),
+    **{"controller-address": f"http://{args.controller_host}:{args.controller_port}"},
+)
+
+if args.gpus:
+    if len(args.gpus.split(",")) < args.num_gpus:
+        raise ValueError(
+            f"Larger --num-gpus ({args.num_gpus}) than --gpus {args.gpus}!"
+        )
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
+
+# 0,controller, model_worker, openai_api_server
+# 1, cmd options
+# 2,LOGDIR
+# 3, log file name
+base_launch_sh = "nohup python3 -m fastchat.serve.{0} {1} >{2}/{3}.log 2>&1 &"
+
+# 0 LOGDIR
+#! 1 log file name
+# 2 controller, worker, openai_api_server
+base_check_sh = """while [ `grep -c "Uvicorn running on" {0}/{1}.log` -eq '0' ];do
+                        sleep 1s;
+                        echo "wait {2} running"
+                done
+                echo '{2} running' """
+
+
+def string_args(args, args_list):
+    args_str = ""
+    for key, value in args._get_kwargs():
+        key = key.replace("_", "-")
+        if key not in args_list:
+            continue
+
+        key = key.split("-")[-1] if re.search("port|host", key) else key
+        if not value:
+            pass
+        # 1==True ->  True
+        elif isinstance(value, bool) and value == True:
+            args_str += f" --{key} "
+        elif (
+            isinstance(value, list)
+            or isinstance(value, tuple)
+            or isinstance(value, set)
+        ):
+            value = " ".join(value)
+            args_str += f" --{key} {value} "
+        else:
+            args_str += f" --{key} {value} "
+
+    return args_str
+
+
+def launch_worker(item):
+    log_name = (
+        item.split("/")[-1]
+        .split("\\")[-1]
+        .replace("-", "_")
+        .replace("@", "_")
+        .replace(".", "_")
+    )
+
+    args.model_path, args.worker_host, args.worker_port = item.split("@")
+    print("*" * 80)
+    worker_str_args = string_args(args, worker_args)
+    print(worker_str_args)
+    worker_sh = base_launch_sh.format(
+        "model_worker", worker_str_args, LOGDIR, f"worker_{log_name}"
+    )
+    worker_check_sh = base_check_sh.format(LOGDIR, f"worker_{log_name}", "model_worker")
+    subprocess.run(worker_sh, shell=True, check=True)
+    subprocess.run(worker_check_sh, shell=True, check=True)
+
+
+def launch_all():
+    controller_str_args = string_args(args, controller_args)
+    controller_sh = base_launch_sh.format(
+        "controller", controller_str_args, LOGDIR, "controller"
+    )
+    controller_check_sh = base_check_sh.format(LOGDIR, "controller", "controller")
+    subprocess.run(controller_sh, shell=True, check=True)
+    subprocess.run(controller_check_sh, shell=True, check=True)
+
+    if isinstance(args.model_path_address, str):
+        launch_worker(args.model_path_address)
+    else:
+        for idx, item in enumerate(args.model_path_address):
+            print(f"loading {idx}th model:{item}")
+            launch_worker(item)
+
+    server_str_args = string_args(args, server_args)
+    server_sh = base_launch_sh.format(
+        "openai_api_server", server_str_args, LOGDIR, "openai_api_server"
+    )
+    server_check_sh = base_check_sh.format(
+        LOGDIR, "openai_api_server", "openai_api_server"
+    )
+    subprocess.run(server_sh, shell=True, check=True)
+    subprocess.run(server_check_sh, shell=True, check=True)
+
+
+if __name__ == "__main__":
+    launch_all()

lightllm_worker.py

@@ -0,0 +1,512 @@
+"""
+A model worker that executes the model based on LightLLM.
+
+See documentations at docs/lightllm_integration.md
+"""
+
+import argparse
+import asyncio
+import json
+import os
+import torch
+import uvicorn
+
+from transformers import AutoConfig
+
+from typing import List
+
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse, JSONResponse
+
+from fastchat.serve.base_model_worker import BaseModelWorker
+from fastchat.serve.model_worker import (
+    logger,
+    worker_id,
+)
+
+from lightllm.server.sampling_params import SamplingParams
+from lightllm.server.multimodal_params import MultimodalParams
+from lightllm.server.httpserver.manager import HttpServerManager
+from lightllm.server.detokenization.manager import start_detokenization_process
+from lightllm.server.router.manager import start_router_process
+from lightllm.server.req_id_generator import ReqIDGenerator
+
+from lightllm.utils.net_utils import alloc_can_use_network_port
+from lightllm.utils.start_utils import start_submodule_processes
+from fastchat.utils import get_context_length, is_partial_stop
+
+app = FastAPI()
+g_id_gen = ReqIDGenerator()
+
+
+class LightLLMWorker(BaseModelWorker):
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        no_register: bool,
+        conv_template: str,
+        tokenizer,
+        context_len,
+    ):
+        super().__init__(
+            controller_addr,
+            worker_addr,
+            worker_id,
+            model_path,
+            model_names,
+            limit_worker_concurrency,
+            conv_template,
+        )
+
+        logger.info(
+            f"Loading the model {self.model_names} on worker {worker_id}, worker type: LightLLM worker..."
+        )
+        self.tokenizer = tokenizer
+        self.context_len = context_len
+
+        self.is_first = True
+
+        if not no_register:
+            self.init_heart_beat()
+
+    async def generate_stream(self, params):
+        self.call_ct += 1
+
+        prompt = params.pop("prompt")
+        request_id = params.pop("request_id")
+        temperature = float(params.get("temperature", 1.0))
+        top_p = float(params.get("top_p", 1.0))
+        top_k = params.get("top_k", -1.0)
+        presence_penalty = float(params.get("presence_penalty", 0.0))
+        frequency_penalty = float(params.get("frequency_penalty", 0.0))
+        repetition_penalty = float(params.get("repetition_penalty", 1.0))
+        max_new_tokens = params.get("max_new_tokens", 256)
+        echo = params.get("echo", True)
+        stop_str = params.get("stop", None)
+        stop_token_ids = params.get("stop_token_ids", None) or []
+        if self.tokenizer.eos_token_id is not None:
+            stop_token_ids.append(self.tokenizer.eos_token_id)
+
+        request = params.get("request", None)
+
+        # Handle stop_str
+        stop = set()
+        if isinstance(stop_str, str) and stop_str != "":
+            stop.add(stop_str)
+        elif isinstance(stop_str, list) and stop_str != []:
+            stop.update(stop_str)
+
+        for tid in stop_token_ids:
+            if tid is not None:
+                s = self.tokenizer.decode(tid)
+                if s != "":
+                    stop.add(s)
+
+        if self.is_first:
+            loop = asyncio.get_event_loop()
+            loop.create_task(httpserver_manager.handle_loop())
+            self.is_first = False
+
+        # make sampling params in vllm
+        top_p = max(top_p, 1e-5)
+        if temperature <= 1e-5:
+            top_p = 1.0
+
+        sampling_params = SamplingParams(
+            do_sample=temperature > 0.0,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            presence_penalty=presence_penalty,
+            frequency_penalty=frequency_penalty,
+            repetition_penalty=repetition_penalty,
+            max_new_tokens=max_new_tokens,
+            stop_sequences=list(stop),
+        )
+        sampling_params.verify()
+
+        results_generator = httpserver_manager.generate(
+            prompt, sampling_params, request_id, MultimodalParams()
+        )
+
+        completion_tokens = 0
+        text_outputs = ""
+        cumulative_logprob = 0.0
+
+        async for request_output, metadata, finish_status in results_generator:
+            text_outputs += request_output
+            completion_tokens += 1
+
+            partial_stop = any(is_partial_stop(text_outputs, i) for i in stop)
+            # prevent yielding partial stop sequence
+            if partial_stop:
+                continue
+
+            if type(finish_status) is bool:  # compatibility with old version
+                finish_reason = "stop" if finish_status else None
+            else:
+                finish_reason = finish_status.get_finish_reason()
+
+            if request and await request.is_disconnected():
+                await httpserver_manager.abort(request_id)
+                finish_reason = "abort"
+
+            logprob = metadata.get("logprob", None)
+            if logprob is not None:
+                cumulative_logprob += logprob
+
+            prompt_tokens = metadata["prompt_tokens"]
+            ret = {
+                "text": prompt + text_outputs if echo else text_outputs,
+                "error_code": 0,
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": prompt_tokens + completion_tokens,
+                },
+                "cumulative_logprob": cumulative_logprob,
+            }
+
+            if finish_reason is not None:
+                yield (
+                    json.dumps({**ret, "finish_reason": None}, ensure_ascii=False)
+                    + "\0"
+                ).encode("utf-8")
+            yield (
+                json.dumps({**ret, "finish_reason": finish_reason}, ensure_ascii=False)
+                + "\0"
+            ).encode("utf-8")
+
+            if finish_reason is not None:  # In case of abort, we need to break the loop
+                break
+
+    async def generate(self, params):
+        async for x in self.generate_stream(params):
+            pass
+        return json.loads(x[:-1].decode())
+
+
+def release_worker_semaphore():
+    worker.semaphore.release()
+
+
+def acquire_worker_semaphore():
+    if worker.semaphore is None:
+        worker.semaphore = asyncio.Semaphore(worker.limit_worker_concurrency)
+    return worker.semaphore.acquire()
+
+
+def create_background_tasks(request_id):
+    async def abort_request() -> None:
+        await httpserver_manager.abort(request_id)
+
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(release_worker_semaphore)
+    background_tasks.add_task(abort_request)
+    return background_tasks
+
+
+@app.post("/worker_generate_stream")
+async def api_generate_stream(request: Request):
+    params = await request.json()
+    await acquire_worker_semaphore()
+    request_id = g_id_gen.generate_id()
+    params["request_id"] = request_id
+    params["request"] = request
+    generator = worker.generate_stream(params)
+    background_tasks = create_background_tasks(request_id)
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_generate")
+async def api_generate(request: Request):
+    params = await request.json()
+    await acquire_worker_semaphore()
+    request_id = g_id_gen.generate_id()
+    params["request_id"] = request_id
+    params["request"] = request
+    output = await worker.generate(params)
+    release_worker_semaphore()
+    await httpserver_manager.abort(request_id)
+    return JSONResponse(output)
+
+
+@app.post("/worker_get_status")
+async def api_get_status(request: Request):
+    return worker.get_status()
+
+
+@app.post("/count_token")
+async def api_count_token(request: Request):
+    params = await request.json()
+    return worker.count_token(params)
+
+
+@app.post("/worker_get_conv_template")
+async def api_get_conv(request: Request):
+    return worker.get_conv_template()
+
+
+@app.post("/model_details")
+async def api_model_details(request: Request):
+    return {"context_length": worker.context_len}
+
+
+if __name__ == "__main__":
+    torch.multiprocessing.set_start_method("spawn")
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+
+    parser.add_argument(
+        "--model-path",
+        dest="model_dir",
+        type=str,
+        default=None,
+        help="the model weight dir path, the app will load config, weights and tokenizer from this dir",
+    )
+    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+    parser.add_argument(
+        "--controller-address", type=str, default="http://localhost:21001"
+    )
+    parser.add_argument(
+        "--conv-template", type=str, default=None, help="Conversation prompt template."
+    )
+    parser.add_argument(
+        "--model-names",
+        type=lambda s: s.split(","),
+        help="Optional display comma separated names",
+    )
+    parser.add_argument("--limit-worker-concurrency", type=int, default=1024)
+    parser.add_argument("--no-register", action="store_true")
+
+    parser.add_argument(
+        "--tokenizer_mode",
+        type=str,
+        default="slow",
+        help="""tokenizer load mode, can be slow or auto, slow mode load fast but run slow, slow mode is good for debug and test,
+                        when you want to get best performance, try auto mode""",
+    )
+    parser.add_argument(
+        "--load_way",
+        type=str,
+        default="HF",
+        help="the way of loading model weights, the default is HF(Huggingface format), llama also supports DS(Deepspeed)",
+    )
+    parser.add_argument(
+        "--max_total_token_num",
+        type=int,
+        default=6000,
+        help="the total token nums the gpu and model can support, equals = max_batch * (input_len + output_len)",
+    )
+    parser.add_argument(
+        "--batch_max_tokens",
+        type=int,
+        default=None,
+        help="max tokens num for new cat batch, it control prefill batch size to Preventing OOM",
+    )
+    parser.add_argument("--eos_id", type=int, default=2, help="eos stop token id")
+    parser.add_argument(
+        "--running_max_req_size",
+        type=int,
+        default=1000,
+        help="the max size for forward requests in the same time",
+    )
+    parser.add_argument(
+        "--tp", type=int, default=1, help="model tp parral size, the default is 1"
+    )
+    parser.add_argument(
+        "--max_req_input_len",
+        type=int,
+        default=None,
+        help="the max value for req input tokens num. If None, it will be derived from the config.",
+    )
+    parser.add_argument(
+        "--max_req_total_len",
+        type=int,
+        default=None,
+        help="the max value for req_input_len + req_output_len. If None, it will be derived from the config.",
+    )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default=[],
+        nargs="+",
+        help="""Model mode: [triton_int8kv | ppl_int8kv | ppl_fp16 | triton_flashdecoding
+                        | triton_gqa_attention | triton_gqa_flashdecoding]
+                        [triton_int8weight | triton_int4weight | lmdeploy_int4weight | ppl_int4weight],
+                        triton_flashdecoding mode is for long context, current support llama llama2 qwen;
+                        triton_gqa_attention and triton_gqa_flashdecoding is fast kernel for model which use GQA;
+                        triton_int8kv mode use int8 to store kv cache, can increase token capacity, use triton kernel;
+                        ppl_int8kv mode use int8 to store kv cache, and use ppl fast kernel;
+                        ppl_fp16 mode use ppl fast fp16 decode attention kernel;
+                        triton_int8weight and triton_int4weight and lmdeploy_int4weight or ppl_int4weight mode use int8 and int4 to store weights;
+                        you need to read source code to make sure the supported detail mode for all models""",
+    )
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        help="Whether or not to allow for custom models defined on the Hub in their own modeling files.",
+    )
+    parser.add_argument(
+        "--disable_log_stats",
+        action="store_true",
+        help="disable logging throughput stats.",
+    )
+    parser.add_argument(
+        "--log_stats_interval",
+        type=int,
+        default=10,
+        help="log stats interval in second.",
+    )
+
+    parser.add_argument(
+        "--router_token_ratio",
+        type=float,
+        default=0.0,
+        help="token ratio to control router dispatch",
+    )
+    parser.add_argument(
+        "--router_max_new_token_len",
+        type=int,
+        default=1024,
+        help="the request max new token len for router",
+    )
+
+    parser.add_argument(
+        "--no_skipping_special_tokens",
+        action="store_true",
+        help="whether to skip special tokens when decoding",
+    )
+    parser.add_argument(
+        "--no_spaces_between_special_tokens",
+        action="store_true",
+        help="whether to add spaces between special tokens when decoding",
+    )
+
+    parser.add_argument(
+        "--splitfuse_mode", action="store_true", help="use splitfuse mode"
+    )
+    parser.add_argument(
+        "--splitfuse_block_size", type=int, default=256, help="splitfuse block size"
+    )
+    parser.add_argument(
+        "--prompt_cache_strs",
+        type=str,
+        default=[],
+        nargs="+",
+        help="""prompt cache strs""",
+    )
+    parser.add_argument(
+        "--cache_capacity",
+        type=int,
+        default=200,
+        help="cache server capacity for multimodal resources",
+    )
+    parser.add_argument(
+        "--cache_reserved_ratio",
+        type=float,
+        default=0.5,
+        help="cache server reserved capacity ratio after clear",
+    )
+    parser.add_argument(
+        "--return_all_prompt_logprobs",
+        action="store_true",
+        help="return all prompt tokens logprobs",
+    )
+    parser.add_argument(
+        "--long_truncation_mode",
+        type=str,
+        choices=[None, "head", "center"],
+        default=None,
+        help="""use to select the handle way when input token len > max_req_input_len.
+                        None : raise Exception
+                        head : remove some head tokens to make input token len <= max_req_input_len
+                        center : remove some tokens in center loc to make input token len <= max_req_input_len""",
+    )
+
+    args = parser.parse_args()
+
+    # 非splitfuse 模式，不支持 prompt cache 特性
+    if not args.splitfuse_mode:
+        assert len(args.prompt_cache_strs) == 0
+
+    model_config = AutoConfig.from_pretrained(args.model_dir)
+    context_length = get_context_length(model_config)
+
+    if args.max_req_input_len is None:
+        args.max_req_input_len = context_length - 1
+    if args.max_req_total_len is None:
+        args.max_req_total_len = context_length
+
+    assert args.max_req_input_len < args.max_req_total_len
+    assert args.max_req_total_len <= args.max_total_token_num
+
+    if not args.splitfuse_mode:
+        # 普通模式下
+        if args.batch_max_tokens is None:
+            batch_max_tokens = int(1 / 6 * args.max_total_token_num)
+            batch_max_tokens = max(batch_max_tokens, args.max_req_total_len)
+            args.batch_max_tokens = batch_max_tokens
+        else:
+            assert (
+                args.batch_max_tokens >= args.max_req_total_len
+            ), "batch_max_tokens must >= max_req_total_len"
+    else:
+        # splitfuse 模式下
+        # assert args.batch_max_tokens is not None, "need to set by yourself"
+        if args.batch_max_tokens is None:
+            batch_max_tokens = int(1 / 6 * args.max_total_token_num)
+            batch_max_tokens = max(batch_max_tokens, args.splitfuse_block_size)
+            args.batch_max_tokens = batch_max_tokens
+
+    can_use_ports = alloc_can_use_network_port(num=6 + args.tp)
+
+    assert can_use_ports is not None, "Can not alloc enough free ports."
+    (
+        router_port,
+        detokenization_port,
+        httpserver_port,
+        visual_port,
+        cache_port,
+        nccl_port,
+    ) = can_use_ports[0:6]
+    args.nccl_port = nccl_port
+    model_rpc_ports = can_use_ports[6:]
+
+    global httpserver_manager
+    httpserver_manager = HttpServerManager(
+        args,
+        router_port=router_port,
+        cache_port=cache_port,
+        visual_port=visual_port,
+        httpserver_port=httpserver_port,
+        enable_multimodal=False,
+    )
+
+    start_submodule_processes(
+        start_funcs=[start_router_process, start_detokenization_process],
+        start_args=[
+            (args, router_port, detokenization_port, model_rpc_ports),
+            (args, detokenization_port, httpserver_port),
+        ],
+    )
+    worker = LightLLMWorker(
+        args.controller_address,
+        args.worker_address,
+        worker_id,
+        args.model_dir,
+        args.model_names,
+        args.limit_worker_concurrency,
+        args.no_register,
+        args.conv_template,
+        httpserver_manager.tokenizer,
+        context_length,
+    )
+
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

mlx_worker.py

@@ -0,0 +1,288 @@
+"""
+A model worker using Apple MLX
+
+https://github.com/ml-explore/mlx-examples/tree/main/llms
+
+Code based on vllm_worker https://github.com/lm-sys/FastChat/blob/main/fastchat/serve/vllm_worker.py
+
+You must install MLX python:
+
+pip install mlx-lm
+"""
+
+import argparse
+import asyncio
+import atexit
+import json
+from typing import List
+import uuid
+
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.concurrency import run_in_threadpool
+from fastapi.responses import StreamingResponse, JSONResponse
+import uvicorn
+
+from fastchat.serve.base_model_worker import BaseModelWorker
+from fastchat.serve.model_worker import (
+    logger,
+    worker_id,
+)
+from fastchat.utils import get_context_length, is_partial_stop
+
+import mlx.core as mx
+from mlx_lm import load, generate
+from mlx_lm.utils import generate_step
+
+app = FastAPI()
+
+
+class MLXWorker(BaseModelWorker):
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        no_register: bool,
+        llm_engine: "MLX",
+        conv_template: str,
+    ):
+        super().__init__(
+            controller_addr,
+            worker_addr,
+            worker_id,
+            model_path,
+            model_names,
+            limit_worker_concurrency,
+            conv_template,
+        )
+
+        logger.info(
+            f"Loading the model {self.model_names} on worker {worker_id}, worker type: MLX worker..."
+        )
+
+        self.model_name = model_path
+        self.mlx_model, self.mlx_tokenizer = load(model_path)
+
+        self.tokenizer = self.mlx_tokenizer
+        # self.context_len = get_context_length(
+        #     llm_engine.engine.model_config.hf_config)
+        self.context_len = 2048  # hard code for now -- not sure how to get in MLX
+
+        if not no_register:
+            self.init_heart_beat()
+
+    async def generate_stream(self, params):
+        self.call_ct += 1
+
+        context = params.pop("prompt")
+        request_id = params.pop("request_id")
+        temperature = float(params.get("temperature", 1.0))
+        top_p = float(params.get("top_p", 1.0))
+        top_k = params.get("top_k", -1.0)
+        presence_penalty = float(params.get("presence_penalty", 0.0))
+        frequency_penalty = float(params.get("frequency_penalty", 0.0))
+        max_new_tokens = params.get("max_new_tokens", 256)
+        stop_str = params.get("stop", None)
+        stop_token_ids = params.get("stop_token_ids", None) or []
+        if self.tokenizer.eos_token_id is not None:
+            stop_token_ids.append(self.tokenizer.eos_token_id)
+        echo = params.get("echo", True)
+        use_beam_search = params.get("use_beam_search", False)
+        best_of = params.get("best_of", None)
+
+        # Handle stop_str
+        stop = set()
+        if isinstance(stop_str, str) and stop_str != "":
+            stop.add(stop_str)
+        elif isinstance(stop_str, list) and stop_str != []:
+            stop.update(stop_str)
+
+        for tid in stop_token_ids:
+            if tid is not None:
+                s = self.tokenizer.decode(tid)
+                if s != "":
+                    stop.add(s)
+
+        print("Stop patterns: ", stop)
+
+        top_p = max(top_p, 1e-5)
+        if temperature <= 1e-5:
+            top_p = 1.0
+
+        tokens = []
+        skip = 0
+
+        context_mlx = mx.array(self.tokenizer.encode(context))
+
+        finish_reason = "length"
+
+        iterator = await run_in_threadpool(
+            generate_step, context_mlx, self.mlx_model, temperature
+        )
+
+        for i in range(max_new_tokens):
+            (token, _) = await run_in_threadpool(next, iterator)
+            if token == self.mlx_tokenizer.eos_token_id:
+                finish_reason = "stop"
+                break
+            tokens.append(token.item())
+            tokens_decoded = self.mlx_tokenizer.decode(tokens)
+            last_token_decoded = self.mlx_tokenizer.decode([token.item()])
+            skip = len(tokens_decoded)
+
+            partial_stop = any(is_partial_stop(tokens_decoded, i) for i in stop)
+
+            if partial_stop:
+                finish_reason = "stop"
+                break
+
+            ret = {
+                "text": tokens_decoded,
+                "error_code": 0,
+                "usage": {
+                    "prompt_tokens": len(context),
+                    "completion_tokens": len(tokens),
+                    "total_tokens": len(context) + len(tokens),
+                },
+                "cumulative_logprob": [],
+                "finish_reason": None,  # hard code for now
+            }
+            # print(ret)
+            yield (json.dumps(ret) + "\0").encode()
+        ret = {
+            "text": self.mlx_tokenizer.decode(tokens),
+            "error_code": 0,
+            "usage": {},
+            "cumulative_logprob": [],
+            "finish_reason": finish_reason,
+        }
+        yield (json.dumps(obj={**ret, **{"finish_reason": None}}) + "\0").encode()
+        yield (json.dumps(ret) + "\0").encode()
+
+    async def generate(self, params):
+        async for x in self.generate_stream(params):
+            pass
+        return json.loads(x[:-1].decode())
+
+
+def release_worker_semaphore():
+    worker.semaphore.release()
+
+
+def acquire_worker_semaphore():
+    if worker.semaphore is None:
+        worker.semaphore = asyncio.Semaphore(worker.limit_worker_concurrency)
+    return worker.semaphore.acquire()
+
+
+def create_background_tasks(request_id):
+    async def abort_request() -> None:
+        print("trying to abort but not implemented")
+
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(release_worker_semaphore)
+    background_tasks.add_task(abort_request)
+    return background_tasks
+
+
+@app.post("/worker_generate_stream")
+async def api_generate_stream(request: Request):
+    params = await request.json()
+    await acquire_worker_semaphore()
+    request_id = uuid.uuid4()
+    params["request_id"] = str(request_id)
+    generator = worker.generate_stream(params)
+    background_tasks = create_background_tasks(request_id)
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_generate")
+async def api_generate(request: Request):
+    params = await request.json()
+    await acquire_worker_semaphore()
+    request_id = uuid.uuid4()
+    params["request_id"] = str(request_id)
+    output = await worker.generate(params)
+    release_worker_semaphore()
+    # await engine.abort(request_id)
+    print("Trying to abort but not implemented")
+    return JSONResponse(output)
+
+
+@app.post("/worker_get_status")
+async def api_get_status(request: Request):
+    return worker.get_status()
+
+
+@app.post("/count_token")
+async def api_count_token(request: Request):
+    params = await request.json()
+    return worker.count_token(params)
+
+
+@app.post("/worker_get_conv_template")
+async def api_get_conv(request: Request):
+    return worker.get_conv_template()
+
+
+@app.post("/model_details")
+async def api_model_details(request: Request):
+    return {"context_length": worker.context_len}
+
+
+worker = None
+
+
+def cleanup_at_exit():
+    global worker
+    print("Cleaning up...")
+    del worker
+
+
+atexit.register(cleanup_at_exit)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+    parser.add_argument(
+        "--controller-address", type=str, default="http://localhost:21001"
+    )
+    parser.add_argument("--model-path", type=str, default="microsoft/phi-2")
+    parser.add_argument(
+        "--model-names",
+        type=lambda s: s.split(","),
+        help="Optional display comma separated names",
+    )
+    parser.add_argument(
+        "--conv-template", type=str, default=None, help="Conversation prompt template."
+    )
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_false",
+        default=True,
+        help="Trust remote code (e.g., from HuggingFace) when"
+        "downloading the model and tokenizer.",
+    )
+
+    args, unknown = parser.parse_known_args()
+
+    if args.model_path:
+        args.model = args.model_path
+
+    worker = MLXWorker(
+        args.controller_address,
+        args.worker_address,
+        worker_id,
+        args.model_path,
+        args.model_names,
+        1024,
+        False,
+        "MLX",
+        args.conv_template,
+    )
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

model_worker.py

@@ -0,0 +1,425 @@
+"""
+A model worker that executes the model.
+"""
+import argparse
+import base64
+import gc
+import json
+import os
+from typing import List, Optional
+import uuid
+
+import torch
+import torch.nn.functional as F
+from transformers import set_seed
+import uvicorn
+
+from fastchat.constants import ErrorCode, SERVER_ERROR_MSG
+from fastchat.model.model_adapter import (
+    load_model,
+    add_model_args,
+    get_generate_stream_function,
+)
+from fastchat.modules.awq import AWQConfig
+from fastchat.modules.exllama import ExllamaConfig
+from fastchat.modules.xfastertransformer import XftConfig
+from fastchat.modules.gptq import GptqConfig
+from fastchat.serve.base_model_worker import BaseModelWorker, app
+from fastchat.utils import (
+    build_logger,
+    get_context_length,
+    str_to_torch_dtype,
+)
+
+worker_id = str(uuid.uuid4())[:8]
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
+
+
+class ModelWorker(BaseModelWorker):
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        no_register: bool,
+        device: str,
+        num_gpus: int,
+        max_gpu_memory: str,
+        revision: str = None,
+        dtype: Optional[torch.dtype] = None,
+        load_8bit: bool = False,
+        cpu_offloading: bool = False,
+        gptq_config: Optional[GptqConfig] = None,
+        awq_config: Optional[AWQConfig] = None,
+        exllama_config: Optional[ExllamaConfig] = None,
+        xft_config: Optional[XftConfig] = None,
+        stream_interval: int = 2,
+        conv_template: Optional[str] = None,
+        embed_in_truncate: bool = False,
+        seed: Optional[int] = None,
+        debug: bool = False,
+        **kwargs,
+    ):
+        super().__init__(
+            controller_addr,
+            worker_addr,
+            worker_id,
+            model_path,
+            model_names,
+            limit_worker_concurrency,
+            conv_template=conv_template,
+        )
+
+        logger.info(f"Loading the model {self.model_names} on worker {worker_id} ...")
+        self.model, self.tokenizer = load_model(
+            model_path,
+            revision=revision,
+            device=device,
+            num_gpus=num_gpus,
+            max_gpu_memory=max_gpu_memory,
+            dtype=dtype,
+            load_8bit=load_8bit,
+            cpu_offloading=cpu_offloading,
+            gptq_config=gptq_config,
+            awq_config=awq_config,
+            exllama_config=exllama_config,
+            xft_config=xft_config,
+            debug=debug,
+        )
+        self.device = device
+        if self.tokenizer.pad_token == None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.context_len = get_context_length(self.model.config)
+        self.generate_stream_func = get_generate_stream_function(self.model, model_path)
+        self.stream_interval = stream_interval
+        self.embed_in_truncate = embed_in_truncate
+        self.seed = seed
+
+        if not no_register:
+            self.init_heart_beat()
+
+    def generate_stream_gate(self, params):
+        if self.device == "npu":
+            import torch_npu
+
+            torch_npu.npu.set_device("npu:0")
+        self.call_ct += 1
+
+        try:
+            if self.seed is not None:
+                set_seed(self.seed)
+            for output in self.generate_stream_func(
+                self.model,
+                self.tokenizer,
+                params,
+                self.device,
+                self.context_len,
+                self.stream_interval,
+            ):
+                ret = {
+                    "text": output["text"],
+                    "error_code": 0,
+                }
+                if "usage" in output:
+                    ret["usage"] = output["usage"]
+                if "finish_reason" in output:
+                    ret["finish_reason"] = output["finish_reason"]
+                if "logprobs" in output:
+                    ret["logprobs"] = output["logprobs"]
+                yield json.dumps(ret).encode() + b"\0"
+        except torch.cuda.OutOfMemoryError as e:
+            ret = {
+                "text": f"{SERVER_ERROR_MSG}\n\n({e})",
+                "error_code": ErrorCode.CUDA_OUT_OF_MEMORY,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except (ValueError, RuntimeError) as e:
+            ret = {
+                "text": f"{SERVER_ERROR_MSG}\n\n({e})",
+                "error_code": ErrorCode.INTERNAL_ERROR,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+    def generate_gate(self, params):
+        for x in self.generate_stream_gate(params):
+            pass
+        return json.loads(x[:-1].decode())
+
+    def __process_embed_chunk(self, input_ids, attention_mask, **model_type_dict):
+        if model_type_dict.get("is_bert"):
+            model_output = self.model(input_ids)
+            if model_type_dict.get("is_robert"):
+                data = model_output.last_hidden_state
+            else:
+                data = model_output[0]
+        elif model_type_dict.get("is_t5"):
+            model_output = self.model(input_ids, decoder_input_ids=input_ids)
+            data = model_output.encoder_last_hidden_state
+        else:
+            model_output = self.model(input_ids, output_hidden_states=True)
+            if model_type_dict.get("is_chatglm"):
+                data = model_output.hidden_states[-1].transpose(0, 1)
+            else:
+                data = model_output.hidden_states[-1]
+
+        if hasattr(self.model, "use_cls_pooling") and self.model.use_cls_pooling:
+            sum_embeddings = data[:, 0]
+        else:
+            mask = attention_mask.unsqueeze(-1).expand(data.size()).float()
+            masked_embeddings = data * mask
+            sum_embeddings = torch.sum(masked_embeddings, dim=1)
+        token_num = torch.sum(attention_mask).item()
+
+        return sum_embeddings, token_num
+
+    def __encode_base64(self, embeddings: torch.Tensor) -> List[str]:
+        embeddings = embeddings.cpu()
+        return [
+            base64.b64encode(e.numpy().tobytes()).decode("utf-8") for e in embeddings
+        ]
+
+    @torch.inference_mode()
+    def get_embeddings(self, params):
+        self.call_ct += 1
+
+        try:
+            tokenizer = self.tokenizer
+            ret = {"embedding": [], "token_num": 0}
+
+            model_type_dict = {
+                "is_llama": "llama" in str(type(self.model)),
+                "is_t5": "t5" in str(type(self.model)),
+                "is_chatglm": "chatglm" in str(type(self.model)),
+                "is_bert": "bert" in str(type(self.model)),
+                "is_robert": "robert" in str(type(self.model)),
+            }
+
+            if self.embed_in_truncate:
+                encoding = tokenizer.batch_encode_plus(
+                    params["input"],
+                    padding=True,
+                    truncation="longest_first",
+                    return_tensors="pt",
+                    max_length=self.context_len,
+                )
+            else:
+                encoding = tokenizer.batch_encode_plus(
+                    params["input"], padding=True, return_tensors="pt"
+                )
+            input_ids = encoding["input_ids"].to(self.device)
+            attention_mask = input_ids != tokenizer.pad_token_id
+
+            base64_encode = params.get("encoding_format", None)
+
+            if self.embed_in_truncate:
+                embedding, token_num = self.__process_embed_chunk(
+                    input_ids, attention_mask, **model_type_dict
+                )
+                if (
+                    not hasattr(self.model, "use_cls_pooling")
+                    or not self.model.use_cls_pooling
+                ):
+                    embedding = embedding / token_num
+                normalized_embeddings = F.normalize(embedding, p=2, dim=1)
+                ret["token_num"] = token_num
+            else:
+                all_embeddings = []
+                all_token_num = 0
+                for i in range(0, input_ids.size(1), self.context_len):
+                    chunk_input_ids = input_ids[:, i : i + self.context_len]
+                    chunk_attention_mask = attention_mask[:, i : i + self.context_len]
+
+                    # add cls token and mask to get cls embedding
+                    if (
+                        hasattr(self.model, "use_cls_pooling")
+                        and self.model.use_cls_pooling
+                    ):
+                        cls_tokens = (
+                            torch.zeros(
+                                (chunk_input_ids.size(0), 1),
+                                dtype=chunk_input_ids.dtype,
+                                device=chunk_input_ids.device,
+                            )
+                            + tokenizer.cls_token_id
+                        )
+                        chunk_input_ids = torch.cat(
+                            [cls_tokens, chunk_input_ids], dim=-1
+                        )
+                        mask = torch.ones(
+                            (chunk_attention_mask.size(0), 1),
+                            dtype=chunk_attention_mask.dtype,
+                            device=chunk_attention_mask.device,
+                        )
+                        chunk_attention_mask = torch.cat(
+                            [mask, chunk_attention_mask], dim=-1
+                        )
+
+                    chunk_embeddings, token_num = self.__process_embed_chunk(
+                        chunk_input_ids, chunk_attention_mask, **model_type_dict
+                    )
+                    if (
+                        hasattr(self.model, "use_cls_pooling")
+                        and self.model.use_cls_pooling
+                    ):
+                        all_embeddings.append(chunk_embeddings * token_num)
+                    else:
+                        all_embeddings.append(chunk_embeddings)
+                    all_token_num += token_num
+
+                all_embeddings_tensor = torch.stack(all_embeddings)
+                embedding = torch.sum(all_embeddings_tensor, dim=0) / all_token_num
+                normalized_embeddings = F.normalize(embedding, p=2, dim=1)
+
+                ret["token_num"] = all_token_num
+
+            if base64_encode == "base64":
+                out_embeddings = self.__encode_base64(normalized_embeddings)
+            else:
+                out_embeddings = normalized_embeddings.tolist()
+            ret["embedding"] = out_embeddings
+
+            gc.collect()
+            torch.cuda.empty_cache()
+            if self.device == "xpu":
+                torch.xpu.empty_cache()
+            if self.device == "npu":
+                torch.npu.empty_cache()
+        except torch.cuda.OutOfMemoryError as e:
+            ret = {
+                "text": f"{SERVER_ERROR_MSG}\n\n({e})",
+                "error_code": ErrorCode.CUDA_OUT_OF_MEMORY,
+            }
+        except (ValueError, RuntimeError) as e:
+            ret = {
+                "text": f"{SERVER_ERROR_MSG}\n\n({e})",
+                "error_code": ErrorCode.INTERNAL_ERROR,
+            }
+        return ret
+
+
+def create_model_worker():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+    parser.add_argument(
+        "--controller-address", type=str, default="http://localhost:21001"
+    )
+    add_model_args(parser)
+    parser.add_argument(
+        "--model-names",
+        type=lambda s: s.split(","),
+        help="Optional display comma separated names",
+    )
+    parser.add_argument(
+        "--conv-template", type=str, default=None, help="Conversation prompt template."
+    )
+    parser.add_argument("--embed-in-truncate", action="store_true")
+    parser.add_argument(
+        "--limit-worker-concurrency",
+        type=int,
+        default=5,
+        help="Limit the model concurrency to prevent OOM.",
+    )
+    parser.add_argument("--stream-interval", type=int, default=2)
+    parser.add_argument("--no-register", action="store_true")
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Overwrite the random seed for each generation.",
+    )
+    parser.add_argument(
+        "--debug", type=bool, default=False, help="Print debugging messages"
+    )
+    parser.add_argument(
+        "--ssl",
+        action="store_true",
+        required=False,
+        default=False,
+        help="Enable SSL. Requires OS Environment variables 'SSL_KEYFILE' and 'SSL_CERTFILE'.",
+    )
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    if args.gpus:
+        if len(args.gpus.split(",")) < args.num_gpus:
+            raise ValueError(
+                f"Larger --num-gpus ({args.num_gpus}) than --gpus {args.gpus}!"
+            )
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
+
+    gptq_config = GptqConfig(
+        ckpt=args.gptq_ckpt or args.model_path,
+        wbits=args.gptq_wbits,
+        groupsize=args.gptq_groupsize,
+        act_order=args.gptq_act_order,
+    )
+    awq_config = AWQConfig(
+        ckpt=args.awq_ckpt or args.model_path,
+        wbits=args.awq_wbits,
+        groupsize=args.awq_groupsize,
+    )
+    if args.enable_exllama:
+        exllama_config = ExllamaConfig(
+            max_seq_len=args.exllama_max_seq_len,
+            gpu_split=args.exllama_gpu_split,
+            cache_8bit=args.exllama_cache_8bit,
+        )
+    else:
+        exllama_config = None
+    if args.enable_xft:
+        xft_config = XftConfig(
+            max_seq_len=args.xft_max_seq_len,
+            data_type=args.xft_dtype,
+        )
+        if args.device != "cpu":
+            print("xFasterTransformer now is only support CPUs. Reset device to CPU")
+            args.device = "cpu"
+    else:
+        xft_config = None
+
+    worker = ModelWorker(
+        args.controller_address,
+        args.worker_address,
+        worker_id,
+        args.model_path,
+        args.model_names,
+        args.limit_worker_concurrency,
+        revision=args.revision,
+        no_register=args.no_register,
+        device=args.device,
+        num_gpus=args.num_gpus,
+        max_gpu_memory=args.max_gpu_memory,
+        dtype=str_to_torch_dtype(args.dtype),
+        load_8bit=args.load_8bit,
+        cpu_offloading=args.cpu_offloading,
+        gptq_config=gptq_config,
+        awq_config=awq_config,
+        exllama_config=exllama_config,
+        xft_config=xft_config,
+        stream_interval=args.stream_interval,
+        conv_template=args.conv_template,
+        embed_in_truncate=args.embed_in_truncate,
+        seed=args.seed,
+        debug=args.debug,
+    )
+    return args, worker
+
+
+if __name__ == "__main__":
+    args, worker = create_model_worker()
+    if args.ssl:
+        uvicorn.run(
+            app,
+            host=args.host,
+            port=args.port,
+            log_level="info",
+            ssl_keyfile=os.environ["SSL_KEYFILE"],
+            ssl_certfile=os.environ["SSL_CERTFILE"],
+        )
+    else:
+        uvicorn.run(app, host=args.host, port=args.port, log_level="info")