修改了单元测试的问题生成代码

2024-09-06 18:40:54 +08:00
parent c1df0d1bba
commit d12800e14e
2 changed files with 174 additions and 33 deletions
@@ -0,0 +1,55 @@
+Attribute_Prompt = (
+    "你是一个电力造价工程相关的项目经理，现在给你一些上下文信息，"
+    "你需要根据现有的上下文信息，来生成{num_questions_per_chunk}个电力造价工程相关的问题和对应的回答，"
+	"现在需要你针对数据中属性一列进行提问和回答。"
+    "问题和回答的示例应该是这种类型的，示例：'工程总投资（万元）,工程总投资（万元）是77469835.590045万元','尖峰及施工基面土石方量,尖峰及施工基面土石方量是8377.6','截止阀的编码,截止阀的编码是F01010203',"
+    "你生成的回答必须严格按照示例中的格式（'问题, 回答'），不允许有丝毫的变动。问题和回答应该在一个单引号内。"
+    "这种类似的问题和答案，生成的问题和答案必须一一对应，要符合文件里的内容，不要生成一些无关的问题，不要生成一些重复的问题，"
+    "不要生成一些过于简单的问题，不要生成一些过于复杂的问题。"
+)
+
+
+Amount_Prompt = (
+    "你是一个电力造价工程相关的项目经理，现在给你一些上下文信息，"
+    "你需要根据现有的上下文信息，来生成{num_questions_per_chunk}个电力造价工程相关的问题和对应的回答，"
+	"现在需要你针对上下文信息中的金额或者合价进行提问和回答。" 	
+    "问题和回答的示例应该是这种类型的，示例：'项目建设技术服务费的金额,项目建设技术服务费的金额是16855957065.4302','项目后评价费的费率,项目后评价费的费率是0.5','架空输电线路本体工程的金额,架空输电线路本体工程的金额是55105688268.5176','工程静态投资的金额,工程静态投资的金额是715035853336.391'"
+    "你生成的回答必须严格按照示例中的格式（'问题, 回答'），不允许有丝毫的变动。问题和回答应该在一个单引号内。"
+    "这种类似的问题和答案，生成的问题和答案必须一一对应，要符合文件里的内容，不要生成一些无关的问题，不要生成一些重复的问题，"
+    "不要生成一些过于简单的问题，不要生成一些过于复杂的问题。"
+)
+
+
+
+Units_Prompt = (
+    "你是一个电力造价工程相关的项目经理，现在给你一些上下文信息，"
+    "你需要根据现有的上下文信息，来生成{num_questions_per_chunk}个电力造价工程相关的问题和对应的回答，"
+	"现在需要你针对上下文信息来进行单位转化问题提问和回答。"  
+    "问题和回答的示例应该是这种类型的，示例：'工程总投资（万元）结果用元表示,工程总投资（万元）是774698355900.45元','本体工程（元）结果用万元表示,本体工程（元）是5490494.261046万元'"
+    "你生成的回答必须严格按照示例中的格式（'问题, 回答'），不允许有丝毫的变动。问题和回答应该在一个单引号内。"
+    "这种类似的问题和答案，生成的问题和答案必须一一对应，要符合文件里的内容，不要生成一些无关的问题，不要生成一些重复的问题，"
+    "不要生成一些过于简单的问题，不要生成一些过于复杂的问题。"
+)
+
+Name_Prompt = (
+    "你是一个电力造价工程相关的项目经理，现在给你一些上下文信息，"
+    "你需要根据现有的上下文信息，来生成{num_questions_per_chunk}个电力造价工程相关的问题和对应的回答，"
+	"现在需要你针对上下文信息中的重名问题进行提问和回答。"    
+    "问题和回答的示例应该是这种类型的，示例：'专业类型为线路的杆塔工程项目划分的合价,专业类型为线路的杆塔工程项目划分的合价是220969744.905856','专业类型为线路清理的杆塔工程项目划分的合价,电缆工程的合价是0'"
+    "你生成的回答必须严格按照示例中的格式（'问题, 回答'），不允许有丝毫的变动。问题和回答应该在一个单引号内。"
+    "这种类似的问题和答案，生成的问题和答案必须一一对应，要符合文件里的内容，不要生成一些无关的问题，不要生成一些重复的问题，"
+    "不要生成一些过于简单的问题，不要生成一些过于复杂的问题。"
+)
+
+
+All_Amount_Prompt = (
+    "你是一个电力造价工程相关的项目经理，现在给你一些上下文信息，"
+    "你需要根据现有的上下文信息，来生成{num_questions_per_chunk}个电力造价工程相关的问题和对应的回答，"
+	"现在需要你针对上下文信息中的总体金额进行提问和回答。" 	
+    "问题和回答的示例应该是这种类型的，示例：'架空输电线路本体工程的总体金额,架空输电线路本体工程的总体金额是7.706703','工程静态投资的总体金额,工程静态投资的总体金额是100'"
+    "你生成的回答必须严格按照示例中的格式（'问题, 回答'），不允许有丝毫的变动。问题和回答应该在一个单引号内。"
+    "这种类似的问题和答案，生成的问题和答案必须一一对应，要符合文件里的内容，不要生成一些无关的问题，不要生成一些重复的问题，"
+    "不要生成一些过于简单的问题，不要生成一些过于复杂的问题。"
+)
+
+
@@ -1,6 +1,11 @@
+
 from dotenv import load_dotenv
 load_dotenv()

