From 06b249f4a39661adbc9841f35202c43b7a3943e1 Mon Sep 17 00:00:00 2001 From: "majiahui@haimaqingfan.com" Date: Mon, 5 Feb 2024 18:31:27 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E5=86=B3bug=EF=BC=9A=20=E5=A4=9A?= =?UTF-8?q?=E7=BA=BF=E7=A8=8B=E6=8A=A2=E8=B5=84=E6=BA=90=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mistral_model_predict_vllm.py | 73 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 60 insertions(+), 13 deletions(-) diff --git a/mistral_model_predict_vllm.py b/mistral_model_predict_vllm.py index 60ba66a..23e4956 100644 --- a/mistral_model_predict_vllm.py +++ b/mistral_model_predict_vllm.py @@ -1,5 +1,5 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "2" +os.environ["CUDA_VISIBLE_DEVICES"] = "1" from transformers import pipeline import redis import uuid @@ -14,20 +14,39 @@ import requests import socket -pool = redis.ConnectionPool(host='localhost', port=63179, max_connections=50,db=5, password="zhicheng123*") +pool = redis.ConnectionPool(host='localhost', port=63179, max_connections=50,db=4, password="zhicheng123*") redis_ = redis.Redis(connection_pool=pool, decode_responses=True) db_key_query = 'query' -db_key_query_articles_directory = 'query_articles_directory' +db_key_querying = 'querying' db_key_result = 'result' -batch_size = 512 +batch_size = 24 # sampling_params = SamplingParams(temperature=0.95, top_p=0.7,presence_penalty=0.9,stop="", max_tokens=4096) -sampling_params = SamplingParams(temperature=0.95, top_p=0.7,stop="", max_tokens=4096) -models_path = "/home/majiahui/project/models-llm/openbuddy-mistral-7b-v13.1-finetune-90000" +sampling_params = SamplingParams(temperature=0.95, top_p=0.7,stop="", presence_penalty=1.1, max_tokens=4096) +models_path = "/home/majiahui/project/LLaMA-Factory-main/lora_openbuddy_zephyr_paper_model_190000" llm = LLM(model=models_path, tokenizer_mode="slow") +class log: + def __init__(self): + pass + + def log(*args, **kwargs): + format = '%Y/%m/%d-%H:%M:%S' + format_h = '%Y-%m-%d' + value = time.localtime(int(time.time())) + dt = time.strftime(format, value) + dt_log_file = time.strftime(format_h, value) + log_file = 'log_file/access-%s' % dt_log_file + ".log" + if not os.path.exists(log_file): + with open(os.path.join(log_file), 'w', encoding='utf-8') as f: + print(dt, *args, file=f, **kwargs) + else: + with open(os.path.join(log_file), 'a+', encoding='utf-8') as f: + print(dt, *args, file=f, **kwargs) + + def classify(batch_size): # 调用模型，设置最大batch_size while True: texts = [] @@ -35,10 +54,28 @@ def classify(batch_size): # 调用模型，设置最大batch_size if redis_.llen(db_key_query) == 0: # 若队列中没有元素就继续获取 time.sleep(2) 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 + + # for i in range(min(redis_.llen(db_key_query), batch_size)): + while True: + query = redis_.lpop(db_key_query) # 获取query的text + if query == None: + break + + query = query.decode('UTF-8') + data_dict_path = json.loads(query) + + path = data_dict_path['path'] + with open(path, encoding='utf8') as f1: + # 加载文件的对象 + data_dict = json.load(f1) + # query_ids.append(json.loads(query)['id']) + # texts.append(json.loads(query)['text']) # 拼接若干text 为batch + query_id = data_dict['id'] + text = data_dict["text"] + query_ids.append(query_id) + texts.append(text) + if len(texts) == batch_size: + break outputs = llm.generate(texts, sampling_params) # 调用模型 generated_text_list = [""] * len(texts) @@ -48,10 +85,20 @@ def classify(batch_size): # 调用模型，设置最大batch_size generated_text = output.outputs[0].text generated_text_list[int(index)] = generated_text - for (id_, output) in zip(query_ids, generated_text_list): - res = output - redis_.set(id_, json.dumps(res)) # 将模型结果送回队列 + + return_text = {"texts": output, "probabilities": None, "status_code": 200} + load_result_path = "./new_data_logs/{}.json".format(id_) + with open(load_result_path, 'w', encoding='utf8') as f2: + # ensure_ascii=False才能输入中文，否则是Unicode字符 + # indent=2 JSON数据的缩进，美观 + json.dump(return_text, f2, ensure_ascii=False, indent=4) + redis_.set(id_, load_result_path, 86400) + # redis_.set(id_, load_result_path, 30) + redis_.srem(db_key_querying, id_) + log.log('start at', + 'query_id:{},load_result_path:{},return_text:{}'.format( + id_, load_result_path, return_text)) if __name__ == '__main__':