diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py index a20828f0..8cb2e055 100644 --- a/rdagent/core/conf.py +++ b/rdagent/core/conf.py @@ -98,6 +98,7 @@ class RDAgentSettings(BaseSettings): # factor extraction conf max_input_duplicate_factor_group: int = 300 max_output_duplicate_factor_group: int = 20 + max_kmeans_group_number: int = 40 # workspace conf workspace_path: Path = Path.cwd() / "git_ignore_folder" / "RD-Agent_workspace" diff --git a/rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py b/rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py index d34fd614..3820ab36 100644 --- a/rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py +++ b/rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py @@ -363,29 +363,49 @@ def check_factor_viability( def __check_factor_duplication_simulate_json_mode( factor_df: pd.DataFrame, ) -> list[list[str]]: - session = APIBackend().build_chat_session( - session_system_prompt=document_process_prompts["factor_duplicate_system"], - ) current_user_prompt = factor_df.to_string() - generated_duplicated_groups = [] - for _ in range(10): - extract_result_resp = session.build_chat_completion( - user_prompt=current_user_prompt, - json_mode=True, - ) - ret_dict = json.loads(extract_result_resp) - if len(ret_dict) == 0: - return generated_duplicated_groups + working_list = [factor_df] + final_list = [] + + while len(working_list) > 0: + current_df = working_list.pop(0) + if ( + APIBackend().build_messages_and_calculate_token( + user_prompt=current_df.to_string(), system_prompt=document_process_prompts["factor_duplicate_system"] + ) + > RD_AGENT_SETTINGS.chat_token_limit + ): + working_list.append(current_df.iloc[: current_df.shape[0] // 2, :]) + working_list.append(current_df.iloc[current_df.shape[0] // 2 :, :]) else: - generated_duplicated_groups.extend(ret_dict) - current_user_prompt = """Continue to extract duplicated groups. If no more duplicated group found please respond empty dict.""" + final_list.append(current_df) + + generated_duplicated_groups = [] + for current_df in final_list: + current_factor_to_string = current_df.to_string() + session = APIBackend().build_chat_session( + session_system_prompt=document_process_prompts["factor_duplicate_system"], + ) + for _ in range(10): + extract_result_resp = session.build_chat_completion( + user_prompt=current_factor_to_string, + json_mode=True, + ) + ret_dict = json.loads(extract_result_resp) + if len(ret_dict) == 0: + return generated_duplicated_groups + else: + generated_duplicated_groups.extend(ret_dict) + current_factor_to_string = """Continue to extract duplicated groups. If no more duplicated group found please respond empty dict.""" return generated_duplicated_groups def __kmeans_embeddings(embeddings: np.ndarray, k: int = 20) -> list[list[str]]: x_normalized = normalize(embeddings) + np.random.seed(42) + kmeans = KMeans( n_clusters=k, init="random", @@ -468,7 +488,7 @@ Factor variables: {variables} else: for k in range( len(full_str_list) // RD_AGENT_SETTINGS.max_input_duplicate_factor_group, - 40, + RD_AGENT_SETTINGS.max_kmeans_group_number, ): kmeans_index_group = __kmeans_embeddings(embeddings=embeddings, k=k) if len(kmeans_index_group[0]) < RD_AGENT_SETTINGS.max_input_duplicate_factor_group: @@ -479,26 +499,22 @@ Factor variables: {variables} duplication_names_list = [] - pool = mp.Pool(target_k) - result_list = [ - pool.apply_async( - __check_factor_duplication_simulate_json_mode, - (factor_df.loc[factor_name_group, :],), - ) - for factor_name_group in factor_name_groups - ] + result_list = multiprocessing_wrapper( + [ + (__check_factor_duplication_simulate_json_mode, (factor_df.loc[factor_name_group, :],)) + for factor_name_group in factor_name_groups + ], + n=RD_AGENT_SETTINGS.multi_proc_n, + ) - pool.close() - pool.join() + duplication_names_list = [] - for result in result_list: - deduplication_factor_names_list = result.get() - for deduplication_factor_names in deduplication_factor_names_list: - filter_factor_names = [ - factor_name for factor_name in set(deduplication_factor_names) if factor_name in factor_dict - ] - if len(filter_factor_names) > 1: - duplication_names_list.append(filter_factor_names) + for deduplication_factor_names_list in result_list: + filter_factor_names = [ + factor_name for factor_name in set(deduplication_factor_names_list) if factor_name in factor_dict + ] + if len(filter_factor_names) > 1: + duplication_names_list.append(filter_factor_names) return duplication_names_list @@ -509,6 +525,8 @@ def deduplicate_factors_by_llm( # noqa: C901, PLR0912 ) -> list[list[str]]: final_duplication_names_list = [] current_round_factor_dict = factor_dict + + # handle multi-round deduplication for _ in range(10): duplication_names_list = __deduplicate_factor_dict(current_round_factor_dict) @@ -524,11 +542,13 @@ def deduplicate_factors_by_llm( # noqa: C901, PLR0912 else: break + # sort the final list of duplicates by their length, largest first final_duplication_names_list = sorted(final_duplication_names_list, key=lambda x: len(x), reverse=True) - to_replace_dict = {} + to_replace_dict = {} # to map duplicates to the target factor names for duplication_names in duplication_names_list: if factor_viability_dict is not None: + # check viability of each factor in the duplicates group viability_list = [factor_viability_dict[name]["viability"] for name in duplication_names] if True not in viability_list: continue @@ -543,6 +563,7 @@ def deduplicate_factors_by_llm( # noqa: C901, PLR0912 llm_deduplicated_factor_dict = {} added_lower_name_set = set() for factor_name in factor_dict: + # only add factors that haven't been replaced and are not duplicates if factor_name not in to_replace_dict and factor_name.lower() not in added_lower_name_set: if factor_viability_dict is not None and not factor_viability_dict[factor_name]["viability"]: continue