You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
75 lines
2.2 KiB
75 lines
2.2 KiB
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
@Time : 2023/3/2 19:31
|
|
@Author :
|
|
@FileName:
|
|
@Software:
|
|
@Describe:
|
|
"""
|
|
#
|
|
# import redis
|
|
#
|
|
# redis_pool = redis.ConnectionPool(host='127.0.0.1', port=6379, password='', db=0)
|
|
# redis_conn = redis.Redis(connection_pool=redis_pool)
|
|
#
|
|
#
|
|
# name_dict = {
|
|
# 'name_4' : 'Zarten_4',
|
|
# 'name_5' : 'Zarten_5'
|
|
# }
|
|
# redis_conn.mset(name_dict)
|
|
|
|
import flask
|
|
import redis
|
|
import uuid
|
|
import json
|
|
from threading import Thread
|
|
import time
|
|
|
|
app = flask.Flask(__name__)
|
|
pool = redis.ConnectionPool(host='localhost', port=6379, max_connections=50)
|
|
redis_ = redis.Redis(connection_pool=pool, decode_responses=True)
|
|
|
|
db_key_query = 'query'
|
|
db_key_result = 'result'
|
|
batch_size = 32
|
|
|
|
|
|
|
|
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 = texts # 调用模型
|
|
time.sleep(20)
|
|
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.json['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, args=(batch_size,))
|
|
t.start()
|
|
app.run(debug=False, host='127.0.0.1', port=9000)
|