x

chat_with_speed.py

@@ -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()