+import json
+import sys
+
+
 from app.observability import init_observability
 from app.settings import init_settings

@@ -9,43 +14,77 @@ nest_asyncio.apply()

 from llama_index.core.node_parser import SentenceSplitter
 from llama_index.core import SimpleDirectoryReader
-
 from llama_index.core.evaluation import DatasetGenerator

-import json
+import prompts  

 init_settings()
 init_observability()

-documents = SimpleDirectoryReader("backend\data-test").load_data()
+# 读取所有文档（即所有表格）
+documents = SimpleDirectoryReader("D:/LLM_model/text2sql/zjdataai-app-test/backend/data-test").load_data()

-splitter = SentenceSplitter(chunk_size=512)
+# 定义表格名称和索引的对应关系
+table_names = {
+    "工程信息表": 0,
+    "其他费用表": 1,
+    "取费表": 2,
+    "项目划分表": 3,
+    "项目划分_费用预览表": 4,
+    "总算表": 5,
+    "工程量表": 6
+}

-# question_generator = DatasetGenerator.from_documents(documents)
-quest_prompt = (
-    "你是一个电力造价工程相关的项目经理，现在给你一些上下文信息，"
-    "你需要根据现有的上下文信息，来生成{num_questions_per_chunk}个电力造价工程相关的问题和对应的回答，"
-    "问题的实例应该是这种类型的：'人工费的费率是多少?，费率是100','前期工作管理费用的金额是多少?，金额是0',"
-    "这种类似的问题和答案，生成的问题和答案必须一一对应，要符合文件里的内容，不要生成一些无关的问题，不要生成一些重复的问题，"
-    "不要生成一些过于简单的问题，不要生成一些过于复杂的问题。"
-)
+# 定义中文提示词和Python代码中提示词名称的映射
+prompt_mapping = {
+    "普通属性": "Attribute_Prompt",
+    "金额查询": "Amount_Prompt",
+    "单位换算": "Units_Prompt",
+    "重名项目划分": "Name_Prompt",
+    "总体金额查询": "All_Amount_Prompt"
+}

