Upload folder using huggingface_hub

README.md

@@ -1,12 +1,144 @@
 ---
-title: Glm4
-emoji: 🌖
-colorFrom: green
-colorTo: pink
+title: glm4
+app_file: trans_web_demo.py
 sdk: gradio
 sdk_version: 4.37.2
-app_file: app.py
-pinned: false
 ---
+# Basic Demo
+
+Read this in [English](README_en.md).
+
+本 demo 中，你将体验到如何使用 GLM-4-9B 开源模型进行基本的任务。
+
+请严格按照文档的步骤进行操作，以避免不必要的错误。
+
+## 设备和依赖检查
+
+### 相关推理测试数据
+
+**本文档的数据均在以下硬件环境测试,实际运行环境需求和运行占用的显存略有不同，请以实际运行环境为准。**
+
+测试硬件信息:
+
++ OS: Ubuntu 22.04
++ Memory: 512GB
++ Python: 3.10.12 （推荐） / 3.12.3 均已测试
++ CUDA Version:  12.3
++ GPU Driver: 535.104.05
++ GPU: NVIDIA A100-SXM4-80GB * 8
+
+相关推理的压力测试数据如下：
+
+**所有测试均在单张GPU上进行测试,所有显存消耗都按照峰值左右进行测算**
+
+#### GLM-4-9B-Chat
+
+| 精度   | 显存占用  | Prefilling | Decode Speed  | Remarks      |
+|------|-------|------------|---------------|--------------|
+| BF16 | 19 GB | 0.2s       | 27.8 tokens/s | 输入长度为 1000   |
+| BF16 | 21 GB | 0.8s       | 31.8 tokens/s | 输入长度为 8000   |
+| BF16 | 28 GB | 4.3s       | 14.4 tokens/s | 输入长度为 32000  |
+| BF16 | 58 GB | 38.1s      | 3.4  tokens/s | 输入长度为 128000 |
+
+| 精度   | 显存占用  | Prefilling | Decode Speed  | Remarks     |
+|------|-------|------------|---------------|-------------|
+| INT4 | 8 GB  | 0.2s       | 23.3 tokens/s | 输入长度为 1000  |
+| INT4 | 10 GB | 0.8s       | 23.4 tokens/s | 输入长度为 8000  |
+| INT4 | 17 GB | 4.3s       | 14.6 tokens/s | 输入长度为 32000 |
+
+#### GLM-4-9B-Chat-1M
+
+| 精度   | 显存占用  | Prefilling | Decode Speed | Remarks      |
+|------|-------|------------|--------------|--------------|
+| BF16 | 75 GB | 98.4s      | 2.3 tokens/s | 输入长度为 200000 |
+
+如果您的输入超过200K，我们建议您使用vLLM后端进行多卡推理，以获得更好的性能。
+
+#### GLM-4V-9B
+
+| 精度   | 显存占用  | Prefilling | Decode Speed  | Remarks    |
+|------|-------|------------|---------------|------------|
+| BF16 | 28 GB | 0.1s       | 33.4 tokens/s | 输入长度为 1000 |
+| BF16 | 33 GB | 0.7s       | 39.2 tokens/s | 输入长度为 8000 |
+
+| 精度   | 显存占用  | Prefilling | Decode Speed  | Remarks    |
+|------|-------|------------|---------------|------------|
+| INT4 | 10 GB | 0.1s       | 28.7 tokens/s | 输入长度为 1000 |
+| INT4 | 15 GB | 0.8s       | 24.2 tokens/s | 输入长度为 8000 |
+
+### 最低硬件要求
+
+如果您希望运行官方提供的最基础代码 (transformers 后端) 您需要：
+
++ Python >= 3.10
++ 内存不少于 32 GB
+
+如果您希望运行官方提供的本文件夹的所有代码，您还需要：
+
++ Linux 操作系统 (Debian 系列最佳)
++ 大于 8GB 显存的，支持 CUDA 或者 ROCM 并且支持 `BF16` 推理的 GPU 设备。(`FP16` 精度无法训练，推理有小概率出现问题)
+
+安装依赖
+
+```shell
+pip install -r requirements.txt
+```
+
+## 基础功能调用
+
+**除非特殊说明，本文件夹所有 demo 并不支持 Function Call 和 All Tools 等进阶用法**
+
+### 使用 transformers 后端代码
+
++ 使用命令行与 GLM-4-9B 模型进行对话。
+
+```shell
+python trans_cli_demo.py # GLM-4-9B-Chat
+python trans_cli_vision_demo.py # GLM-4V-9B
+```
+
++ 使用 Gradio 网页端与 GLM-4-9B 模型进行对话。
+
+```shell
+python trans_web_demo.py  # GLM-4-9B-Chat
+python trans_web_vision_demo.py # GLM-4V-9B
+```
+
++ 使用 Batch 推理。
+
+```shell
+python trans_batch_demo.py
+```
+
+### 使用 vLLM 后端代码
+
++ 使用命令行与 GLM-4-9B-Chat 模型进行对话。
+
+```shell
+python vllm_cli_demo.py
+```
+
++ 自行构建服务端，并使用 `OpenAI API` 的请求格式与 GLM-4-9B-Chat 模型进行对话。本 demo 支持 Function Call 和 All Tools功能。
+
+启动服务端：
+
+```shell
+python openai_api_server.py
+```
+
+客户端请求：
+
+```shell
+python openai_api_request.py
+```
+
+## 压力测试
+
+用户可以在自己的设备上使用本代码测试模型在 transformers后端的生成速度:
+
+```shell
+python trans_stress_test.py
+```
+
+
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

README_en.md

