mirror of
https://github.com/DrHo1y/ezrknn-llm.git
synced 2026-10-15 12:06:40 +07:00
298 lines
11 KiB
Python
298 lines
11 KiB
Python
import ctypes
|
||
import sys
|
||
import os
|
||
import subprocess
|
||
import resource
|
||
import threading
|
||
import time
|
||
import argparse
|
||
import json
|
||
from flask import Flask, request, jsonify, Response
|
||
|
||
app = Flask(__name__)
|
||
|
||
# 创建一个锁，用于控制多人访问Server
|
||
lock = threading.Lock()
|
||
|
||
# 创建一个全局变量，用于标识服务器当前是否处于阻塞状态
|
||
is_blocking = False
|
||
|
||
# 设置动态库路径
|
||
rkllm_lib = ctypes.CDLL('lib/librkllmrt.so')
|
||
|
||
# 定义全局变量，用于保存回调函数的输出，便于在gradio界面中输出
|
||
global_text = []
|
||
global_state = -1
|
||
split_byte_data = bytes(b"") # 用于保存分割的字节数据
|
||
|
||
# 定义动态库中的结构体
|
||
class Token(ctypes.Structure):
|
||
_fields_ = [
|
||
("logprob", ctypes.c_float),
|
||
("id", ctypes.c_int32)
|
||
]
|
||
|
||
class RKLLMResult(ctypes.Structure):
|
||
_fields_ = [
|
||
("text", ctypes.c_char_p),
|
||
("tokens", ctypes.POINTER(Token)),
|
||
("num", ctypes.c_int32)
|
||
]
|
||
|
||
|
||
# 定义回调函数
|
||
def callback(result, userdata, state):
|
||
global global_text, global_state, split_byte_data
|
||
if state == 0:
|
||
# 保存输出的token文本及RKLLM运行状态
|
||
global_state = state
|
||
# 需要监控当前的字节数据是否完整，不完整则进行记录，后续进行解析
|
||
try:
|
||
global_text.append((split_byte_data + result.contents.text).decode('utf-8'))
|
||
print((split_byte_data + result.contents.text).decode('utf-8'), end='')
|
||
split_byte_data = bytes(b"")
|
||
except:
|
||
split_byte_data += result.contents.text
|
||
sys.stdout.flush()
|
||
elif state == 1:
|
||
# 保存RKLLM运行状态
|
||
global_state = state
|
||
print("\n")
|
||
sys.stdout.flush()
|
||
else:
|
||
print("run error")
|
||
|
||
# Python端与C++端的回调函数连接
|
||
callback_type = ctypes.CFUNCTYPE(None, ctypes.POINTER(RKLLMResult), ctypes.c_void_p, ctypes.c_int)
|
||
c_callback = callback_type(callback)
|
||
|
||
# 定义动态库中的结构体
|
||
class RKNNllmParam(ctypes.Structure):
|
||
_fields_ = [
|
||
("model_path", ctypes.c_char_p),
|
||
("num_npu_core", ctypes.c_int32),
|
||
("max_context_len", ctypes.c_int32),
|
||
("max_new_tokens", ctypes.c_int32),
|
||
("top_k", ctypes.c_int32),
|
||
("top_p", ctypes.c_float),
|
||
("temperature", ctypes.c_float),
|
||
("repeat_penalty", ctypes.c_float),
|
||
("frequency_penalty", ctypes.c_float),
|
||
("presence_penalty", ctypes.c_float),
|
||
("mirostat", ctypes.c_int32),
|
||
("mirostat_tau", ctypes.c_float),
|
||
("mirostat_eta", ctypes.c_float),
|
||
("logprobs", ctypes.c_bool),
|
||
("top_logprobs", ctypes.c_int32),
|
||
("use_gpu", ctypes.c_bool)
|
||
]
|
||
|
||
# 定义RKLLM_Handle_t和userdata
|
||
RKLLM_Handle_t = ctypes.c_void_p
|
||
userdata = ctypes.c_void_p(None)
|
||
|
||
# 设置提示文本
|
||
PROMPT_TEXT_PREFIX = "<|im_start|>system You are a helpful assistant. <|im_end|> <|im_start|>user"
|
||
PROMPT_TEXT_POSTFIX = "<|im_end|><|im_start|>assistant"
|
||
|
||
# 定义Python端的RKLLM类，其中包括了对动态库中RKLLM模型的初始化、推理及释放操作
|
||
class RKLLM(object):
|
||
def __init__(self, model_path, target_platform):
|
||
rknnllm_param = RKNNllmParam()
|
||
rknnllm_param.model_path = bytes(model_path, 'utf-8')
|
||
if target_platform == "rk3588":
|
||
rknnllm_param.num_npu_core = 3
|
||
elif target_platform == "rk3576":
|
||
rknnllm_param.num_npu_core = 1
|
||
rknnllm_param.max_context_len = 320
|
||
rknnllm_param.max_new_tokens = 512
|
||
rknnllm_param.top_k = 1
|
||
rknnllm_param.top_p = 0.9
|
||
rknnllm_param.temperature = 0.8
|
||
rknnllm_param.repeat_penalty = 1.1
|
||
rknnllm_param.frequency_penalty = 0.0
|
||
rknnllm_param.presence_penalty = 0.0
|
||
rknnllm_param.mirostat = 0
|
||
rknnllm_param.mirostat_tau = 5.0
|
||
rknnllm_param.mirostat_eta = 0.1
|
||
rknnllm_param.logprobs = False
|
||
rknnllm_param.top_logprobs = 5
|
||
rknnllm_param.use_gpu = True
|
||
self.handle = RKLLM_Handle_t()
|
||
|
||
self.rkllm_init = rkllm_lib.rkllm_init
|
||
self.rkllm_init.argtypes = [ctypes.POINTER(RKLLM_Handle_t), ctypes.POINTER(RKNNllmParam), callback_type]
|
||
self.rkllm_init.restype = ctypes.c_int
|
||
self.rkllm_init(ctypes.byref(self.handle), rknnllm_param, c_callback)
|
||
|
||
self.rkllm_run = rkllm_lib.rkllm_run
|
||
self.rkllm_run.argtypes = [RKLLM_Handle_t, ctypes.POINTER(ctypes.c_char), ctypes.c_void_p]
|
||
self.rkllm_run.restype = ctypes.c_int
|
||
|
||
self.rkllm_destroy = rkllm_lib.rkllm_destroy
|
||
self.rkllm_destroy.argtypes = [RKLLM_Handle_t]
|
||
self.rkllm_destroy.restype = ctypes.c_int
|
||
|
||
def run(self, prompt):
|
||
prompt = bytes(PROMPT_TEXT_PREFIX + prompt + PROMPT_TEXT_POSTFIX, 'utf-8')
|
||
self.rkllm_run(self.handle, prompt, ctypes.byref(userdata))
|
||
return
|
||
|
||
def release(self):
|
||
self.rkllm_destroy(self.handle)
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--target_platform', help='目标平台: 如rk3588/rk3576;')
|
||
parser.add_argument('--rkllm_model_path', help='Linux板端上已转换好的rkllm模型的绝对路径')
|
||
args = parser.parse_args()
|
||
|
||
if not (args.target_platform in ["rk3588", "rk3576"]):
|
||
print("====== Error: 请指定正确的目标平台: rk3588/rk3576 ======")
|
||
sys.stdout.flush()
|
||
exit()
|
||
|
||
if not os.path.exists(args.rkllm_model_path):
|
||
print("====== Error: 请给出准确的rkllm模型路径，需注意是板端的绝对路径 ======")
|
||
sys.stdout.flush()
|
||
exit()
|
||
|
||
# 定频设置
|
||
command = "sudo bash fix_freq_{}.sh".format(args.target_platform)
|
||
subprocess.run(command, shell=True)
|
||
|
||
# 设置文件描述符限制
|
||
resource.setrlimit(resource.RLIMIT_NOFILE, (102400, 102400))
|
||
|
||
# 初始化RKLLM模型
|
||
print("=========init....===========")
|
||
sys.stdout.flush()
|
||
target_platform = args.target_platform
|
||
model_path = args.rkllm_model_path
|
||
rkllm_model = RKLLM(model_path, target_platform)
|
||
print("RKLLM初始化成功！")
|
||
print("==============================")
|
||
sys.stdout.flush()
|
||
|
||
# 创建一个函数用于接受用户使用 request 发送的数据
|
||
@app.route('/rkllm_chat', methods=['POST'])
|
||
def receive_message():
|
||
# 链接全局变量，获取回调函数的输出信息
|
||
global global_text, global_state
|
||
global is_blocking
|
||
|
||
# 如果服务器正在阻塞状态，则返回特定响应
|
||
if is_blocking or global_state==0:
|
||
return jsonify({'status': 'error', 'message': 'RKLLM_Server is busy! Maybe you can try again later.'}), 503
|
||
|
||
# 加锁
|
||
lock.acquire()
|
||
try:
|
||
# 设置服务器为阻塞状态
|
||
is_blocking = True
|
||
|
||
# 获取 POST 请求中的 JSON 数据
|
||
data = request.json
|
||
if data and 'messages' in data:
|
||
# 重置全局变量
|
||
global_text = []
|
||
global_state = -1
|
||
|
||
# 定义返回的结构体
|
||
rkllm_responses = {
|
||
"id": "rkllm_chat",
|
||
"object": "rkllm_chat",
|
||
"created": None,
|
||
"choices": [],
|
||
"usage": {
|
||
"prompt_tokens": None,
|
||
"completion_tokens": None,
|
||
"total_tokens": None
|
||
}
|
||
}
|
||
|
||
if not "stream" in data.keys() or data["stream"] == False:
|
||
# 在这里处理收到的数据
|
||
messages = data['messages']
|
||
print("Received messages:", messages)
|
||
for index, message in enumerate(messages):
|
||
input_prompt = message['content']
|
||
rkllm_output = ""
|
||
|
||
# 创建模型推理的线程
|
||
model_thread = threading.Thread(target=rkllm_model.run, args=(input_prompt,))
|
||
model_thread.start()
|
||
|
||
# 等待模型运行完成，定时检查模型的推理线程
|
||
model_thread_finished = False
|
||
while not model_thread_finished:
|
||
while len(global_text) > 0:
|
||
rkllm_output += global_text.pop(0)
|
||
time.sleep(0.005)
|
||
|
||
model_thread.join(timeout=0.005)
|
||
model_thread_finished = not model_thread.is_alive()
|
||
|
||
rkllm_responses["choices"].append(
|
||
{"index": index,
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": rkllm_output,
|
||
},
|
||
"logprobs": None,
|
||
"finish_reason": "stop"
|
||
}
|
||
)
|
||
return jsonify(rkllm_responses), 200
|
||
else:
|
||
# 在这里处理收到的数据
|
||
messages = data['messages']
|
||
print("Received messages:", messages)
|
||
for index, message in enumerate(messages):
|
||
input_prompt = message['content']
|
||
rkllm_output = ""
|
||
|
||
def generate():
|
||
# 创建模型推理的线程
|
||
model_thread = threading.Thread(target=rkllm_model.run, args=(input_prompt,))
|
||
model_thread.start()
|
||
|
||
# 等待模型运行完成，定时检查模型的推理线程
|
||
model_thread_finished = False
|
||
while not model_thread_finished:
|
||
while len(global_text) > 0:
|
||
rkllm_output = global_text.pop(0)
|
||
|
||
rkllm_responses["choices"].append(
|
||
{"index": index,
|
||
"delta": {
|
||
"role": "assistant",
|
||
"content": rkllm_output,
|
||
},
|
||
"logprobs": None,
|
||
"finish_reason": "stop" if global_state == 1 else None,
|
||
}
|
||
)
|
||
yield f"{json.dumps(rkllm_responses)}\n\n"
|
||
|
||
model_thread.join(timeout=0.005)
|
||
model_thread_finished = not model_thread.is_alive()
|
||
|
||
return Response(generate(), content_type='text/plain')
|
||
else:
|
||
return jsonify({'status': 'error', 'message': 'Invalid JSON data!'}), 400
|
||
finally:
|
||
# 释放锁
|
||
lock.release()
|
||
# 将服务器状态设置为非阻塞
|
||
is_blocking = False
|
||
|
||
# 启动 Flask 应用程序
|
||
# app.run(host='0.0.0.0', port=8080)
|
||
app.run(host='0.0.0.0', port=8080, threaded=True, debug=False)
|
||
|
||
print("====================")
|
||
print("RKLLM模型推理结束, 释放RKLLM模型资源...")
|
||
rkllm_model.release()
|
||
print("====================")
|