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",
|
||||
"trust_remote": true,
|
||||
"valid_tp": [1, 2],
|
||||
"max_num_seqs": "8",
|
||||
"max_num_seqs": "16",
|
||||
"max_tokens": "32768",
|
||||
"gpu_util": "0.98",
|
||||
"tool_call_parser": "qwen3_coder",
|
||||
@@ -102,7 +102,9 @@
|
||||
"valid_tp": [2],
|
||||
"max_num_seqs": "12",
|
||||
"max_tokens": "65536",
|
||||
"gpu_util": "0.94",
|
||||
"max_num_batched_tokens": 32768,
|
||||
"kv_cache_dtype": "fp8",
|
||||
"gpu_util": "0.92",
|
||||
"enforce_eager": false,
|
||||
"enable_prefix_caching": true,
|
||||
"tool_call_parser": "qwen3_coder",
|
||||
|
||||
Reference in New Issue
Block a user