@@ -0,0 +1,141 @@
+# Basic Demo
+
+In this demo, you will experience how to use the GLM-4-9B open source model to perform basic tasks.
+
+Please follow the steps in the document strictly to avoid unnecessary errors.
+
+## Device and dependency check
+
+### Related inference test data
+
+**The data in this document are tested in the following hardware environment. The actual operating environment
+requirements and the GPU memory occupied by the operation are slightly different. Please refer to the actual operating
+environment.**
+
+Test hardware information:
+
++ OS: Ubuntu 22.04
++ Memory: 512GB
++ Python: 3.10.12 (recommend) / 3.12.3 have been tested
++ CUDA Version: 12.3
++ GPU Driver: 535.104.05
++ GPU: NVIDIA A100-SXM4-80GB * 8
+
+The stress test data of relevant inference are as follows:
+
+**All tests are performed on a single GPU, and all GPU memory consumption is calculated based on the peak value**
+
+#
+
+### GLM-4-9B-Chat
+
+| Dtype | GPU Memory | Prefilling | Decode Speed  | Remarks                |
+|-------|------------|------------|---------------|------------------------|
+| BF16  | 19 GB      | 0.2s       | 27.8 tokens/s | Input length is 1000   |
+| BF16  | 21 GB      | 0.8s       | 31.8 tokens/s | Input length is 8000   |
+| BF16  | 28 GB      | 4.3s       | 14.4 tokens/s | Input length is 32000  |
+| BF16  | 58 GB      | 38.1s      | 3.4  tokens/s | Input length is 128000 |
+
+| Dtype | GPU Memory | Prefilling | Decode Speed  | Remarks               |
+|-------|------------|------------|---------------|-----------------------|
+| INT4  | 8 GB       | 0.2s       | 23.3 tokens/s | Input length is 1000  |
+| INT4  | 10 GB      | 0.8s       | 23.4 tokens/s | Input length is 8000  |
+| INT4  | 17 GB      | 4.3s       | 14.6 tokens/s | Input length is 32000 |
+
+### GLM-4-9B-Chat-1M
+
+| Dtype | GPU Memory | Prefilling | Decode Speed     | Remarks                |
+|-------|------------|------------|------------------|------------------------|
+| BF16  | 74497MiB   | 98.4s      | 2.3653  tokens/s | Input length is 200000 |
+
+If your input exceeds 200K, we recommend that you use the vLLM backend with multi gpus for inference to get better
+performance.
+
+#### GLM-4V-9B
+
+| Dtype | GPU Memory | Prefilling | Decode Speed  | Remarks              |
+|-------|------------|------------|---------------|----------------------|
+| BF16  | 28 GB      | 0.1s       | 33.4 tokens/s | Input length is 1000 |
+| BF16  | 33 GB      | 0.7s       | 39.2 tokens/s | Input length is 8000 |
+
+| Dtype | GPU Memory | Prefilling | Decode Speed  | Remarks              |
+|-------|------------|------------|---------------|----------------------|
+| INT4  | 10 GB      | 0.1s       | 28.7 tokens/s | Input length is 1000 |
+| INT4  | 15 GB      | 0.8s       | 24.2 tokens/s | Input length is 8000 |
+
+### Minimum hardware requirements
+
+If you want to run the most basic code provided by the official (transformers backend) you need:
+
++ Python >= 3.10
++ Memory of at least 32 GB
+
+If you want to run all the codes in this folder provided by the official, you also need:
+
++ Linux operating system (Debian series is best)
++ GPU device with more than 8GB GPU memory, supporting CUDA or ROCM and supporting `BF16` reasoning (`FP16` precision
+  cannot be finetuned, and there is a small probability of problems in infering)
+
+Install dependencies
+
+```shell
+pip install -r requirements.txt
+```
+
+## Basic function calls
+
+**Unless otherwise specified, all demos in this folder do not support advanced usage such as Function Call and All Tools
+**
+
+### Use transformers backend code
+
++ Use the command line to communicate with the GLM-4-9B model.
+
+```shell
+python trans_cli_demo.py # GLM-4-9B-Chat
+python trans_cli_vision_demo.py # GLM-4V-9B
+```
+
++ Use the Gradio web client to communicate with the  GLM-4-9B model.
+
+```shell
+python trans_web_demo.py  # GLM-4-9B-Chat
+python trans_web_vision_demo.py # GLM-4V-9B
+```
+
++ Use Batch inference.
+
+```shell
+python trans_batch_demo.py
+```
+
+### Use vLLM backend code
+
++ Use the command line to communicate with the GLM-4-9B-Chat model.
+
+```shell
+python vllm_cli_demo.py
+```
+
++ Build the server by yourself and use the request format of `OpenAI API` to communicate with the glm-4-9b model. This
+  demo supports Function Call and All Tools functions.
+
+Start the server:
+
+```shell
+python openai_api_server.py
+```
+
+Client request:
+
+```shell
+python openai_api_request.py
+```
+
+## Stress test
+
+Users can use this code to test the generation speed of the model on the transformers backend on their own devices:
+
+```shell
+python trans_stress_test.py
+```

openai_api_request.py

@@ -0,0 +1,127 @@
+"""
+This script creates a OpenAI Request demo for the glm-4-9b model, just Use OpenAI API to interact with the model.
+"""
+
+from openai import OpenAI
+
+base_url = "http://127.0.0.1:8000/v1/"
+client = OpenAI(api_key="EMPTY", base_url=base_url)
+
+
+def function_chat(use_stream=False):
+    messages = [
+        {
+            "role": "user", "content": "What's the Celsius temperature in San Francisco?"
+        },
+
+        # Give Observations
+        # {
+        #     "role": "assistant",
+        #         "content": None,
+        #         "function_call": None,
+        #         "tool_calls": [
+        #             {
+        #                 "id": "call_1717912616815",
+        #                 "function": {
+        #                     "name": "get_current_weather",
+        #                     "arguments": "{\"location\": \"San Francisco, CA\", \"format\": \"celsius\"}"
+        #                 },
+        #                 "type": "function"
+        #             }
+        #         ]
+        # },
+        # {
+        #     "tool_call_id": "call_1717912616815",
+        #     "role": "tool",
+        #     "name": "get_current_weather",
+        #     "content": "23°C",
+        # }
+    ]
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_current_weather",
+                "description": "Get the current weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g. San Francisco, CA",
+                        },
+                        "format": {
+                            "type": "string",
+                            "enum": ["celsius", "fahrenheit"],
+                            "description": "The temperature unit to use. Infer this from the users location.",
+                        },
+                    },
+                    "required": ["location", "format"],
+                },
+            }
+        },
+    ]
+
+    # All Tools: CogView
+    # messages = [{"role": "user", "content": "帮我画一张天空的画画吧"}]
+    # tools = [{"type": "cogview"}]
+
+    # All Tools: Searching
+    # messages = [{"role": "user", "content": "今天黄金的价格"}]
+    # tools = [{"type": "simple_browser"}]
+
+    response = client.chat.completions.create(
+        model="glm-4",
+        messages=messages,
+        tools=tools,
+        stream=use_stream,
+        max_tokens=256,
+        temperature=0.9,
+        presence_penalty=1.2,
+        top_p=0.1,
+        tool_choice="auto"
+    )
+    if response:
+        if use_stream:
+            for chunk in response:
+                print(chunk)
+        else:
+            print(response)
+    else:
+        print("Error:", response.status_code)
+
+
+def simple_chat(use_stream=False):
+    messages = [
+        {
+            "role": "system",
+            "content": "请在你输出的时候都带上“喵喵喵”三个字，放在开头。",
+        },
+        {
+            "role": "user",
+            "content": "你是谁"
+        }
+    ]
+    response = client.chat.completions.create(
+        model="glm-4",
+        messages=messages,
+        stream=use_stream,
+        max_tokens=256,
+        temperature=0.4,
+        presence_penalty=1.2,
+        top_p=0.8,
+    )
+    if response:
+        if use_stream:
+            for chunk in response:
+                print(chunk)
+        else:
+            print(response)
+    else:
+        print("Error:", response.status_code)
+
+
+if __name__ == "__main__":
+    # simple_chat(use_stream=False)
+    function_chat(use_stream=False)
+

openai_api_server.py

