mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2025-09-05 20:19:51 +00:00
146 lines
10 KiB
Python
146 lines
10 KiB
Python
import asyncio
|
||
import json
|
||
import sys
|
||
import aiohttp
|
||
import random
|
||
import argparse
|
||
import yaml
|
||
import os
|
||
import time
|
||
from time import sleep
|
||
|
||
decodesz = 128
|
||
# Server URL (replace with your server URL)
|
||
decodesz_list = [128]
|
||
ktansformer_prompt1024="""在遥远的翡翠森林里，住着各种各样的神奇生物。其中，有一只名叫露露的小狐狸，她与其他狐狸不同，天生长着一双晶莹剔透的翅膀。然而，这双翅膀却从未带她飞翔过。
|
||
一天，森林里传来一个惊人的消息：藏在森林深处的魔法泉水干涸了，所有生物赖以生存的泉水即将枯竭。他们说，只有传说中的“天空之羽”才能唤醒泉水，让它重新流淌。然而，“天空之羽”藏在一座高耸入云的山峰上，没有任何动物能抵达那里。
|
||
露露听到这个消息后，决定亲自去寻找“天空之羽”，即便她的翅膀无法飞翔，她也要尝试。最终，露露来到了传说中的高峰脚下，根本无法攀爬。她望着天空，心里充满了不甘：“如果我能飞起来，就不会被这座山挡住了……”
|
||
正当她感到迷茫时，一只年迈的白鹰出现在她面前。
|
||
“孩子，你为什么到这里来？”白鹰用苍老但慈祥的声音问道。
|
||
露露将森林的困境告诉了白鹰，并说自己愿意付出一切，只要能拯救森林。
|
||
白鹰沉思了一会儿，缓缓说道：“你的翅膀并不是没有力量，而是你一直害怕它们不能飞翔。相信自己，勇敢跳下去。”
|
||
露露听后，心跳加速，她望着万丈深渊，犹豫不决就在那一瞬间，她竟然真的飞了起来！露露兴奋极了，她终于看到了“天空之羽”——一根散发着金光的羽毛，轻盈地悬浮在空中。露露小心翼翼地将“天空之羽”叼住，振翅返回森林。
|
||
当她将羽毛放入干涸的泉水中时，一道金光闪耀。整个森林恢复了生机，花草重新绽放，动物们欢欣鼓舞。从那以后，露露成为了森林的英雄，她是翱翔天空的勇士。她让所有动物都明白：只要相信自己，勇敢前行，就能实现自己的梦想。
|
||
请简述这个故事的内涵 写10000个字。
|
||
在遥远的翡翠森林里，住着各种各样的神奇生物。其中，有一只名叫露露的小狐狸，她与其他狐狸不同，天生长着一双晶莹剔透的翅膀。然而，这双翅膀却从未带她飞翔过。
|
||
一天，森林里传来一个惊人的消息：藏在森林深处的魔法泉水干涸了，所有生物赖以生存的泉水即将枯竭。他们说，只有传说中的“天空之羽”才能唤醒泉水，让它重新流淌。然而，“天空之羽”藏在一座高耸入云的山峰上，没有任何动物能抵达那里。
|
||
露露听到这个消息后，决定亲自去寻找“天空之羽”，即便她的翅膀无法飞翔，她也要尝试。最终，露露来到了传说中的高峰脚下，根本无法攀爬。她望着天空，心里充满了不甘：“如果我能飞起来，就不会被这座山挡住了……”
|
||
正当她感到迷茫时，一只年迈的白鹰出现在她面前。
|
||
“孩子，你为什么到这里来？”白鹰用苍老但慈祥的声音问道。
|
||
露露将森林的困境告诉了白鹰，并说自己愿意付出一切，只要能拯救森林。
|
||
白鹰沉思了一会儿，缓缓说道：“你的翅膀并不是没有力量，而是你一直害怕它们不能飞翔。相信自己，勇敢跳下去。”
|
||
露露听后，心跳加速，她望着万丈深渊，犹豫不决就在那一瞬间，她竟然真的飞了起来！露露兴奋极了，她终于看到了“天空之羽”——一根散发着金光的羽毛，轻盈地悬浮在空中。露露小心翼翼地将“天空之羽”叼住，振翅返回森林。
|
||
当她将羽毛放入干涸的泉水中时，一道金光闪耀。整个森林恢复了生机，花草重新绽放，动物们欢欣鼓舞。从那以后，露露成为了森林的英雄，她是翱翔天空的勇士。她让所有动物都明白：只要相信自己，勇敢前行，就能实现自己的梦想。
|
||
请简述这个故事的内涵 写10000个字。
|
||
露露将森林的困境告诉了白鹰，并说自己愿意付出一切，只要能拯救森林。
|
||
白鹰沉思了一会儿，缓缓说道：“你的翅膀并不是没有力量，而是你一直害怕它们不能飞翔。相信自己，勇敢跳下去。”
|
||
露露听后，心跳加速，她望着万丈深渊，犹豫不决就在那一瞬间，她竟然真的飞了起来！露露兴奋极了，她终于看到了“天空之羽”——一根散发着金光的羽毛，轻盈地悬浮在空中。露露小心翼翼地将“天空之羽”叼住，振翅返回森林。
|
||
当她将羽毛放入干涸的泉水中时，一道金光闪耀。整个森林恢复了生机，花草重新绽放，动物们欢欣鼓舞。从那以后，露露成为了森林的英雄，她是翱翔天空的勇士。她让所有动物都明白：只要相信自己，勇敢前行，就能实现自己的梦想。
|
||
请简述这个故事的内涵 写10000个字。想。
|
||
请简述这个故事的内涵 故事的内涵这个故事的内涵写10000个字"""
|
||
async def fetch_event_stream(session, request_id , prompt):
|
||
try:
|
||
payload = {
|
||
"messages": [
|
||
{"role": "system", "content": ""},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
"model": "DeepSeek-V3",
|
||
"temperature": 0.3,
|
||
"top_p": 1.0,
|
||
"stream": True # 开启流式输出
|
||
}
|
||
|
||
headers = {
|
||
'accept': 'application/json',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
|
||
async with session.post(SERVER_URL, json=payload, headers=headers, timeout=500000) as response:
|
||
print(f"Request {request_id}: Connected, status {response.status}")
|
||
|
||
if response.status != 200:
|
||
print(f"Request {request_id}: Error, status {response.status}")
|
||
return
|
||
|
||
output_text = "" # 存储当前 response 的所有 token
|
||
total_tokens = 0 # 统计总 tokens 数
|
||
decode_start_time = None # 记录 decode 阶段开始时间
|
||
decode_end_time = None # 记录 decode 结束时间
|
||
|
||
async for line in response.content:
|
||
try:
|
||
decoded_line = line.decode("utf-8").strip()
|
||
|
||
# 过滤空行
|
||
if not decoded_line or not decoded_line.startswith("data: "):
|
||
continue
|
||
|
||
decoded_line = decoded_line[6:].strip() # 去掉 `data: `
|
||
|
||
# 确保 JSON 数据是合法的
|
||
if not decoded_line:
|
||
continue
|
||
|
||
response_data = json.loads(decoded_line) # 解析 JSON
|
||
|
||
# 确保 choices 存在
|
||
choices = response_data.get("choices", [])
|
||
if not choices:
|
||
continue
|
||
|
||
delta = choices[0].get("delta", {})
|
||
token = delta.get("content", "")
|
||
|
||
if token:
|
||
if decode_start_time is None:
|
||
decode_start_time = time.time() # 记录 decode 开始时间
|
||
|
||
output_text += token # 追加 token
|
||
sys.stdout.write(str(request_id))
|
||
sys.stdout.write(token) # 直接输出 token
|
||
sys.stdout.flush() # 立即刷新，确保 token 立刻出现在终端
|
||
total_tokens += 1 # 增加 token 计数
|
||
decode_end_time = time.time() # 每次收到 token，更新 decode 结束时间
|
||
|
||
# 检查是否完成
|
||
finish_reason = choices[0].get("finish_reason", None)
|
||
if finish_reason:
|
||
# print(f"\nRequest {request_id}: Done")
|
||
break # 结束流式处理
|
||
|
||
except json.JSONDecodeError as e:
|
||
print(f"\nRequest {request_id}: JSON Decode Error - {e}")
|
||
except IndexError:
|
||
print(f"\nRequest {request_id}: List Index Error - choices is empty")
|
||
except Exception as e:
|
||
print(f"\nRequest {request_id}: Error parsing stream - {e}")
|
||
|
||
# 计算 decode 速度
|
||
if decode_start_time and decode_end_time and total_tokens > 0:
|
||
decode_time = decode_end_time - decode_start_time
|
||
decode_speed = total_tokens / decode_time if decode_time > 0 else 0
|
||
# print(f"Request {request_id}: Decode Speed = {decode_speed:.2f} tokens/s")
|
||
|
||
except Exception as e:
|
||
print(f"\nRequest {request_id}: Exception - {e}")
|
||
|
||
async def main(concurrent_requests , prompt ):
|
||
async with aiohttp.ClientSession() as session:
|
||
tasks = [fetch_event_stream(session, i , prompt) for i in range(concurrent_requests)]
|
||
await asyncio.gather(*tasks)
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="Event Stream Request Tester")
|
||
parser.add_argument("--concurrent", type=int, default=1, help="Number of concurrent requests")
|
||
parser.add_argument("--prompt_lens", type=int, default=1024, help="prefill prompt lens, 1024 or 2048")
|
||
parser.add_argument("--api_url", type=str, default="http://localhost:10002/v1/chat/completions", help="API URL")
|
||
|
||
args = parser.parse_args()
|
||
SERVER_URL = args.api_url
|
||
if args.prompt_lens == 1024:
|
||
prompt = ktansformer_prompt1024
|
||
elif args.prompt_lens == 2048:
|
||
prompt = ktansformer_prompt1024 * 2
|
||
asyncio.run(main(args.concurrent, prompt))
|
||
|