+# 定义表格与其对应的查询类别
+table_prompt_mapping = {
+    "工程信息表": ["普通属性", "单位换算"],
+    "其他费用表": ["金额查询", "单位换算"],
+    "取费表": ["金额查询"],
+    "总算表": ["金额查询", "总体金额查询"],
+    "工程量表": ["普通属性", "重名项目划分"]
+}
+
+# 根据表格名称选择特定的表格
+def select_document(documents, table_name):
+    if table_name not in table_names:
+        raise ValueError(f"未找到名为 '{table_name}' 的表格")
+    index = table_names[table_name]
+    return [documents[index]]  # 返回一个包含所选表格的列表
+
+# 选择提示词
+def select_prompt(prompt_category):
+    prompt_name = prompt_mapping.get(prompt_category)
+    if not prompt_name:
+        raise ValueError(f"未找到名为 '{prompt_category}' 的提示词")
+    try:
+        return getattr(prompts, prompt_name)
+    except AttributeError:
+        raise ValueError(f"未找到提示词 '{prompt_name}' 对应的函数")
+
+# 生成问题和答案
+def generate_questions_from_document(document, quest_prompt, num_questions):
    question_generator = DatasetGenerator.from_documents(
-    documents=documents,
+        documents=document,
        question_gen_query=quest_prompt,
-    num_questions_per_chunk=5  #生成的问题数
+        num_questions_per_chunk=num_questions
    )

-eval_questions = question_generator.generate_questions_from_nodes(5)
+    eval_questions = question_generator.generate_questions_from_nodes(num_questions)
+    print(eval_questions)

-# print(eval_questions)
-
-# 处理生成的问题和答案，转换为JSON格式
    qa_pairs = []
    for qa in eval_questions:
-    # 处理可能没有 '，' 的情况
-    if '？' in qa:
-        question, answer = qa.split("？", 1)
+        if ',' in qa:
+            question, answer = qa.split(",", 1)
            qa_pairs.append({
                "question": question.strip(),
                "answer": answer.strip()
@@ -53,6 +92,53 @@ for qa in eval_questions:
        else:
            print(f"无法处理的问题和答案: {qa}")

-# 保存为JSON文件
-with open("backend/unit_test/questions_and_answers.json", "w", encoding="utf-8") as f:
-    json.dump(qa_pairs, f, ensure_ascii=False, indent=4)
+    return qa_pairs
+
+# 主函数，控制生成多个表格的问题和使用多个提示词，并将结果合并到一个文件中
+def main(documents, table_names_input, prompt_categories_input, num_questions_per_prompt):
+    if table_names_input == "all":
+        selected_tables = list(table_prompt_mapping.keys())
+    else:
+        selected_tables = table_names_input.strip('[]').split(',')
+
+    all_results = {}
+
+    for table_name in selected_tables:
+        table_name = table_name.strip()  # 去掉前后空格
+        document = select_document(documents, table_name)
+
+        if prompt_categories_input == "all":
+            selected_prompts = table_prompt_mapping[table_name]
+        else:
+            selected_prompts = prompt_categories_input.strip('[]').split(',')
+            selected_prompts = [p.strip() for p in selected_prompts]  # 去掉前后空格
+
+        for prompt_category in selected_prompts:
+            if prompt_category not in table_prompt_mapping[table_name]:
+                print(f"跳过表格 '{table_name}' 的提示词 '{prompt_category}'，因为该表中不包含该类别的信息")
+                continue
+
+            quest_prompt = select_prompt(prompt_category).format(num_questions_per_chunk=num_questions_per_prompt)
+            qa_pairs = generate_questions_from_document(document, quest_prompt, num_questions_per_prompt)
+
+            label = f"test:{table_name}_{prompt_category}"
+            all_results[label] = qa_pairs
+
+    # 自动生成输出文件名
+    output_file = "combined_test.json"
+
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(all_results, f, ensure_ascii=False, indent=4)
+
+    print(f"All questions and answers have been saved to '{output_file}'")
+
+# 获取命令行参数
+if __name__ == "__main__":
+    if len(sys.argv) != 4:
+        print("Usage: python script.py <table_names_input> <prompt_categories_input> <num_questions_per_prompt>")
+    else:
+        table_names_input = sys.argv[1]
+        prompt_categories_input = sys.argv[2]
+        num_questions_per_prompt = int(sys.argv[3])
+
+        main(documents, table_names_input, prompt_categories_input, num_questions_per_prompt)