import flask
from transformers import pipeline
import redis
import uuid
import json
from threading import Thread
import time

app = flask.Flask(__name__)
pool = redis.ConnectionPool(host='localhost', port=63179, max_connections=100, db=5, password="zhicheng123*")
redis_ = redis.Redis(connection_pool=pool, decode_responses=True)

db_key_query = 'query'
db_key_result = 'result'

sampling_params = SamplingParams(temperature=0.95, top_p=0.7,presence_penalty=1.1,stop="</s>", max_tokens=4096)
models_path = "/home/majiahui/model-llm/openbuddy-mistral-7b-v13.1"
llm = LLM(model=models_path, tokenizer_mode="slow")

def mistral_vllm_models(texts):
    outputs = llm.generate(texts, sampling_params)  # 调用模型

    generated_text_list = [""] * len(texts)

    # generated_text_list = ["" if len(i[0]) > 5 else i[0] for i in text_list]

    for i, output in enumerate(outputs):
        index = output.request_id
        generated_text = output.outputs[0].text
        generated_text_list[int(index)] = generated_text


def classify(batch_size):  # 调用模型，设置最大batch_size
    while True:
        texts = []
        query_ids = []
        if redis_.llen(db_key_query) == 0:  # 若队列中没有元素就继续获取
            continue
        for i in range(min(redis_.llen(db_key_query), batch_size)):
            query = redis_.lpop(db_key_query).decode('UTF-8')  # 获取query的text
            query_ids.append(json.loads(query)['id'])
            texts.append(json.loads(query)['text'])  # 拼接若干text 为batch
        result = mistral_vllm_models(texts)  # 调用模型
        for (id_, res) in zip(query_ids, result):
            res['score'] = str(res['score'])
            redis_.set(id_, json.dumps(res))  # 将模型结果送回队列


@app.route("/predict", methods=["POST"])
def handle_query():
    text = flask.request.form['text']  # 获取用户query中的文本 例如"I love you"
    id_ = str(uuid.uuid1())  # 为query生成唯一标识
    d = {'id': id_, 'text': text}  # 绑定文本和query id
    redis_.rpush(db_key_query, json.dumps(d))  # 加入redis
    while True:
        result = redis_.get(id_)  # 获取该query的模型结果
        if result is not None:
            redis_.delete(id_)
            result_text = {'code': "200", 'data': result.decode('UTF-8')}
            break
    return flask.jsonify(result_text)  # 返回结果


if __name__ == "__main__":
    t = Thread(target=classify)
    t.start()
    app.run(debug=False, host='127.0.0.1', port=9000)