@@ -0,0 +1,635 @@
+import time
+from asyncio.log import logger
+import re
+import uvicorn
+import gc
+import json
+import torch
+import random
+import string
+
+from vllm import SamplingParams, AsyncEngineArgs, AsyncLLMEngine
+from fastapi import FastAPI, HTTPException, Response
+from fastapi.middleware.cors import CORSMiddleware
+from contextlib import asynccontextmanager
+from typing import List, Literal, Optional, Union
+from pydantic import BaseModel, Field
+from transformers import AutoTokenizer, LogitsProcessor
+from sse_starlette.sse import EventSourceResponse
+
+EventSourceResponse.DEFAULT_PING_INTERVAL = 1000
+import os
+
+MODEL_PATH = os.environ.get('MODEL_PATH', 'THUDM/glm-4-9b-chat')
+MAX_MODEL_LENGTH = 8192
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    yield
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def generate_id(prefix: str, k=29) -> str:
+    suffix = ''.join(random.choices(string.ascii_letters + string.digits, k=k))
+    return f"{prefix}{suffix}"
+
+
+class ModelCard(BaseModel):
+    id: str = ""
+    object: str = "model"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    owned_by: str = "owner"
+    root: Optional[str] = None
+    parent: Optional[str] = None
+    permission: Optional[list] = None
+
+
+class ModelList(BaseModel):
+    object: str = "list"
+    data: List[ModelCard] = ["glm-4"]
+
+
+class FunctionCall(BaseModel):
+    name: Optional[str] = None
+    arguments: Optional[str] = None
+
+
+class ChoiceDeltaToolCallFunction(BaseModel):
+    name: Optional[str] = None
+    arguments: Optional[str] = None
+
+
+class UsageInfo(BaseModel):
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+    completion_tokens: Optional[int] = 0
+
+
+class ChatCompletionMessageToolCall(BaseModel):
+    index: Optional[int] = 0
+    id: Optional[str] = None
+    function: FunctionCall
+    type: Optional[Literal["function"]] = 'function'
+
+
+class ChatMessage(BaseModel):
+    # “function” 字段解释：
+    # 使用较老的OpenAI API版本需要注意在这里添加 function 字段并在 process_messages函数中添加相应角色转换逻辑为 observation
+
+    role: Literal["user", "assistant", "system", "tool"]
+    content: Optional[str] = None
+    function_call: Optional[ChoiceDeltaToolCallFunction] = None
+    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
+
+
+class DeltaMessage(BaseModel):
+    role: Optional[Literal["user", "assistant", "system"]] = None
+    content: Optional[str] = None
+    function_call: Optional[ChoiceDeltaToolCallFunction] = None
+    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
+
+
+class ChatCompletionResponseChoice(BaseModel):
+    index: int
+    message: ChatMessage
+    finish_reason: Literal["stop", "length", "tool_calls"]
+
+
+class ChatCompletionResponseStreamChoice(BaseModel):
+    delta: DeltaMessage
+    finish_reason: Optional[Literal["stop", "length", "tool_calls"]]
+    index: int
+
+
+class ChatCompletionResponse(BaseModel):
+    model: str
+    id: Optional[str] = Field(default_factory=lambda: generate_id('chatcmpl-', 29))
+    object: Literal["chat.completion", "chat.completion.chunk"]
+    choices: List[Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]]
+    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
+    system_fingerprint: Optional[str] = Field(default_factory=lambda: generate_id('fp_', 9))
+    usage: Optional[UsageInfo] = None
+
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    temperature: Optional[float] = 0.8
+    top_p: Optional[float] = 0.8
+    max_tokens: Optional[int] = None
+    stream: Optional[bool] = False
+    tools: Optional[Union[dict, List[dict]]] = None
+    tool_choice: Optional[Union[str, dict]] = None
+    repetition_penalty: Optional[float] = 1.1
+
+
+class InvalidScoreLogitsProcessor(LogitsProcessor):
+    def __call__(
+            self, input_ids: torch.LongTensor, scores: torch.FloatTensor
+    ) -> torch.FloatTensor:
+        if torch.isnan(scores).any() or torch.isinf(scores).any():
+            scores.zero_()
+            scores[..., 5] = 5e4
+        return scores
+
+
+def process_response(output: str, tools: dict | List[dict] = None, use_tool: bool = False) -> Union[str, dict]:
+    lines = output.strip().split("\n")
+    arguments_json = None
+    special_tools = ["cogview", "simple_browser"]
+    tools = {tool['function']['name'] for tool in tools} if tools else {}
+
+    # 这是一个简单的工具比较函数，不能保证拦截所有非工具输出的结果，比如参数未对齐等特殊情况。
+    ##TODO 如果你希望做更多判断，可以在这里进行逻辑完善。
+
+    if len(lines) >= 2 and lines[1].startswith("{"):
+        function_name = lines[0].strip()
+        arguments = "\n".join(lines[1:]).strip()
+        if function_name in tools or function_name in special_tools:
+            try:
+                arguments_json = json.loads(arguments)
+                is_tool_call = True
+            except json.JSONDecodeError:
+                is_tool_call = function_name in special_tools
+
+            if is_tool_call and use_tool:
+                content = {
+                    "name": function_name,
+                    "arguments": json.dumps(arguments_json if isinstance(arguments_json, dict) else arguments,
+                                            ensure_ascii=False)
+                }
+                if function_name == "simple_browser":
+                    search_pattern = re.compile(r'search\("(.+?)"\s*,\s*recency_days\s*=\s*(\d+)\)')
+                    match = search_pattern.match(arguments)
+                    if match:
+                        content["arguments"] = json.dumps({
+                            "query": match.group(1),
+                            "recency_days": int(match.group(2))
+                        }, ensure_ascii=False)
+                elif function_name == "cogview":
+                    content["arguments"] = json.dumps({
+                        "prompt": arguments
+                    }, ensure_ascii=False)
+
+                return content
+    return output.strip()
+
+
+@torch.inference_mode()
+async def generate_stream_glm4(params):
+    messages = params["messages"]
+    tools = params["tools"]
+    tool_choice = params["tool_choice"]
+    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))
+    max_new_tokens = int(params.get("max_tokens", 8192))
+
+    messages = process_messages(messages, tools=tools, tool_choice=tool_choice)
+    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    params_dict = {
+        "n": 1,
+        "best_of": 1,
+        "presence_penalty": 1.0,
+        "frequency_penalty": 0.0,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": -1,
+        "repetition_penalty": repetition_penalty,
+        "use_beam_search": False,
+        "length_penalty": 1,
+        "early_stopping": False,
+        "stop_token_ids": [151329, 151336, 151338],
+        "ignore_eos": False,
+        "max_tokens": max_new_tokens,
+        "logprobs": None,
+        "prompt_logprobs": None,
+        "skip_special_tokens": True,
+    }
+    sampling_params = SamplingParams(**params_dict)
+    async for output in engine.generate(inputs=inputs, sampling_params=sampling_params, request_id=f"{time.time()}"):
+        output_len = len(output.outputs[0].token_ids)
+        input_len = len(output.prompt_token_ids)
+        ret = {
+            "text": output.outputs[0].text,
+            "usage": {
+                "prompt_tokens": input_len,
+                "completion_tokens": output_len,
+                "total_tokens": output_len + input_len
+            },
+            "finish_reason": output.outputs[0].finish_reason,
+        }
+        yield ret
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+def process_messages(messages, tools=None, tool_choice="none"):
+    _messages = messages
+    processed_messages = []
+    msg_has_sys = False
+
+    def filter_tools(tool_choice, tools):
+        function_name = tool_choice.get('function', {}).get('name', None)
+        if not function_name:
+            return []
+        filtered_tools = [
+            tool for tool in tools
+            if tool.get('function', {}).get('name') == function_name
+        ]
+        return filtered_tools
+
+    if tool_choice != "none":
+        if isinstance(tool_choice, dict):
+            tools = filter_tools(tool_choice, tools)
+        if tools:
+            processed_messages.append(
+                {
+                    "role": "system",
+                    "content": None,
+                    "tools": tools
+                }
+            )
+            msg_has_sys = True
+
+    if isinstance(tool_choice, dict) and tools:
+        processed_messages.append(
+            {
+                "role": "assistant",
+                "metadata": tool_choice["function"]["name"],
+                "content": ""
+            }
+        )
+
+    for m in _messages:
+        role, content, func_call = m.role, m.content, m.function_call
+        tool_calls = getattr(m, 'tool_calls', None)
+
+        if role == "function":
+            processed_messages.append(
+                {
+                    "role": "observation",
+                    "content": content
+                }
+            )
+        elif role == "tool":
+            processed_messages.append(
+                {
+                    "role": "observation",
+                    "content": content,
+                    "function_call": True
+                }
+            )
+        elif role == "assistant":
+            if tool_calls:
+                for tool_call in tool_calls:
+                    processed_messages.append(
+                        {
+                            "role": "assistant",
+                            "metadata": tool_call.function.name,
+                            "content": tool_call.function.arguments
+                        }
+                    )
+            else:
+                for response in content.split("\n"):
+                    if "\n" in response:
+                        metadata, sub_content = response.split("\n", maxsplit=1)
+                    else:
+                        metadata, sub_content = "", response
+                    processed_messages.append(
+                        {
+                            "role": role,
+                            "metadata": metadata,
+                            "content": sub_content.strip()
+                        }
+                    )
+        else:
+            if role == "system" and msg_has_sys:
+                msg_has_sys = False
+                continue
+            processed_messages.append({"role": role, "content": content})
+
+    if not tools or tool_choice == "none":
+        for m in _messages:
+            if m.role == 'system':
+                processed_messages.insert(0, {"role": m.role, "content": m.content})
+                break
+    return processed_messages
+
+
+@app.get("/health")
+async def health() -> Response:
+    """Health check."""
+    return Response(status_code=200)
+
+
+@app.get("/v1/models", response_model=ModelList)
+async def list_models():
+    model_card = ModelCard(id="glm-4")
+    return ModelList(data=[model_card])
+
+
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+async def create_chat_completion(request: ChatCompletionRequest):
+    if len(request.messages) < 1 or request.messages[-1].role == "assistant":
+        raise HTTPException(status_code=400, detail="Invalid request")
+
+    gen_params = dict(
+        messages=request.messages,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        max_tokens=request.max_tokens or 1024,
+        echo=False,
+        stream=request.stream,
+        repetition_penalty=request.repetition_penalty,
+        tools=request.tools,
+        tool_choice=request.tool_choice,
+    )
+    logger.debug(f"==== request ====\n{gen_params}")
+
+    if request.stream:
+        predict_stream_generator = predict_stream(request.model, gen_params)
+        output = await anext(predict_stream_generator)
+        if output:
+            return EventSourceResponse(predict_stream_generator, media_type="text/event-stream")
+        logger.debug(f"First result output：\n{output}")
+
+        function_call = None
+        if output and request.tools:
+            try:
+                function_call = process_response(output, request.tools, use_tool=True)
+            except:
+                logger.warning("Failed to parse tool call")
+
+        if isinstance(function_call, dict):
+            function_call = ChoiceDeltaToolCallFunction(**function_call)
+            generate = parse_output_text(request.model, output, function_call=function_call)
+            return EventSourceResponse(generate, media_type="text/event-stream")
+        else:
+            return EventSourceResponse(predict_stream_generator, media_type="text/event-stream")
+    response = ""
+    async for response in generate_stream_glm4(gen_params):
+        pass
+
+    if response["text"].startswith("\n"):
+        response["text"] = response["text"][1:]
+    response["text"] = response["text"].strip()
+
+    usage = UsageInfo()
+
+    function_call, finish_reason = None, "stop"
+    tool_calls = None
+    if request.tools:
+        try:
+            function_call = process_response(response["text"], request.tools, use_tool=True)
+        except Exception as e:
+            logger.warning(f"Failed to parse tool call: {e}")
+    if isinstance(function_call, dict):
+        finish_reason = "tool_calls"
+        function_call_response = ChoiceDeltaToolCallFunction(**function_call)
+        function_call_instance = FunctionCall(
+            name=function_call_response.name,
+            arguments=function_call_response.arguments
+        )
+        tool_calls = [
+            ChatCompletionMessageToolCall(
+                id=generate_id('call_', 24),
+                function=function_call_instance,
+                type="function")]
+
+    message = ChatMessage(
+        role="assistant",
+        content=None if tool_calls else response["text"],
+        function_call=None,
+        tool_calls=tool_calls,
+    )
+
+    logger.debug(f"==== message ====\n{message}")
+
+    choice_data = ChatCompletionResponseChoice(
+        index=0,
+        message=message,
+        finish_reason=finish_reason,
+    )
+    task_usage = UsageInfo.model_validate(response["usage"])
+    for usage_key, usage_value in task_usage.model_dump().items():
+        setattr(usage, usage_key, getattr(usage, usage_key) + usage_value)
+
+    return ChatCompletionResponse(
+        model=request.model,
+        choices=[choice_data],
+        object="chat.completion",
+        usage=usage
+    )
+
+
+async def predict_stream(model_id, gen_params):
+    output = ""
+    is_function_call = False
+    has_send_first_chunk = False
+    created_time = int(time.time())
+    function_name = None
+    response_id = generate_id('chatcmpl-', 29)
+    system_fingerprint = generate_id('fp_', 9)
+    tools = {tool['function']['name'] for tool in gen_params['tools']} if gen_params['tools'] else {}
+    async for new_response in generate_stream_glm4(gen_params):
+        decoded_unicode = new_response["text"]
+        delta_text = decoded_unicode[len(output):]
+        output = decoded_unicode
+        lines = output.strip().split("\n")
+
+        # 检查是否为工具
+        # 这是一个简单的工具比较函数，不能保证拦截所有非工具输出的结果，比如参数未对齐等特殊情况。
+        ##TODO 如果你希望做更多处理，可以在这里进行逻辑完善。
+
+        if not is_function_call and len(lines) >= 2:
+            first_line = lines[0].strip()
+            if first_line in tools:
+                is_function_call = True
+                function_name = first_line
+
+        # 工具调用返回
+        if is_function_call:
+            if not has_send_first_chunk:
+                function_call = {"name": function_name, "arguments": ""}
+                tool_call = ChatCompletionMessageToolCall(
+                    index=0,
+                    id=generate_id('call_', 24),
+                    function=FunctionCall(**function_call),
+                    type="function"
+                )
+                message = DeltaMessage(
+                    content=None,
+                    role="assistant",
+                    function_call=None,
+                    tool_calls=[tool_call]
+                )
+                choice_data = ChatCompletionResponseStreamChoice(
+                    index=0,
+                    delta=message,
+                    finish_reason=None
+                )
+                chunk = ChatCompletionResponse(
+                    model=model_id,
+                    id=response_id,
+                    choices=[choice_data],
+                    created=created_time,
+                    system_fingerprint=system_fingerprint,
+                    object="chat.completion.chunk"
+                )
+                yield ""
+                yield chunk.model_dump_json(exclude_unset=True)
+                has_send_first_chunk = True
+
+            function_call = {"name": None, "arguments": delta_text}
+            tool_call = ChatCompletionMessageToolCall(
+                index=0,
+                id=None,
+                function=FunctionCall(**function_call),
+                type="function"
+            )
+            message = DeltaMessage(
+                content=None,
+                role=None,
+                function_call=None,
+                tool_calls=[tool_call]
+            )
+            choice_data = ChatCompletionResponseStreamChoice(
+                index=0,
+                delta=message,
+                finish_reason=None
+            )
+            chunk = ChatCompletionResponse(
+                model=model_id,
+                id=response_id,
+                choices=[choice_data],
+                created=created_time,
+                system_fingerprint=system_fingerprint,
+                object="chat.completion.chunk"
+            )
+            yield chunk.model_dump_json(exclude_unset=True)
+
+        # 用户请求了 Function Call 但是框架还没确定是否为Function Call
+        elif (gen_params["tools"] and gen_params["tool_choice"] != "none") or is_function_call:
+            continue
+
+        # 常规返回
+        else:
+            finish_reason = new_response.get("finish_reason", None)
+            if not has_send_first_chunk:
+                message = DeltaMessage(
+                    content="",
+                    role="assistant",
+                    function_call=None,
+                )
+                choice_data = ChatCompletionResponseStreamChoice(
+                    index=0,
+                    delta=message,
+                    finish_reason=finish_reason
+                )
+                chunk = ChatCompletionResponse(
+                    model=model_id,
+                    id=response_id,
+                    choices=[choice_data],
+                    created=created_time,
+                    system_fingerprint=system_fingerprint,
+                    object="chat.completion.chunk"
+                )
+                yield chunk.model_dump_json(exclude_unset=True)
+                has_send_first_chunk = True
+
+            message = DeltaMessage(
+                content=delta_text,
+                role="assistant",
+                function_call=None,
+            )
+            choice_data = ChatCompletionResponseStreamChoice(
+                index=0,
+                delta=message,
+                finish_reason=finish_reason
+            )
+            chunk = ChatCompletionResponse(
+                model=model_id,
+                id=response_id,
+                choices=[choice_data],
+                created=created_time,
+                system_fingerprint=system_fingerprint,
+                object="chat.completion.chunk"
+            )
+            yield chunk.model_dump_json(exclude_unset=True)
+
+    # 工具调用需要额外返回一个字段以对齐 OpenAI 接口
+    if is_function_call:
+        yield ChatCompletionResponse(
+            model=model_id,
+            id=response_id,
+            system_fingerprint=system_fingerprint,
+            choices=[
+                ChatCompletionResponseStreamChoice(
+                    index=0,
+                    delta=DeltaMessage(
+                        content=None,
+                        role=None,
+                        function_call=None,
+                    ),
+                    finish_reason="tool_calls"
+                )],
+            created=created_time,
+            object="chat.completion.chunk",
+            usage=None
+        ).model_dump_json(exclude_unset=True)
+    yield '[DONE]'
+
+
+async def parse_output_text(model_id: str, value: str, function_call: ChoiceDeltaToolCallFunction = None):
+    delta = DeltaMessage(role="assistant", content=value)
+    if function_call is not None:
+        delta.function_call = function_call
+
+    choice_data = ChatCompletionResponseStreamChoice(
+        index=0,
+        delta=delta,
+        finish_reason=None
+    )
+    chunk = ChatCompletionResponse(
+        model=model_id,
+        choices=[choice_data],
+        object="chat.completion.chunk"
+    )
+    yield "{}".format(chunk.model_dump_json(exclude_unset=True))
+    yield '[DONE]'
+
+
+if __name__ == "__main__":
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+    engine_args = AsyncEngineArgs(
+        model=MODEL_PATH,
+        tokenizer=MODEL_PATH,
+        # 如果你有多张显卡，可以在这里设置成你的显卡数量
+        tensor_parallel_size=1,
+        dtype="bfloat16",
+        trust_remote_code=True,
+        # 占用显存的比例，请根据你的显卡显存大小设置合适的值，例如，如果你的显卡有80G，您只想使用24G，请按照24/80=0.3设置
+        gpu_memory_utilization=0.9,
+        enforce_eager=True,
+        worker_use_ray=False,
+        engine_use_ray=False,
+        disable_log_requests=True,
+        max_model_len=MAX_MODEL_LENGTH,
+    )
+    engine = AsyncLLMEngine.from_engine_args(engine_args)
+    uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)

