From 1e9375725481639a4840656976e4f9dbae00d5bd Mon Sep 17 00:00:00 2001 From: "majiahui@haimaqingfan.com" Date: Tue, 15 Apr 2025 14:45:31 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E6=89=B9=E9=87=8F=E4=B8=8A?= =?UTF-8?q?=E4=BC=A0=E5=87=BA=E9=94=99=E9=97=AE=E9=A2=98=EF=BC=8C=E5=B9=B6?= =?UTF-8?q?=E5=8F=91=E6=89=BF=E8=BD=BD=E5=8A=9B=E6=9B=B4=E5=BC=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 59 ++++++++++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 44 insertions(+), 15 deletions(-) diff --git a/main.py b/main.py index 87f8823..c38157d 100644 --- a/main.py +++ b/main.py @@ -15,21 +15,29 @@ from flask_cors import CORS import pandas as pd import concurrent.futures import json - +from threading import Thread +import redis app = Flask(__name__) CORS(app) app.config["JSON_AS_ASCII"] = False -openai_api_key = "token-abc123" -openai_api_base = "http://127.0.0.1:12011/v1" +pool = redis.ConnectionPool(host='localhost', port=63179, max_connections=100, db=1, password="zhicheng123*") +redis_ = redis.Redis(connection_pool=pool, decode_responses=True) + +db_key_query = 'query' +db_key_querying = 'querying' +batch_size = 32 -client = OpenAI( - api_key=openai_api_key, - base_url=openai_api_base, -) +# openai_api_key = "token-abc123" +# openai_api_base = "http://127.0.0.1:12011/v1" +# +# client = OpenAI( +# api_key=openai_api_key, +# base_url=openai_api_base, +# ) -models = client.models.list() +# models = client.models.list() # model = models.data[0].id model = "1" model_encode = SentenceTransformer('/home/majiahui/project/models-llm/bge-large-zh-v1.5') @@ -288,7 +296,6 @@ def main(question, title, top): 构造prompt ''' print("paper_list_str", paper_list_str) - 9/0 propmt_connect_ziliao_input = [] for i in db_type_list: propmt_connect_ziliao_input.append(propmt_connect_ziliao.format(i, paper_list_str)) @@ -300,6 +307,20 @@ def main(question, title, top): ''' return model_generate_stream(propmt_connect_input) +def classify(): # 调用模型,设置最大batch_size + while True: + if redis_.llen(db_key_query) == 0: # 若队列中没有元素就继续获取 + time.sleep(3) + continue + query = redis_.lpop(db_key_query).decode('UTF-8') # 获取query的text + data_dict = json.loads(query) + if data_dict["state"] == "1": + new_id = data_dict["id"] + sentence = data_dict["sentence"] + title = data_dict["title"] + df = ulit_request_file(new_id, sentence, title) + Building_vector_database(title, df) + def model_generate_stream(prompt): messages = [ @@ -339,10 +360,10 @@ def model_generate_stream(prompt): @app.route("/upload_file_check", methods=["POST"]) def upload_file_check(): print(request.remote_addr) - sentence = request.form.get('sentence') - title = request.form.get("title") - new_id = request.form.get("id") - state = request.form.get("state") + sentence = request.json['sentence'] + title = request.json["title"] + new_id = request.json["id"] + state = request.json["state"] ''' { "1": "csv", @@ -353,8 +374,14 @@ def upload_file_check(): ''' state_res = "" if state == "1": - df = ulit_request_file(new_id, sentence, title) - Building_vector_database(title, df) + # df = ulit_request_file(new_id, sentence, title) + # Building_vector_database(title, df) + redis_.rpush(db_key_query, json.dumps({ + "id": new_id, + "sentence": sentence, + "state": state, + "title": title + })) # 加入redis state_res = "上传完成" elif state == "2": delete_data(title, new_id) @@ -375,6 +402,8 @@ def search(): response = main(texts, title, top) return Response(response, mimetype='text/plain; charset=utf-8') # 返回结果 +t = Thread(target=classify) +t.start() if __name__ == "__main__": app.run(host="0.0.0.0", port=27000, threaded=True, debug=False)