mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 23:47:46 +00:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user