import time
from flask import Flask, jsonify
from flask import request
from vllm import LLM, SamplingParams

app = Flask(__name__)
app.config["JSON_AS_ASCII"] = False
# prompts = [
#     "生成论文小标题内容#问：论文题目是“大学生村官管理研究”，目录是“一、大学生村官管理现状分析\\n1.1 村官数量及分布情况\\n1.2 村官岗位设置及职责\\n1.3 村官工作绩效评估\\n\\n二、大学生村官管理存在的问题\\n2.1 村官队伍结构不合理\\n2.2 村官工作能力不足\\n2.3 村官管理制度不健全\\n\\n三、大学生村官管理对策研究\\n3.1 加强村官队伍建设\\n3.2 提高村官工作能力\\n3.3 完善村官管理制度\\n\\n四、大学生村官管理案例分析\\n4.1 案例一：某村大学生村官工作情况分析\\n4.2 案例二：某村大学生村官管理策略探讨\\n\\n五、大学生村官管理的未来发展趋势\\n5.1 多元化村官队伍建设\\n5.2 信息化村官管理模式\\n5.3 村官职业化发展\\n\\n六、大学生村官管理的政策建议\\n6.1 加强对大学生村官的培训和管理\\n6.2 完善大学生村官管理制度\\n6.3 提高大学生村官的待遇和福利\\n\\n七、结论与展望”，请把其中的小标题“3.3 完善村官管理制度”的内容补充完整，补充内容字数在800字左右\n答：\n"
# ] * 10
sampling_params = SamplingParams(temperature=0.95, top_p=0.7,presence_penalty=0.9,stop="</s>")
#
model_path = '/home/majiahui/project/models-llm/openbuddy-llama-7b-finetune'
llm = LLM(model=model_path)
#
# t1 = time.time()
# outputs = llm.generate(prompts, sampling_params)


@app.route("/predict", methods=["POST"])
def handle_query():
    print(request.remote_addr)
    texts = request.json["texts"]
    print(type(texts))
    outputs = llm.generate(texts, sampling_params)
    print(outputs)
    print("===========================================================================================================")
    # generated_text = outputs[0].outputs[0].text
    print(len(texts))
    generated_text_list = [""] * len(texts)
    for i,output in enumerate(outputs):
        index = output.request_id
        print(index)
        try:
            generated_text = output.outputs[0].text
            generated_text_list[int(index)] = generated_text
        except:
            print(output)
    result_text = {'code': "200", 'data': generated_text_list}
    return jsonify(result_text)  # 返回结果

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=15001, threaded=True, debug=False)