x
This commit is contained in:
@@ -0,0 +1,99 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
交互式对话脚本,用于与vLLM模型进行对话并计算token生成速度
|
||||||
|
|
||||||
|
使用方法:
|
||||||
|
1. 进入Docker容器:docker exec -it rocm-vllm-openai bash
|
||||||
|
2. 运行:python chat_with_speed.py
|
||||||
|
3. 输入提示词与模型对话
|
||||||
|
4. 输入 'exit' 退出
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
# 模型服务地址
|
||||||
|
API_URL = "http://localhost:8001/v1/chat/completions"
|
||||||
|
# API密钥
|
||||||
|
API_KEY = "sk-szcjw"
|
||||||
|
# 模型名称
|
||||||
|
MODEL_NAME = "Qwen_local_model"
|
||||||
|
|
||||||
|
def chat_with_model():
|
||||||
|
"""交互式对话函数"""
|
||||||
|
print("=== vLLM 交互式对话工具 ===")
|
||||||
|
print("输入提示词与模型对话,输入 'exit' 退出")
|
||||||
|
print("=" * 50)
|
||||||
|
|
||||||
|
# 对话历史
|
||||||
|
messages = []
|
||||||
|
|
||||||
|
while True:
|
||||||
|
# 获取用户输入
|
||||||
|
user_input = input("用户: ").strip()
|
||||||
|
|
||||||
|
if user_input.lower() == "exit":
|
||||||
|
print("退出对话...")
|
||||||
|
break
|
||||||
|
|
||||||
|
if not user_input:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 添加用户消息到对话历史
|
||||||
|
messages.append({"role": "user", "content": user_input})
|
||||||
|
|
||||||
|
# 准备请求数据
|
||||||
|
payload = {
|
||||||
|
"model": MODEL_NAME,
|
||||||
|
"messages": messages,
|
||||||
|
"max_tokens": 1000,
|
||||||
|
"temperature": 0.7,
|
||||||
|
"top_p": 0.8,
|
||||||
|
"top_k": 20
|
||||||
|
}
|
||||||
|
|
||||||
|
headers = {
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
"Authorization": f"Bearer {API_KEY}"
|
||||||
|
}
|
||||||
|
|
||||||
|
print("模型: ", end="", flush=True)
|
||||||
|
|
||||||
|
# 记录开始时间
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# 发送请求
|
||||||
|
response = httpx.post(API_URL, json=payload, headers=headers, timeout=300.0)
|
||||||
|
response.raise_for_status()
|
||||||
|
|
||||||
|
# 解析响应
|
||||||
|
result = response.json()
|
||||||
|
|
||||||
|
# 获取模型回复
|
||||||
|
assistant_message = result["choices"][0]["message"]["content"]
|
||||||
|
print(assistant_message)
|
||||||
|
|
||||||
|
# 添加模型回复到对话历史
|
||||||
|
messages.append({"role": "assistant", "content": assistant_message})
|
||||||
|
|
||||||
|
# 计算token速度
|
||||||
|
usage = result.get("usage", {})
|
||||||
|
completion_tokens = usage.get("completion_tokens", 0)
|
||||||
|
end_time = time.time()
|
||||||
|
elapsed_time = end_time - start_time
|
||||||
|
|
||||||
|
if completion_tokens > 0 and elapsed_time > 0:
|
||||||
|
tokens_per_second = completion_tokens / elapsed_time
|
||||||
|
print(f"\n[速度统计] 生成 {completion_tokens} tokens,用时 {elapsed_time:.2f} 秒,速度: {tokens_per_second:.2f} tokens/s")
|
||||||
|
else:
|
||||||
|
print("\n[速度统计] 无法计算速度")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\n错误: {e}")
|
||||||
|
|
||||||
|
print("=" * 50)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
chat_with_model()
|
||||||
+4
-2
@@ -27,7 +27,7 @@
|
|||||||
"ctx": "32768",
|
"ctx": "32768",
|
||||||
"trust_remote": true,
|
"trust_remote": true,
|
||||||
"valid_tp": [1, 2],
|
"valid_tp": [1, 2],
|
||||||
"max_num_seqs": "8",
|
"max_num_seqs": "16",
|
||||||
"max_tokens": "32768",
|
"max_tokens": "32768",
|
||||||
"gpu_util": "0.98",
|
"gpu_util": "0.98",
|
||||||
"tool_call_parser": "qwen3_coder",
|
"tool_call_parser": "qwen3_coder",
|
||||||
@@ -102,7 +102,9 @@
|
|||||||
"valid_tp": [2],
|
"valid_tp": [2],
|
||||||
"max_num_seqs": "12",
|
"max_num_seqs": "12",
|
||||||
"max_tokens": "65536",
|
"max_tokens": "65536",
|
||||||
"gpu_util": "0.94",
|
"max_num_batched_tokens": 32768,
|
||||||
|
"kv_cache_dtype": "fp8",
|
||||||
|
"gpu_util": "0.92",
|
||||||
"enforce_eager": false,
|
"enforce_eager": false,
|
||||||
"enable_prefix_caching": true,
|
"enable_prefix_caching": true,
|
||||||
"tool_call_parser": "qwen3_coder",
|
"tool_call_parser": "qwen3_coder",
|
||||||
|
|||||||
Reference in New Issue
Block a user