File size: 2,017 Bytes
7a10dd5 c95976d 7a10dd5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
# 服务器端代码
# 导入Flask库和其他必要的库
from flask import Flask, request, jsonify
import threading
from flask_cors import CORS
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
# 设置可见的GPU设备
os.environ["CUDA_VISIBLE_DEVICES"] = '0, 1'
# 创建Flask应用对象
app = Flask(__name__)
# 允许跨域请求
CORS(app)
# 加载百川大模型的分词器和模型
tokenizer = AutoTokenizer.from_pretrained("baichuan-inc/Baichuan-13B-Chat", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("baichuan-inc/Baichuan-13B-Chat", device_map="auto", trust_remote_code=True)
# 创建线程锁和计数器
lock = threading.Lock()
counter = 0
MAX_CONCURRENT_REQUESTS = 50 # 最大并发请求数
# 定义服务接口的路由和方法
@app.route('/baichuan/conversation', methods=['POST'])
def conversation():
global counter # 使用全局变量
# 请求过载，返回提示信息
if counter >= MAX_CONCURRENT_REQUESTS:
return jsonify({'message': '请稍等再试'})
# 获取线程锁
with lock:
counter += 1 # 增加计数器
try:
# 接收POST请求的数据
question = request.json['question']
question += '->' # 添加分隔符
# 对输入进行分词和编码
inputs = tokenizer(question, return_tensors='pt')
inputs = inputs.to('cuda:0') # 移动到GPU上
# 调用模型进行生成
pred = model.generate(**inputs, max_new_tokens=1024, repetition_penalty=1.1)
text = tokenizer.decode(pred.cpu()[0], skip_special_tokens=True) # 对输出进行解码
# 返回结果
response = {'result': text[len(question):]} # 去掉输入部分
return jsonify(response)
finally:
# 释放线程锁并减少计数器
with lock:
counter -= 1
# 主函数
if __name__ == '__main__':
print("Flask服务器已启动")
app.run(host='0.0.0.0', port=30908) # 设置主机地址和端口号
|