# 服务器端代码

# 导入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) # 设置主机地址和端口号