import argparse from copy import deepcopy import json import random import re from http import HTTPStatus from pathlib import Path import dashscope import requests import yaml from tqdm import tqdm def set_api_key(api_type, api_yaml_path): """设置 api key Args: api_type (str): api 类型 api_yaml_path (str): api yaml 文件路径 """ # 读取 yaml 文件 with open(api_yaml_path, "r", encoding="utf-8") as f: api_yaml = yaml.safe_load(f) # 设置 api key if api_type == "qwen": api_key = api_yaml["ali_qwen_api_key"] dashscope.api_key = api_key elif api_type == "ernie": api_key = api_yaml["baidu_ernie_api_key"] else: raise ValueError("api_type must be qwen or ernie") return api_key def call_qwen_message(content_str, model_type=dashscope.Generation.Models.qwen_turbo): try: response = dashscope.Generation.call(model_type, prompt=content_str) except Exception as e: print(f"Maybe connect error , try again : {e}") response = dashscope.Generation.call(model_type, prompt=content_str) if response.status_code == HTTPStatus.OK: print("Used token: ", response.usage) response_str = response.output.text else: print( "Request id: %s, Status code: %s, error code: %s, error message: %s" % ( response.request_id, response.status_code, response.code, response.message, ) ) response_str = "Error" return response_str def call_ernie_message(content_str, access_token): url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token={access_token}" payload = json.dumps( { "messages": [ {"role": "user", "content": content_str}, ], "disable_search": False, "enable_citation": False, } ) headers = {"Content-Type": "application/json"} response = requests.request("POST", url, headers=headers, data=payload) if response.status_code == HTTPStatus.OK: # 获取 body 中的数据 response_json = response.json() print("Used token: ", response_json["usage"]) response_str = response_json["result"] else: response_str = "Error" return response_str def format_json_from_response(func, content_str, func_args, model_name): response = func(content_str, func_args) if "```json" in response: response = re.findall(r"```json(.*)```", response, flags=re.DOTALL)[0] # 去掉导致 json 格式化失败的字符 response = response.replace("\\", "\\\\").replace("\n\n", "\n").replace("”", '"').replace("“", '"') if model_name == "qwen": # qwen 需要检查文案中是否有 " ，并替换为单引号 ' # 查找第一个 output 的字符串 output_start = response.find('"output": "') if output_start != -1: # 查找第二个 output 的字符位置 output_end = response.find("}", output_start + 1) if output_end != -1: response = list(response) # 截取第二个 output 的字符串 check_len = len(response[output_start + len('"output": "') : output_end - 10]) for idx in range(check_len): str_idx = output_start + len('"output": "') + idx if response[str_idx] == '"': response[str_idx] = "'" response = "".join(response) # 加上 strict=False 解决 decode Invalid control character format_json = json.loads(response, strict=False) return format_json, response def process_request(func, content_str, func_args, model_name): """_summary_ Args: func (_type_): _description_ content_str (_type_): _description_ func_args (str): qwen: model_type ernie: api_key Returns: _type_: _description_ """ try: format_json, response = format_json_from_response(func, content_str, func_args, model_name) except Exception as e: try: # 再试一次 print(f"\n Got error, try again <== {e} \n") if isinstance(e, json.decoder.JSONDecodeError): print(f"JSONDecodeError doc 1: {str(e.doc)} \n") format_json, response = format_json_from_response(func, content_str, func_args, model_name) except Exception as e: print(f"\n Got error <== {e} \n") if isinstance(e, json.decoder.JSONDecodeError): print(f"JSONDecodeError doc 2: {str(e.doc)} \n") with open(f"error-{model_name}.log", "a+", encoding="utf-8") as f_error: if isinstance(e, json.decoder.JSONDecodeError): f_error.write(f"JSONDecodeError doc: {str(e.doc)} \n") f_error.write(str(e)) f_error.flush() format_json = {"Error": "Error"} return format_json def gen_product_highlights(dastset_yaml_path, api_yaml_path): """根据产品的 yaml 文件生成每个产品的特点描述 Args: dastset_yaml_path (str): 数据集的 yaml 文件路径 api_yaml_path (_type_): api 的 yaml 文件路径 """ # 读取 yaml 文件 with open(dastset_yaml_path, "r", encoding="utf-8") as f: dataset_yaml = yaml.safe_load(f) set_api_key("qwen", api_yaml_path) for _, products in dataset_yaml["product_list"].items(): for product_class, product in products.items(): product_str = str(product).replace("'", "") print(f"Process: {product_str}") product_highlights = call_qwen_message( content_str=product_str, system_str="现在你精通医院里的各种事物，你帮我举例每个科室中五个细分专业治疗方法中每个细分治疗方法的六个优势或者特点，然后用python-dic的形式输出：{类名：[特点1，特点2，...]},去掉1，2的字样，除python字典外的其他都不要输出，不要有任何的警告信息", model_type=dashscope.Generation.Models.qwen_turbo, ) code_block = re.findall(r"```python(.*)```", product_highlights, flags=re.DOTALL)[0] if " = " in code_block[:20]: code_block = code_block.split(" = ")[1] products[product_class] = eval(re.findall(r"```python(.*)```", product_highlights, flags=re.DOTALL)[0]) # 保存 yaml 文件 with open(f"{dastset_yaml_path}", "w", encoding="utf-8") as f: yaml.dump(dataset_yaml, f, allow_unicode=True) def gen_dataset(dastset_yaml_path: str, api_yaml_path: str, save_json_root: Path, model_name: str, specific_name=""): # 确保文件夹存在 save_json_root.mkdir(parents=True, exist_ok=True) # 读取 yaml 文件 with open(dastset_yaml_path, "r", encoding="utf-8") as f: dataset_yaml = yaml.safe_load(f) if specific_name != "": assert ( specific_name in dataset_yaml["role_type"] ), f"{specific_name} not in dataset_yaml['role_type'] ({dataset_yaml['role_type']}), pls check dataset yaml!" # 设置 api key api_key = set_api_key(model_name, api_yaml_path) data_gen_setting = dataset_yaml["data_generation_setting"] gen_num = data_gen_setting["each_product_gen"] each_pick_hightlight = data_gen_setting["each_pick_hightlight"] each_pick_question = data_gen_setting["each_pick_question"] # qwen 配置调取的模型种类，确保有个一是最强模型 # gen_model_type = [dashscope.Generation.Models.qwen_plus] * (gen_num - 2) # gen_model_type += [dashscope.Generation.Models.qwen_max] * 2 qwen_model_type = [dashscope.Generation.Models.qwen_max] * gen_num for role_type, role_character in dataset_yaml["role_type"].items(): if specific_name != "" and role_type != specific_name: # 只生成特定人物的 print(f"specific_name = {specific_name}, skipping for {role_type}") continue gen_json = dict() save_json_path = save_json_root.joinpath(f"{model_name}_{role_type}_train.json") bk_json_path = save_json_root.joinpath(f"{model_name}_{role_type}_train.json.bk") # 加载之前已经有的 json if save_json_path.exists(): with open(save_json_path, "r", encoding="utf-8") as f: gen_json = json.load(f) # 加载成功的话，再删除备份的 json if bk_json_path.exists(): bk_json_path.unlink() # 遍历所有产品，方便进度条显示 list_product = [ product_name for _, products in dataset_yaml["product_list"].items() for _, product_name_list in products.items() for product_name in product_name_list.keys() ] # 生成人物性格 character = "、".join(role_character) pbar = tqdm(total=len(list_product)) # 遍历产品 for _, products in dataset_yaml["product_list"].items(): for _, product_name_list in products.items(): for product, hightlights in product_name_list.items(): pbar.set_description(product) if product in gen_json: # 跳过已经有的 pbar.update(1) continue gen_json.update({product: []}) # 生成数据 for idx in range(gen_num): # 随机抽取 ${each_pick_hightlight} 个产品特性 if each_pick_hightlight >= len(hightlights): # 超过打乱，增加随机性 hightlights_list = random.shuffle(hightlights) else: hightlights_list = random.sample(hightlights, each_pick_hightlight) hightlight_str = "、".join(hightlights_list) # 随机抽取 ${each_pick_question} 个提问角度 if each_pick_question >= len(dataset_yaml["customer_question_type"]): # 超过打乱，增加随机性 customer_question_type = random.shuffle(dataset_yaml["customer_question_type"]) else: customer_question_type = random.sample(dataset_yaml["customer_question_type"], each_pick_question) customer_question_str = "、".join(customer_question_type) # 商品信息 product_info_str = dataset_yaml["product_info_struct"][0].replace("{name}", product) product_info_str += dataset_yaml["product_info_struct"][1].replace("{highlights}", hightlight_str) content_str = ( data_gen_setting["dataset_gen_prompt"] .replace("{role_type}", role_type) .replace("{character}", character) .replace("{product_info}", product_info_str) .replace("{customer_question}", customer_question_str) .replace("{each_conversation_qa}", str(data_gen_setting["each_conversation_qa"])) .replace( "{dataset_json_format}", data_gen_setting["dataset_json_format"].replace("{product_info}", product_info_str), ) ) print(f"\n Resquest [ {model_name} ] {idx + 1}/{gen_num} ==> {content_str} \n") if model_name == "qwen": format_json = process_request(call_qwen_message, content_str, qwen_model_type[idx], model_name) elif model_name == "ernie": format_json = process_request(call_ernie_message, content_str, api_key, model_name) else: raise ValueError(f"model_name {model_name} not support") if "conversation" in format_json and len(format_json["conversation"]) > 0: # 第一个结果因为节省 token，需要将 system 和 input 放回去 conversation_setting = deepcopy(dataset_yaml["conversation_setting"]) system_str = ( conversation_setting["system"].replace("{role_type}", role_type).replace("{character}", character) ) input_str = conversation_setting["first_input"].replace("{product_info}", product_info_str) # 将第一个对话加入必要信息 format_json["conversation"][0] = { "system": system_str, "input": input_str, "output": format_json["conversation"][0]["output"], } else: format_json = {"Error": "Error"} print(f"\n Response [ {model_name} ] {idx + 1}/{gen_num} <== {format_json} \n") gen_json[product].append(format_json) pbar.update(1) # 备份旧的 if save_json_path.exists(): save_json_path.rename(bk_json_path) # 保存 json with open(save_json_path, "w", encoding="utf-8") as f: json.dump(gen_json, f, indent=4, ensure_ascii=False) # 如果保存成功，删掉旧的 if bk_json_path.exists(): bk_json_path.unlink() if __name__ == "__main__": # 例子：全部人物使用 Qwen api 生成数据 # cd /path/to/Streamer-Sales/dataset/gen_dataset # python gen_dataset.py qwen # 命令行输入参数 parser = argparse.ArgumentParser(description="Gen Dataset") parser.add_argument("model_name", type=str, choices=["qwen", "ernie"], help="Model name for data generation") parser.add_argument("--data_yaml", type=str, default="../../configs/conversation_cfg.yaml", help="data setting file path") parser.add_argument("--api_yaml", type=str, default="../../configs/api_cfg.yaml", help="api setting file path") parser.add_argument("--output_dir", type=str, default="./train_dataset/response", help="generation json output dir") parser.add_argument("--specific_name", type=str, default="", help="Character name for data generation") args = parser.parse_args() # 生成产品特性（可选） # gen_product_highlights(args.data_yaml, args.api_yaml) # 生成对话数据集 gen_dataset( args.data_yaml, args.api_yaml, Path(args.output_dir), model_name=args.model_name, specific_name=args.specific_name )