requirements.txt

@@ -0,0 +1,29 @@
+# use vllm
+# vllm>=0.5.0
+
+torch>=2.3.0
+torchvision>=0.18.0
+transformers==4.40.0
+huggingface-hub>=0.23.1
+sentencepiece>=0.2.0
+pydantic>=2.7.1
+timm>=0.9.16
+tiktoken>=0.7.0
+accelerate>=0.30.1
+sentence_transformers>=2.7.0
+
+# web demo
+gradio>=4.33.0
+
+# openai demo
+openai>=1.34.0
+einops>=0.7.0
+sse-starlette>=2.1.0
+
+# INT4
+# bitsandbytes>=0.43.1
+bitsandbytes>=0.42.0
+
+
+# PEFT model, not need if you don't use PEFT finetune model.
+peft>=0.11.0

trans_batch_demo.py

@@ -0,0 +1,90 @@
+"""
+
+Here is an example of using batch request glm-4-9b,
+here you need to build the conversation format yourself and then call the batch function to make batch requests.
+Please note that in this demo, the memory consumption is significantly higher.
+
+"""
+
+from typing import Optional, Union
+from transformers import AutoModel, AutoTokenizer, LogitsProcessorList
+
+MODEL_PATH = 'THUDM/glm-4-9b-chat'
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    encode_special_tokens=True)
+model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True, device_map="auto").eval()
+
+
+def process_model_outputs(inputs, outputs, tokenizer):
+    responses = []
+    for input_ids, output_ids in zip(inputs.input_ids, outputs):
+        response = tokenizer.decode(output_ids[len(input_ids):], skip_special_tokens=True).strip()
+        responses.append(response)
+    return responses
+
+
+def batch(
+        model,
+        tokenizer,
+        messages: Union[str, list[str]],
+        max_input_tokens: int = 8192,
+        max_new_tokens: int = 8192,
+        num_beams: int = 1,
+        do_sample: bool = True,
+        top_p: float = 0.8,
+        temperature: float = 0.8,
+        logits_processor: Optional[LogitsProcessorList] = LogitsProcessorList(),
+):
+    messages = [messages] if isinstance(messages, str) else messages
+    batched_inputs = tokenizer(messages, return_tensors="pt", padding="max_length", truncation=True,
+                               max_length=max_input_tokens).to(model.device)
+
+    gen_kwargs = {
+        "max_new_tokens": max_new_tokens,
+        "num_beams": num_beams,
+        "do_sample": do_sample,
+        "top_p": top_p,
+        "temperature": temperature,
+        "logits_processor": logits_processor,
+        "eos_token_id": model.config.eos_token_id
+    }
+    batched_outputs = model.generate(**batched_inputs, **gen_kwargs)
+    batched_response = process_model_outputs(batched_inputs, batched_outputs, tokenizer)
+    return batched_response
+
+
+if __name__ == "__main__":
+
+    batch_message = [
+        [
+            {"role": "user", "content": "我的爸爸和妈妈结婚为什么不能带我去"},
+            {"role": "assistant", "content": "因为他们结婚时你还没有出生"},
+            {"role": "user", "content": "我刚才的提问是"}
+        ],
+        [
+            {"role": "user", "content": "你好，你是谁"}
+        ]
+    ]
+
+    batch_inputs = []
+    max_input_tokens = 1024
+    for i, messages in enumerate(batch_message):
+        new_batch_input = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        max_input_tokens = max(max_input_tokens, len(new_batch_input))
+        batch_inputs.append(new_batch_input)
+    gen_kwargs = {
+        "max_input_tokens": max_input_tokens,
+        "max_new_tokens": 8192,
+        "do_sample": True,
+        "top_p": 0.8,
+        "temperature": 0.8,
+        "num_beams": 1,
+    }
+
+    batch_responses = batch(model, tokenizer, batch_inputs, **gen_kwargs)
+    for response in batch_responses:
+        print("=" * 10)
+        print(response)

