++

main.py

@@ -7,11 +7,13 @@ import openai
 import httpx
 import csv
 from datetime import datetime
+import jieba
 
 sys.path.append(os.path.join(os.path.abspath(__file__).split('AlphaGenerator')[0] + 'AlphaGenerator'))
 PROJECT_PATH = os.path.join(os.path.abspath(__file__).split('AlphaGenerator')[0] + 'AlphaGenerator')
 
 PREPARE_PROMPT = os.path.join(PROJECT_PATH, 'prepare_prompt')
+KEYS_TEXT = os.path.join(PREPARE_PROMPT, 'keys_text.txt')
 
 SELECT_DATA_SET_QTY = 30
 
@@ -19,15 +21,48 @@ SILICONFLOW_API_KEY = "sk-pvdiisdowmuwkrpnxsrlhxaovicqibmlljwrwwvbbdjaitdl"
 SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
 MODELS = [
     'deepseek-ai/DeepSeek-V3.2-Exp',
+    'Qwen/Qwen3-VL-235B-A22B-Instruct',
     # 'MiniMaxAI/MiniMax-M2',
     # 'zai-org/GLM-4.6',
-    'Qwen/Qwen3-VL-235B-A22B-Instruct',
     # 'inclusionAI/Ring-flash-2.0',
-    'zai-org/GLM-4.6',
-    'inclusionAI/Ling-flash-2.0',
-    'inclusionAI/Ring-flash-2.0',
+    # 'zai-org/GLM-4.6',
+    # 'inclusionAI/Ling-flash-2.0',
+    # 'inclusionAI/Ring-flash-2.0',
 ]
 
+def process_text(text):
+    filter_list = ['\n', '\t', '\r', '\b', '\f', '\v', '：', '的', '或', '10', '天', '了', '可', '是', '该', '，', ' ', '、', '让', '和', '集',
+                   '/', '日', '在', '（', '_', '-', ')', '(', '上', '距', '与', '比', '下', '及', '）', '...', '；', '%', '&', '+', ',', '.', 
+                   ':', ';', '<', '=', '>', '?', '[', ']', '|', '—', '。'
+    ]
+    
+    text_list = jieba.lcut(text)
+    results = []
+    for tl in text_list:
+        should_include = True
+        for fl in filter_list:
+            if fl == tl:
+                should_include = False
+                break
+        if should_include:
+            results.append(tl)
+
+    if results:
+        return list(set(results))
+    else:
+        return None
+
+def load_keys_text():
+    if not os.path.exists(KEYS_TEXT):
+        print(f"文件不存在: {KEYS_TEXT}")
+        exit(1)
+    with open(KEYS_TEXT, 'r', encoding='utf-8') as f:
+        text_list = [line.strip() for line in f if line.strip()]
+        if not text_list:
+            print('关键词文本无数据, 程序退出')
+            exit(1)
+
+        return process_text(';'.join(text_list))
 
 def txtFileLoader(file_path):
     if not os.path.exists(file_path):
@@ -37,16 +72,30 @@ def txtFileLoader(file_path):
         return [line.strip() for line in f if line.strip()]
 
 
-def csvFileLoader(file_path):
+def csvFileLoader(file_path, keys_text):
     if not os.path.exists(file_path):
         print(f"文件不存在: {file_path}")
         exit(1)
-    data = []
+    
+    data_dict = {}  # 使用字典来存储，以id为键
+    
     with open(file_path, 'r', encoding='utf-8') as f:
         reader = csv.reader(f)
         for row in reader:
-            data.append(row)
-    return data
+            for key in keys_text:
+                if key in row[11] or key in row[12]:
+                    item_id = row[0]
+                    # 如果id不存在，或者想要保留第一个出现的记录
+                    if item_id not in data_dict:
+                        data_dict[item_id] = {
+                            'id': item_id,
+                            'data_set_name': row[1],
+                            'description': row[2],
+                            'description_cn': row[11],
+                        }
+    
+    # 将字典的值转换为列表
+    return list(data_dict.values())
 
 
 def read_prompt(alpha_prompt_path):
@@ -169,7 +218,7 @@ def call_ai(prompt, model):
         print("AI调用失败")
 
 
-def prepare_prompt():
+def prepare_prompt(data_sets):
     prompt = ''
 
     # 读取基础提示词
@@ -186,32 +235,10 @@ def prepare_prompt():
     prompt += operator
     prompt += "\n========================= 操作符结束 =======================================\n\n"
 
-    # 读取数据字段, 数据字段数量庞大, 通过 dataset_id 分组读取, 然后每组里面随机选择 {SELECT_DATA_SET_QTY} 个
-    data_sets_path = os.path.join(PREPARE_PROMPT, "all_data_combined.csv")
-    data_sets = csvFileLoader(data_sets_path)
-
-    data_groups = {}
-    for index, data_set in enumerate(data_sets):
-        if index == 0:
-            continue
-
-        if data_set[2] not in data_groups:
-            data_groups[data_set[2]] = []
-        data_groups[data_set[2]].append({data_set[0]: data_set[1]})
-
-    selected_data_sets = []
-    for key, value in data_groups.items():
-        if len(value) < SELECT_DATA_SET_QTY:
-            selected_data_sets.extend(value)
-        else:
-            selected_data_sets.extend(random.sample(value, SELECT_DATA_SET_QTY))
-
     prompt += "========================= 数据字段开始 ======================================="
-    prompt += "注意: DataField: 后面的是数据字段, DataFieldDescription: 此字段后面的是数据字段对应的描述或使用说明, DataFieldDescription字段后面的内容是使用说明, 不是数据字段\n\n"
-    for data_set in selected_data_sets:
-        for key, value in data_set.items():
-            prompt += f"DataField: {key}\n"
-            prompt += f"DataFieldDescription: {value}\n"
+    prompt += "注意: data_set_name: 后面的是数据字段(可以使用), description: 此字段后面的是数据字段对应的描述或使用说明(不能使用), description_cn字段后面的内容是中文使用说明(不能使用)\n\n"
+    for data_set in data_sets:
+        prompt += str(data_set) + '\n'
 
     prompt += "========================= 数据字段结束 =======================================\n\n"
 
@@ -219,13 +246,34 @@ def prepare_prompt():
 
 
 def main():
+    # 将金融逻辑, 分割成标签
+    keys_text = load_keys_text()
+
+    # 分割好的标签, 搜索对应的数据集, 返回匹配到的结果
+    data_sets_path = os.path.join(PREPARE_PROMPT, "all_data_combined.csv")
+    result_data_sets = csvFileLoader(data_sets_path, keys_text)
+
+    if not result_data_sets:
+        print(f'搜索数据集为空, 程序退出')
+        exit(1)
+
+
+    data_sets = 0
+    print(f'从数据集中提取了 {len(result_data_sets)} 条数据')
+    if len(result_data_sets) > 500:
+        data_sets = random.sample(my_list, 10)
+    else:
+        data_sets = result_data_sets
+
+
+    # 组合提示词
+    prompt = prepare_prompt(data_sets)
+
+    # # 如果需要手动在页面段模型, 使用提示词, 打开这个, 将生成的提示词存到本地
+    manual_prompt(prompt)
+
     for model in MODELS:
-        prompt = prepare_prompt()
-
-        # # 如果需要手动在页面段模型, 使用提示词, 打开这个, 将生成的提示词存到本地
-        manual_prompt(prompt)
-
-        # # 如果需要使用模型, 打开这个
+        # 如果需要使用模型, 打开这个
         call_ai(prompt, model)