100 lines
2.9 KiB
Python
100 lines
2.9 KiB
Python
#!/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()
|