trans_cli_demo.py

@@ -0,0 +1,119 @@
+"""
+This script creates a CLI demo with transformers backend for the glm-4-9b model,
+allowing users to interact with the model through a command-line interface.
+
+Usage:
+- Run the script to start the CLI demo.
+- Interact with the model by typing questions and receiving responses.
+
+Note: The script includes a modification to handle markdown to plain text conversion,
+ensuring that the CLI interface displays formatted text correctly.
+
+If you use flash attention, you should install the flash-attn and  add attn_implementation="flash_attention_2" in model loading.
+"""
+
+import os
+import torch
+from threading import Thread
+from transformers import AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer, AutoModel
+
+MODEL_PATH = os.environ.get('MODEL_PATH', 'THUDM/glm-4-9b-chat')
+
+print("MODEL_PATH: " + MODEL_PATH)
+
+## If use peft model.
+# def load_model_and_tokenizer(model_dir, trust_remote_code: bool = True):
+#     if (model_dir / 'adapter_config.json').exists():
+#         model = AutoModel.from_pretrained(
+#             model_dir, trust_remote_code=trust_remote_code, device_map='auto'
+#         )
+#         tokenizer_dir = model.peft_config['default'].base_model_name_or_path
+#     else:
+#         model = AutoModel.from_pretrained(
+#             model_dir, trust_remote_code=trust_remote_code, device_map='auto'
+#         )
+#         tokenizer_dir = model_dir
+#     tokenizer = AutoTokenizer.from_pretrained(
+#         tokenizer_dir, trust_remote_code=trust_remote_code, use_fast=False
+#     )
+#     return model, tokenizer
+
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    encode_special_tokens=True
+)
+
+model = AutoModel.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    # attn_implementation="flash_attention_2", # Use Flash Attention
+    # torch_dtype=torch.bfloat16, #using flash-attn must use bfloat16 or float16
+    device_map="auto").eval()
+
+
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = model.config.eos_token_id
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+
+if __name__ == "__main__":
+    history = []
+    max_length = 8192
+    top_p = 0.8
+    temperature = 0.6
+    stop = StopOnTokens()
+
+    print("Welcome to the GLM-4-9B CLI chat. Type your messages below.")
+    while True:
+        user_input = input("\nYou: ")
+        if user_input.lower() in ["exit", "quit"]:
+            break
+        history.append([user_input, ""])
+
+        messages = []
+        for idx, (user_msg, model_msg) in enumerate(history):
+            if idx == len(history) - 1 and not model_msg:
+                messages.append({"role": "user", "content": user_msg})
+                break
+            if user_msg:
+                messages.append({"role": "user", "content": user_msg})
+            if model_msg:
+                messages.append({"role": "assistant", "content": model_msg})
+        model_inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_tensors="pt"
+        ).to(model.device)
+        streamer = TextIteratorStreamer(
+            tokenizer=tokenizer,
+            timeout=60,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+        generate_kwargs = {
+            "input_ids": model_inputs,
+            "streamer": streamer,
+            "max_new_tokens": max_length,
+            "do_sample": True,
+            "top_p": top_p,
+            "temperature": temperature,
+            "stopping_criteria": StoppingCriteriaList([stop]),
+            "repetition_penalty": 1.2,
+            "eos_token_id": model.config.eos_token_id,
+        }
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
+        print("GLM-4:", end="", flush=True)
+        for new_token in streamer:
+            if new_token:
+                print(new_token, end="", flush=True)
+                history[-1][1] += new_token
+
+        history[-1][1] = history[-1][1].strip()

trans_cli_vision_demo.py

@@ -0,0 +1,123 @@
+"""
+This script creates a CLI demo with transformers backend for the glm-4v-9b model,
+allowing users to interact with the model through a command-line interface.
+
+Usage:
+- Run the script to start the CLI demo.
+- Interact with the model by typing questions and receiving responses.
+
+Note: The script includes a modification to handle markdown to plain text conversion,
+ensuring that the CLI interface displays formatted text correctly.
+"""
+
+import os
+import torch
+from threading import Thread
+from transformers import (
+    AutoTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer, AutoModel, BitsAndBytesConfig
+)
+
+from PIL import Image
+
+MODEL_PATH = os.environ.get('MODEL_PATH', 'THUDM/glm-4v-9b')
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    encode_special_tokens=True
+)
+model = AutoModel.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    # attn_implementation="flash_attention_2",  # Use Flash Attention
+    # torch_dtype=torch.bfloat16,  # using flash-attn must use bfloat16 or float16,
+    device_map="auto",
+).eval()
+
+
+## For INT4 inference
+# model = AutoModel.from_pretrained(
+#     MODEL_PATH,
+#     trust_remote_code=True,
+#     quantization_config=BitsAndBytesConfig(load_in_4bit=True),
+#     torch_dtype=torch.bfloat16,
+#     low_cpu_mem_usage=True
+# ).eval()
+
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = model.config.eos_token_id
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+
+if __name__ == "__main__":
+    history = []
+    max_length = 1024
+    top_p = 0.8
+    temperature = 0.6
+    stop = StopOnTokens()
+    uploaded = False
+    image = None
+    print("Welcome to the GLM-4-9B CLI chat. Type your messages below.")
+    image_path = input("Image Path:")
+    try:
+        image = Image.open(image_path).convert("RGB")
+    except:
+        print("Invalid image path. Continuing with text conversation.")
+    while True:
+        user_input = input("\nYou: ")
+        if user_input.lower() in ["exit", "quit"]:
+            break
+        history.append([user_input, ""])
+
+        messages = []
+        for idx, (user_msg, model_msg) in enumerate(history):
+            if idx == len(history) - 1 and not model_msg:
+                messages.append({"role": "user", "content": user_msg})
+                if image and not uploaded:
+                    messages[-1].update({"image": image})
+                    uploaded = True
+                break
+            if user_msg:
+                messages.append({"role": "user", "content": user_msg})
+            if model_msg:
+                messages.append({"role": "assistant", "content": model_msg})
+        model_inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_tensors="pt",
+            return_dict=True
+        ).to(next(model.parameters()).device)
+        streamer = TextIteratorStreamer(
+            tokenizer=tokenizer,
+            timeout=60,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+        generate_kwargs = {
+            **model_inputs,
+            "streamer": streamer,
+            "max_new_tokens": max_length,
+            "do_sample": True,
+            "top_p": top_p,
+            "temperature": temperature,
+            "stopping_criteria": StoppingCriteriaList([stop]),
+            "repetition_penalty": 1.2,
+            "eos_token_id": [151329, 151336, 151338],
+        }
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
+        print("GLM-4V:", end="", flush=True)
+        for new_token in streamer:
+            if new_token:
+                print(new_token, end="", flush=True)
+                history[-1][1] += new_token
+
+        history[-1][1] = history[-1][1].strip()

trans_stress_test.py

@@ -0,0 +1,135 @@
+import argparse
+import time
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
+import torch
+from threading import Thread
+
+MODEL_PATH = 'THUDM/glm-4-9b-chat'
+
+
+def stress_test(token_len, n, num_gpu):
+    device = torch.device(f"cuda:{num_gpu - 1}" if torch.cuda.is_available() and num_gpu > 0 else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_PATH,
+        trust_remote_code=True,
+        padding_side="left"
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_PATH,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16
+    ).to(device).eval()
+    
+    # Use INT4 weight infer
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     MODEL_PATH,
+    #     trust_remote_code=True,
+    #     quantization_config=BitsAndBytesConfig(load_in_4bit=True),
+    #     low_cpu_mem_usage=True,
+    # ).eval()
+
+    times = []
+    decode_times = []
+
+    print("Warming up...")
+    vocab_size = tokenizer.vocab_size
+    warmup_token_len = 20
+    random_token_ids = torch.randint(3, vocab_size - 200, (warmup_token_len - 5,), dtype=torch.long)
+    start_tokens = [151331, 151333, 151336, 198]
+    end_tokens = [151337]
+    input_ids = torch.tensor(start_tokens + random_token_ids.tolist() + end_tokens, dtype=torch.long).unsqueeze(0).to(
+        device)
+    attention_mask = torch.ones_like(input_ids, dtype=torch.bfloat16).to(device)
+    position_ids = torch.arange(len(input_ids[0]), dtype=torch.bfloat16).unsqueeze(0).to(device)
+    warmup_inputs = {
+        'input_ids': input_ids,
+        'attention_mask': attention_mask,
+        'position_ids': position_ids
+    }
+    with torch.no_grad():
+        _ = model.generate(
+            input_ids=warmup_inputs['input_ids'],
+            attention_mask=warmup_inputs['attention_mask'],
+            max_new_tokens=2048,
+            do_sample=False,
+            repetition_penalty=1.0,
+            eos_token_id=[151329, 151336, 151338]
+        )
+    print("Warming up complete. Starting stress test...")
+
+    for i in range(n):
+        random_token_ids = torch.randint(3, vocab_size - 200, (token_len - 5,), dtype=torch.long)
+        input_ids = torch.tensor(start_tokens + random_token_ids.tolist() + end_tokens, dtype=torch.long).unsqueeze(
+            0).to(device)
+        attention_mask = torch.ones_like(input_ids, dtype=torch.bfloat16).to(device)
+        position_ids = torch.arange(len(input_ids[0]), dtype=torch.bfloat16).unsqueeze(0).to(device)
+        test_inputs = {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask,
+            'position_ids': position_ids
+        }
+
+        streamer = TextIteratorStreamer(
+            tokenizer=tokenizer,
+            timeout=36000,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+
+        generate_kwargs = {
+            "input_ids": test_inputs['input_ids'],
+            "attention_mask": test_inputs['attention_mask'],
+            "max_new_tokens": 512,
+            "do_sample": False,
+            "repetition_penalty": 1.0,
+            "eos_token_id": [151329, 151336, 151338],
+            "streamer": streamer
+        }
+
+        start_time = time.time()
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
+
+        first_token_time = None
+        all_token_times = []
+
+        for token in streamer:
+            current_time = time.time()
+            if first_token_time is None:
+                first_token_time = current_time
+                times.append(first_token_time - start_time)
+            all_token_times.append(current_time)
+
+        t.join()
+        end_time = time.time()
+
+        avg_decode_time_per_token = len(all_token_times) / (end_time - first_token_time) if all_token_times else 0
+        decode_times.append(avg_decode_time_per_token)
+        print(
+            f"Iteration {i + 1}/{n} - Prefilling Time: {times[-1]:.4f} seconds - Average Decode Time: {avg_decode_time_per_token:.4f} tokens/second")
+
+        torch.cuda.empty_cache()
+
+    avg_first_token_time = sum(times) / n
+    avg_decode_time = sum(decode_times) / n
+    print(f"\nAverage First Token Time over {n} iterations: {avg_first_token_time:.4f} seconds")
+    print(f"Average Decode Time per Token over {n} iterations: {avg_decode_time:.4f} tokens/second")
+    return times, avg_first_token_time, decode_times, avg_decode_time
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Stress test for model inference")
+    parser.add_argument('--token_len', type=int, default=1000, help='Number of tokens for each test')
+    parser.add_argument('--n', type=int, default=3, help='Number of iterations for the stress test')
+    parser.add_argument('--num_gpu', type=int, default=1, help='Number of GPUs to use for inference')
+    args = parser.parse_args()
+
+    token_len = args.token_len
+    n = args.n
+    num_gpu = args.num_gpu
+
+    stress_test(token_len, n, num_gpu)
+
+
+if __name__ == "__main__":
+    main()

trans_web_demo.py

@@ -0,0 +1,179 @@
+"""
+This script creates an interactive web demo for the GLM-4-9B model using Gradio,
+a Python library for building quick and easy UI components for machine learning models.
+It's designed to showcase the capabilities of the GLM-4-9B model in a user-friendly interface,
+allowing users to interact with the model through a chat-like interface.
+"""
+
+import os
+from pathlib import Path
+from threading import Thread
+from typing import Union
+
+import gradio as gr
+import torch
+from peft import AutoPeftModelForCausalLM, PeftModelForCausalLM
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer
+)
+
+ModelType = Union[PreTrainedModel, PeftModelForCausalLM]
+TokenizerType = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
+
+MODEL_PATH = os.environ.get('MODEL_PATH', 'THUDM/glm-4-9b-chat')
+MODEL_PATH = "/Users/zmac/.cache/huggingface/hub/models--THUDM--glm-4-9b-chat/snapshots/04419001bc63e05e70991ade6da1f91c4aeec278"
+TOKENIZER_PATH = os.environ.get("TOKENIZER_PATH", MODEL_PATH)
+
+
+def _resolve_path(path: Union[str, Path]) -> Path:
+    return Path(path).expanduser().resolve()
+
+
+def load_model_and_tokenizer(
+        model_dir: Union[str, Path], trust_remote_code: bool = True
+) -> tuple[ModelType, TokenizerType]:
+    model_dir = _resolve_path(model_dir)
+    if (model_dir / 'adapter_config.json').exists():
+        model = AutoPeftModelForCausalLM.from_pretrained(
+            model_dir, trust_remote_code=trust_remote_code, device_map='auto'
+        )
+        tokenizer_dir = model.peft_config['default'].base_model_name_or_path
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_dir, trust_remote_code=trust_remote_code, device_map='auto'
+        )
+        tokenizer_dir = model_dir
+    tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_dir, trust_remote_code=trust_remote_code, use_fast=False
+    )
+    return model, tokenizer
+
+
+model, tokenizer = load_model_and_tokenizer(MODEL_PATH, trust_remote_code=True)
+
+
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = model.config.eos_token_id
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+
+def parse_text(text):
+    lines = text.split("\n")
+    lines = [line for line in lines if line != ""]
+    count = 0
+    for i, line in enumerate(lines):
+        if "```" in line:
+            count += 1
+            items = line.split('`')
+            if count % 2 == 1:
+                lines[i] = f'<pre><code class="language-{items[-1]}">'
+            else:
+                lines[i] = f'<br></code></pre>'
+        else:
+            if i > 0:
+                if count % 2 == 1:
+                    line = line.replace("`", "\`")
+                    line = line.replace("<", "&lt;")
+                    line = line.replace(">", "&gt;")
+                    line = line.replace(" ", "&nbsp;")
+                    line = line.replace("*", "&ast;")
+                    line = line.replace("_", "&lowbar;")
+                    line = line.replace("-", "&#45;")
+                    line = line.replace(".", "&#46;")
+                    line = line.replace("!", "&#33;")
+                    line = line.replace("(", "&#40;")
+                    line = line.replace(")", "&#41;")
+                    line = line.replace("$", "&#36;")
+                lines[i] = "<br>" + line
+    text = "".join(lines)
+    return text
+
+
+def predict(history, prompt, max_length, top_p, temperature):
+    stop = StopOnTokens()
+    messages = []
+    if prompt:
+        messages.append({"role": "system", "content": prompt})
+    for idx, (user_msg, model_msg) in enumerate(history):
+        if prompt and idx == 0:
+            continue
+        if idx == len(history) - 1 and not model_msg:
+            messages.append({"role": "user", "content": user_msg})
+            break
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if model_msg:
+            messages.append({"role": "assistant", "content": model_msg})
+
+    model_inputs = tokenizer.apply_chat_template(messages,
+                                                 add_generation_prompt=True,
+                                                 tokenize=True,
+                                                 return_tensors="pt").to(next(model.parameters()).device)
+    streamer = TextIteratorStreamer(tokenizer, timeout=60, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = {
+        "input_ids": model_inputs,
+        "streamer": streamer,
+        "max_new_tokens": max_length,
+        "do_sample": True,
+        "top_p": top_p,
+        "temperature": temperature,
+        "stopping_criteria": StoppingCriteriaList([stop]),
+        "repetition_penalty": 1.2,
+        "eos_token_id": model.config.eos_token_id,
+    }
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+    for new_token in streamer:
+        if new_token:
+            history[-1][1] += new_token
+        yield history
+
+
+with gr.Blocks() as demo:
+    gr.HTML("""<h1 align="center">GLM-4-9B Gradio Simple Chat Demo</h1>""")
+    chatbot = gr.Chatbot()
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            with gr.Column(scale=12):
+                user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10, container=False)
+            with gr.Column(min_width=32, scale=1):
+                submitBtn = gr.Button("Submit")
+        with gr.Column(scale=1):
+            prompt_input = gr.Textbox(show_label=False, placeholder="Prompt", lines=10, container=False)
+            pBtn = gr.Button("Set Prompt")
+        with gr.Column(scale=1):
+            emptyBtn = gr.Button("Clear History")
+            max_length = gr.Slider(0, 32768, value=8192, step=1.0, label="Maximum length", interactive=True)
+            top_p = gr.Slider(0, 1, value=0.8, step=0.01, label="Top P", interactive=True)
+            temperature = gr.Slider(0.01, 1, value=0.6, step=0.01, label="Temperature", interactive=True)
+
+
+    def user(query, history):
+        return "", history + [[parse_text(query), ""]]
+
+
+    def set_prompt(prompt_text):
+        return [[parse_text(prompt_text), "成功设置prompt"]]
+
+
+    pBtn.click(set_prompt, inputs=[prompt_input], outputs=chatbot)
+
+    submitBtn.click(user, [user_input, chatbot], [user_input, chatbot], queue=False).then(
+        predict, [chatbot, prompt_input, max_length, top_p, temperature], chatbot
+    )
+    emptyBtn.click(lambda: (None, None), None, [chatbot, prompt_input], queue=False)
+
+demo.queue()
+demo.launch(server_name="127.0.0.1", server_port=8000, inbrowser=True, share=True)

