fix: fix some bugs in llm calling (#217)

* fix some bugs in llm calling

* fix a CI error

* fix a small bug in kmeans group number choosing

* fix a ci bug
This commit is contained in:
WinstonLiyt
2024-08-21 16:48:09 +08:00
committed by GitHub
parent e743b2161f
commit c8e8989edc
2 changed files with 56 additions and 34 deletions
+1
View File
@@ -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"
@@ -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