trans_web_vision_demo.py

@@ -0,0 +1,121 @@
+"""
+This script creates a Gradio demo with a Transformers backend for the glm-4v-9b model, allowing users to interact with the model through a Gradio web UI.
+
+Usage:
+- Run the script to start the Gradio server.
+- Interact with the model via the web UI.
+
+Requirements:
+- Gradio package
+  - Type `pip install gradio` to install Gradio.
+"""
+
+import os
+import torch
+import gradio as gr
+from threading import Thread
+from transformers import (
+    AutoTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer, AutoModel, BitsAndBytesConfig
+)
+from PIL import Image
+import requests
+from io import BytesIO
+
+MODEL_PATH = os.environ.get('MODEL_PATH', 'THUDM/glm-4v-9b')
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    encode_special_tokens=True
+)
+model = AutoModel.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    device_map="auto",
+    torch_dtype=torch.bfloat16
+).eval()
+
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = model.config.eos_token_id
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+def get_image(image_path=None, image_url=None):
+    if image_path:
+        return Image.open(image_path).convert("RGB")
+    elif image_url:
+        response = requests.get(image_url)
+        return Image.open(BytesIO(response.content)).convert("RGB")
+    return None
+
+def chatbot(image_path=None, image_url=None, assistant_prompt=""):
+    image = get_image(image_path, image_url)
+
+    messages = [
+        {"role": "assistant", "content": assistant_prompt},
+        {"role": "user", "content": "", "image": image}
+    ]
+
+    model_inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_tensors="pt",
+        return_dict=True
+    ).to(next(model.parameters()).device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer=tokenizer,
+        timeout=60,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
+
+    generate_kwargs = {
+        **model_inputs,
+        "streamer": streamer,
+        "max_new_tokens": 1024,
+        "do_sample": True,
+        "top_p": 0.8,
+        "temperature": 0.6,
+        "stopping_criteria": StoppingCriteriaList([StopOnTokens()]),
+        "repetition_penalty": 1.2,
+        "eos_token_id": [151329, 151336, 151338],
+    }
+
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    response = ""
+    for new_token in streamer:
+        if new_token:
+            response += new_token
+
+    return image, response.strip()
+
+with gr.Blocks() as demo:
+    demo.title = "GLM-4V-9B Image Recognition Demo"
+    demo.description = """
+    This demo uses the GLM-4V-9B model to got image infomation.
+    """
+    with gr.Row():
+        with gr.Column():
+            image_path_input = gr.File(label="Upload Image (High-Priority)", type="filepath")
+            image_url_input = gr.Textbox(label="Image URL (Low-Priority)")
+            assistant_prompt_input = gr.Textbox(label="Assistant Prompt (You Can Change It)", value="这是什么？")
+            submit_button = gr.Button("Submit")
+        with gr.Column():
+            chatbot_output = gr.Textbox(label="GLM-4V-9B Model Response")
+            image_output = gr.Image(label="Image Preview")
+
+    submit_button.click(chatbot,
+                        inputs=[image_path_input, image_url_input, assistant_prompt_input],
+                        outputs=[image_output, chatbot_output])
+
+demo.launch(server_name="127.0.0.1", server_port=8911, inbrowser=True, share=False)

vllm_cli_demo.py

@@ -0,0 +1,111 @@
+"""
+This script creates a CLI demo with vllm backand for the glm-4-9b model,
+allowing users to interact with the model through a command-line interface.
+
+Usage:
+- Run the script to start the CLI demo.
+- Interact with the model by typing questions and receiving responses.
+
+Note: The script includes a modification to handle markdown to plain text conversion,
+ensuring that the CLI interface displays formatted text correctly.
+"""
+import time
+import asyncio
+from transformers import AutoTokenizer
+from vllm import SamplingParams, AsyncEngineArgs, AsyncLLMEngine
+from typing import List, Dict
+
+MODEL_PATH = 'THUDM/glm-4-9b'
+
+
+def load_model_and_tokenizer(model_dir: str):
+    engine_args = AsyncEngineArgs(
+        model=model_dir,
+        tokenizer=model_dir,
+        tensor_parallel_size=1,
+        dtype="bfloat16",
+        trust_remote_code=True,
+        gpu_memory_utilization=0.3,
+        enforce_eager=True,
+        worker_use_ray=True,
+        engine_use_ray=False,
+        disable_log_requests=True
+        # 如果遇见 OOM 现象，建议开启下述参数
+        # enable_chunked_prefill=True,
+        # max_num_batched_tokens=8192
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_dir,
+        trust_remote_code=True,
+        encode_special_tokens=True
+    )
+    engine = AsyncLLMEngine.from_engine_args(engine_args)
+    return engine, tokenizer
+
+
+engine, tokenizer = load_model_and_tokenizer(MODEL_PATH)
+
+
+async def vllm_gen(messages: List[Dict[str, str]], top_p: float, temperature: float, max_dec_len: int):
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=False
+    )
+    params_dict = {
+        "n": 1,
+        "best_of": 1,
+        "presence_penalty": 1.0,
+        "frequency_penalty": 0.0,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": -1,
+        "use_beam_search": False,
+        "length_penalty": 1,
+        "early_stopping": False,
+        "stop_token_ids": [151329, 151336, 151338],
+        "ignore_eos": False,
+        "max_tokens": max_dec_len,
+        "logprobs": None,
+        "prompt_logprobs": None,
+        "skip_special_tokens": True,
+    }
+    sampling_params = SamplingParams(**params_dict)
+    async for output in engine.generate(inputs=inputs, sampling_params=sampling_params, request_id=f"{time.time()}"):
+        yield output.outputs[0].text
+
+
+async def chat():
+    history = []
+    max_length = 8192
+    top_p = 0.8
+    temperature = 0.6
+
+    print("Welcome to the GLM-4-9B CLI chat. Type your messages below.")
+    while True:
+        user_input = input("\nYou: ")
+        if user_input.lower() in ["exit", "quit"]:
+            break
+        history.append([user_input, ""])
+
+        messages = []
+        for idx, (user_msg, model_msg) in enumerate(history):
+            if idx == len(history) - 1 and not model_msg:
+                messages.append({"role": "user", "content": user_msg})
+                break
+            if user_msg:
+                messages.append({"role": "user", "content": user_msg})
+            if model_msg:
+                messages.append({"role": "assistant", "content": model_msg})
+
+        print("\nGLM-4: ", end="")
+        current_length = 0
+        output = ""
+        async for output in vllm_gen(messages, top_p, temperature, max_length):
+            print(output[current_length:], end="", flush=True)
+            current_length = len(output)
+        history[-1][1] = output
+
+
+if __name__ == "__main__":
+    asyncio.run(chat())