mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-08-02 18:07:43 +00:00
Align factor coder into new framework (#47)
* use CoSTEER as component name * rename factorimplementation to avoid confusion * rename modelimplementation * align benchmark and evolving evaluators * add scenario to evaluator init function * rename all factorimplementationknowledge in CoSTEER * remove all scenario related information in component * remove useless code --------- Co-authored-by: xuyang1 <xuyang1@microsoft.com>
This commit is contained in:
Xu Yang
@@ -0,0 +1,102 @@
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
|
||||
from rdagent.components.coder.factor_coder.config import FACTOR_IMPLEMENT_SETTINGS
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.evaluators import (
|
||||
FactorEvaluatorForCoder,
|
||||
FactorMultiEvaluator,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.evolvable_subjects import (
|
||||
FactorEvolvingItem,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.evolving_strategy import (
|
||||
FactorEvolvingStrategyWithGraph,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.knowledge_management import (
|
||||
FactorGraphKnowledgeBase,
|
||||
FactorGraphRAGStrategy,
|
||||
FactorKnowledgeBaseV1,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.factor import FactorExperiment
|
||||
from rdagent.core.evolving_agent import RAGEvoAgent
|
||||
from rdagent.core.scenario import Scenario
|
||||
from rdagent.core.task_generator import TaskGenerator
|
||||
|
||||
|
||||
class FactorCoSTEER(TaskGenerator[FactorExperiment]):
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
with_knowledge: bool = True,
|
||||
with_feedback: bool = True,
|
||||
knowledge_self_gen: bool = True,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.max_loop = FACTOR_IMPLEMENT_SETTINGS.max_loop
|
||||
self.knowledge_base_path = (
|
||||
Path(FACTOR_IMPLEMENT_SETTINGS.knowledge_base_path)
|
||||
if FACTOR_IMPLEMENT_SETTINGS.knowledge_base_path is not None
|
||||
else None
|
||||
)
|
||||
self.new_knowledge_base_path = (
|
||||
Path(FACTOR_IMPLEMENT_SETTINGS.new_knowledge_base_path)
|
||||
if FACTOR_IMPLEMENT_SETTINGS.new_knowledge_base_path is not None
|
||||
else None
|
||||
)
|
||||
self.with_knowledge = with_knowledge
|
||||
self.with_feedback = with_feedback
|
||||
self.knowledge_self_gen = knowledge_self_gen
|
||||
self.evolving_strategy = FactorEvolvingStrategyWithGraph(scen=self.scen)
|
||||
# declare the factor evaluator
|
||||
self.factor_evaluator = FactorMultiEvaluator(FactorEvaluatorForCoder(scen=self.scen), scen=self.scen)
|
||||
self.evolving_version = 2
|
||||
|
||||
def load_or_init_knowledge_base(self, former_knowledge_base_path: Path = None, component_init_list: list = []):
|
||||
if former_knowledge_base_path is not None and former_knowledge_base_path.exists():
|
||||
factor_knowledge_base = pickle.load(open(former_knowledge_base_path, "rb"))
|
||||
if self.evolving_version == 1 and not isinstance(factor_knowledge_base, FactorKnowledgeBaseV1):
|
||||
raise ValueError("The former knowledge base is not compatible with the current version")
|
||||
elif self.evolving_version == 2 and not isinstance(
|
||||
factor_knowledge_base,
|
||||
FactorGraphKnowledgeBase,
|
||||
):
|
||||
raise ValueError("The former knowledge base is not compatible with the current version")
|
||||
else:
|
||||
factor_knowledge_base = (
|
||||
FactorGraphKnowledgeBase(
|
||||
init_component_list=component_init_list,
|
||||
)
|
||||
if self.evolving_version == 2
|
||||
else FactorKnowledgeBaseV1()
|
||||
)
|
||||
return factor_knowledge_base
|
||||
|
||||
def generate(self, exp: FactorExperiment) -> FactorExperiment:
|
||||
# init knowledge base
|
||||
factor_knowledge_base = self.load_or_init_knowledge_base(
|
||||
former_knowledge_base_path=self.knowledge_base_path,
|
||||
component_init_list=[],
|
||||
)
|
||||
# init rag method
|
||||
self.rag = FactorGraphRAGStrategy(factor_knowledge_base)
|
||||
|
||||
# init intermediate items
|
||||
factor_experiment = FactorEvolvingItem(sub_tasks=exp.sub_tasks)
|
||||
|
||||
self.evolve_agent = RAGEvoAgent(max_loop=self.max_loop, evolving_strategy=self.evolving_strategy, rag=self.rag)
|
||||
|
||||
factor_experiment = self.evolve_agent.multistep_evolve(
|
||||
factor_experiment,
|
||||
self.factor_evaluator,
|
||||
with_knowledge=self.with_knowledge,
|
||||
with_feedback=self.with_feedback,
|
||||
knowledge_self_gen=self.knowledge_self_gen,
|
||||
)
|
||||
|
||||
# save new knowledge base
|
||||
if self.new_knowledge_base_path is not None:
|
||||
pickle.dump(factor_knowledge_base, open(self.new_knowledge_base_path, "wb"))
|
||||
self.knowledge_base = factor_knowledge_base
|
||||
self.latest_factor_implementations = exp.sub_tasks
|
||||
return factor_experiment
|
||||
@@ -0,0 +1,610 @@
|
||||
import io
|
||||
import json
|
||||
import re
|
||||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple
|
||||
|
||||
import pandas as pd
|
||||
from jinja2 import Environment, StrictUndefined
|
||||
|
||||
from rdagent.components.coder.factor_coder.config import FACTOR_IMPLEMENT_SETTINGS
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.evolvable_subjects import (
|
||||
FactorEvolvingItem,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.factor import FactorTask
|
||||
from rdagent.core.conf import RD_AGENT_SETTINGS
|
||||
from rdagent.core.evaluation import Evaluator
|
||||
from rdagent.core.evolving_framework import Feedback, QueriedKnowledge
|
||||
from rdagent.core.experiment import Implementation, Task
|
||||
from rdagent.core.log import RDAgentLog
|
||||
from rdagent.core.prompts import Prompts
|
||||
from rdagent.core.utils import multiprocessing_wrapper
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
evaluate_prompts = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml")
|
||||
|
||||
|
||||
class FactorEvaluator(Evaluator):
|
||||
# TODO:
|
||||
# I think we should have unified interface for all evaluates, for examples.
|
||||
# So we should adjust the interface of other factors
|
||||
@abstractmethod
|
||||
def evaluate(
|
||||
self,
|
||||
target_task: Task,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
**kwargs,
|
||||
) -> Tuple[str, object]:
|
||||
"""You can get the dataframe by
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
_, gen_df = implementation.execute()
|
||||
_, gt_df = gt_implementation.execute()
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tuple[str, object]
|
||||
- str: the text-based description of the evaluation result
|
||||
- object: a comparable metric (bool, integer, float ...) None for evaluator with only text-based result
|
||||
|
||||
"""
|
||||
raise NotImplementedError("Please implement the `evaluator` method")
|
||||
|
||||
def _get_df(self, gt_implementation: Implementation, implementation: Implementation):
|
||||
if gt_implementation is not None:
|
||||
_, gt_df = gt_implementation.execute()
|
||||
if isinstance(gt_df, pd.Series):
|
||||
gt_df = gt_df.to_frame("gt_factor")
|
||||
if isinstance(gt_df, pd.DataFrame):
|
||||
gt_df = gt_df.sort_index()
|
||||
else:
|
||||
gt_df = None
|
||||
|
||||
_, gen_df = implementation.execute()
|
||||
if isinstance(gen_df, pd.Series):
|
||||
gen_df = gen_df.to_frame("source_factor")
|
||||
if isinstance(gen_df, pd.DataFrame):
|
||||
gen_df = gen_df.sort_index()
|
||||
return gt_df, gen_df
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.__class__.__name__
|
||||
|
||||
|
||||
class FactorCodeEvaluator(FactorEvaluator):
|
||||
def evaluate(
|
||||
self,
|
||||
target_task: FactorTask,
|
||||
implementation: Implementation,
|
||||
execution_feedback: str,
|
||||
factor_value_feedback: str = "",
|
||||
gt_implementation: Implementation = None,
|
||||
**kwargs,
|
||||
):
|
||||
factor_information = target_task.get_factor_information()
|
||||
code = implementation.code
|
||||
|
||||
system_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(evaluate_prompts["evaluator_code_feedback_v1_system"])
|
||||
.render(scenario=self.scen.get_scenario_all_desc() if self.scen is not None else "No scenario description.")
|
||||
)
|
||||
|
||||
execution_feedback_to_render = execution_feedback
|
||||
for _ in range(10): # 10 times to split the content is enough
|
||||
user_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
evaluate_prompts["evaluator_code_feedback_v1_user"],
|
||||
)
|
||||
.render(
|
||||
factor_information=factor_information,
|
||||
code=code,
|
||||
execution_feedback=execution_feedback_to_render,
|
||||
factor_value_feedback=factor_value_feedback,
|
||||
gt_code=gt_implementation.code if gt_implementation else None,
|
||||
)
|
||||
)
|
||||
if (
|
||||
APIBackend().build_messages_and_calculate_token(
|
||||
user_prompt=user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
)
|
||||
> RD_AGENT_SETTINGS.chat_token_limit
|
||||
):
|
||||
execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :]
|
||||
else:
|
||||
break
|
||||
critic_response = APIBackend().build_messages_and_create_chat_completion(
|
||||
user_prompt=user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
json_mode=False,
|
||||
)
|
||||
|
||||
return critic_response, None
|
||||
|
||||
|
||||
class FactorSingleColumnEvaluator(FactorEvaluator):
|
||||
def evaluate(
|
||||
self,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
) -> Tuple[str, object]:
|
||||
_, gen_df = self._get_df(gt_implementation, implementation)
|
||||
|
||||
if len(gen_df.columns) == 1:
|
||||
return "The source dataframe has only one column which is correct.", True
|
||||
else:
|
||||
return (
|
||||
"The source dataframe has more than one column. Please check the implementation. We only evaluate the first column.",
|
||||
False,
|
||||
)
|
||||
|
||||
|
||||
class FactorOutputFormatEvaluator(FactorEvaluator):
|
||||
def evaluate(
|
||||
self,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
) -> Tuple[str, object]:
|
||||
gt_df, gen_df = self._get_df(gt_implementation, implementation)
|
||||
if gen_df is None:
|
||||
return (
|
||||
"The source dataframe is None. Skip the evaluation of the output format.",
|
||||
False,
|
||||
)
|
||||
buffer = io.StringIO()
|
||||
gen_df.info(buf=buffer)
|
||||
gen_df_info_str = buffer.getvalue()
|
||||
system_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
evaluate_prompts["evaluator_output_format_system"],
|
||||
)
|
||||
.render(scenario=self.scen.get_scenario_all_desc() if self.scen is not None else "No scenario description.")
|
||||
)
|
||||
resp = APIBackend().build_messages_and_create_chat_completion(
|
||||
user_prompt=gen_df_info_str, system_prompt=system_prompt, json_mode=True
|
||||
)
|
||||
resp_dict = json.loads(resp)
|
||||
if isinstance(resp_dict["output_format_decision"], str) and resp_dict["output_format_decision"].lower() in (
|
||||
"true",
|
||||
"false",
|
||||
):
|
||||
resp_dict["output_format_decision"] = bool(resp_dict["output_format_decision"])
|
||||
return (
|
||||
resp_dict["output_format_feedback"],
|
||||
resp_dict["output_format_decision"],
|
||||
)
|
||||
|
||||
|
||||
class FactorRowCountEvaluator(FactorEvaluator):
|
||||
def evaluate(
|
||||
self,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
) -> Tuple[str, object]:
|
||||
gt_df, gen_df = self._get_df(gt_implementation, implementation)
|
||||
|
||||
if gen_df.shape[0] == gt_df.shape[0]:
|
||||
return "Both dataframes have the same rows count.", True
|
||||
else:
|
||||
return (
|
||||
f"The source dataframe and the ground truth dataframe have different rows count. The source dataframe has {gen_df.shape[0]} rows, while the ground truth dataframe has {gt_df.shape[0]} rows. Please check the implementation.",
|
||||
False,
|
||||
)
|
||||
|
||||
|
||||
class FactorIndexEvaluator(FactorEvaluator):
|
||||
def evaluate(
|
||||
self,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
) -> Tuple[str, object]:
|
||||
gt_df, gen_df = self._get_df(gt_implementation, implementation)
|
||||
|
||||
if gen_df.index.equals(gt_df.index):
|
||||
return "Both dataframes have the same index.", True
|
||||
else:
|
||||
return (
|
||||
"The source dataframe and the ground truth dataframe have different index. Please check the implementation.",
|
||||
False,
|
||||
)
|
||||
|
||||
|
||||
class FactorMissingValuesEvaluator(FactorEvaluator):
|
||||
def evaluate(
|
||||
self,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
) -> Tuple[str, object]:
|
||||
gt_df, gen_df = self._get_df(gt_implementation, implementation)
|
||||
|
||||
if gen_df.isna().sum().sum() == gt_df.isna().sum().sum():
|
||||
return "Both dataframes have the same missing values.", True
|
||||
else:
|
||||
return (
|
||||
f"The dataframes do not have the same missing values. The source dataframe has {gen_df.isna().sum().sum()} missing values, while the ground truth dataframe has {gt_df.isna().sum().sum()} missing values. Please check the implementation.",
|
||||
False,
|
||||
)
|
||||
|
||||
|
||||
class FactorEqualValueCountEvaluator(FactorEvaluator):
|
||||
def evaluate(
|
||||
self,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
) -> Tuple[str, object]:
|
||||
gt_df, gen_df = self._get_df(gt_implementation, implementation)
|
||||
|
||||
try:
|
||||
close_values = gen_df.sub(gt_df).abs().lt(1e-6)
|
||||
result_int = close_values.astype(int)
|
||||
pos_num = result_int.sum().sum()
|
||||
acc_rate = pos_num / close_values.size
|
||||
except:
|
||||
close_values = gen_df
|
||||
if close_values.all().iloc[0]:
|
||||
return (
|
||||
"All values in the dataframes are equal within the tolerance of 1e-6.",
|
||||
acc_rate,
|
||||
)
|
||||
else:
|
||||
return (
|
||||
"Some values differ by more than the tolerance of 1e-6. Check for rounding errors or differences in the calculation methods.",
|
||||
acc_rate,
|
||||
)
|
||||
|
||||
|
||||
class FactorCorrelationEvaluator(FactorEvaluator):
|
||||
def __init__(self, hard_check: bool, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.hard_check = hard_check
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
) -> Tuple[str, object]:
|
||||
gt_df, gen_df = self._get_df(gt_implementation, implementation)
|
||||
|
||||
concat_df = pd.concat([gen_df, gt_df], axis=1)
|
||||
concat_df.columns = ["source", "gt"]
|
||||
ic = concat_df.groupby("datetime").apply(lambda df: df["source"].corr(df["gt"])).dropna().mean()
|
||||
ric = (
|
||||
concat_df.groupby("datetime")
|
||||
.apply(lambda df: df["source"].corr(df["gt"], method="spearman"))
|
||||
.dropna()
|
||||
.mean()
|
||||
)
|
||||
|
||||
if self.hard_check:
|
||||
if ic > 0.99 and ric > 0.99:
|
||||
return (
|
||||
f"The dataframes are highly correlated. The ic is {ic:.6f} and the rankic is {ric:.6f}.",
|
||||
True,
|
||||
)
|
||||
else:
|
||||
return (
|
||||
f"The dataframes are not sufficiently high correlated. The ic is {ic:.6f} and the rankic is {ric:.6f}. Investigate the factors that might be causing the discrepancies and ensure that the logic of the factor calculation is consistent.",
|
||||
False,
|
||||
)
|
||||
else:
|
||||
return f"The ic is ({ic:.6f}) and the rankic is ({ric:.6f}).", ic
|
||||
|
||||
|
||||
class FactorValueEvaluator(FactorEvaluator):
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation,
|
||||
**kwargs,
|
||||
) -> Tuple:
|
||||
conclusions = []
|
||||
|
||||
# Check if both dataframe has only one columns
|
||||
feedback_str, _ = FactorSingleColumnEvaluator(self.scen).evaluate(implementation, gt_implementation)
|
||||
conclusions.append(feedback_str)
|
||||
|
||||
# Check if the index of the dataframe is ("datetime", "instrument")
|
||||
feedback_str, _ = FactorOutputFormatEvaluator(self.scen).evaluate(implementation, gt_implementation)
|
||||
conclusions.append(feedback_str)
|
||||
|
||||
# Check if both dataframe have the same rows count
|
||||
if gt_implementation is not None:
|
||||
feedback_str, _ = FactorRowCountEvaluator(self.scen).evaluate(implementation, gt_implementation)
|
||||
conclusions.append(feedback_str)
|
||||
|
||||
feedback_str, same_index_result = FactorIndexEvaluator(self.scen).evaluate(
|
||||
implementation, gt_implementation
|
||||
)
|
||||
conclusions.append(feedback_str)
|
||||
|
||||
feedback_str, _ = FactorMissingValuesEvaluator(self.scen).evaluate(implementation, gt_implementation)
|
||||
conclusions.append(feedback_str)
|
||||
|
||||
feedback_str, equal_value_ratio_result = FactorEqualValueCountEvaluator(self.scen).evaluate(
|
||||
implementation, gt_implementation
|
||||
)
|
||||
conclusions.append(feedback_str)
|
||||
|
||||
if same_index_result:
|
||||
feedback_str, high_correlation_result = FactorCorrelationEvaluator(
|
||||
hard_check=True, scen=self.scen
|
||||
).evaluate(implementation, gt_implementation)
|
||||
else:
|
||||
high_correlation_result = False
|
||||
feedback_str = "The source dataframe and the ground truth dataframe have different index. Give up comparing the values and correlation because it's useless"
|
||||
conclusions.append(feedback_str)
|
||||
|
||||
# Combine all conclusions into a single string
|
||||
conclusion_str = "\n".join(conclusions)
|
||||
|
||||
same_value_or_high_correlation = (
|
||||
((equal_value_ratio_result > 0.99) or high_correlation_result) if gt_implementation is not None else False
|
||||
)
|
||||
return conclusion_str, same_value_or_high_correlation
|
||||
|
||||
|
||||
class FactorFinalDecisionEvaluator(Evaluator):
|
||||
def evaluate(
|
||||
self,
|
||||
target_task: FactorTask,
|
||||
execution_feedback: str,
|
||||
value_feedback: str,
|
||||
code_feedback: str,
|
||||
**kwargs,
|
||||
) -> Tuple:
|
||||
system_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(evaluate_prompts["evaluator_final_decision_v1_system"])
|
||||
.render(scenario=self.scen.get_scenario_all_desc() if self.scen is not None else "No scenario description.")
|
||||
)
|
||||
execution_feedback_to_render = execution_feedback
|
||||
|
||||
for _ in range(10): # 10 times to split the content is enough
|
||||
user_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
evaluate_prompts["evaluator_final_decision_v1_user"],
|
||||
)
|
||||
.render(
|
||||
factor_information=target_task.get_factor_information(),
|
||||
execution_feedback=execution_feedback_to_render,
|
||||
code_feedback=code_feedback,
|
||||
factor_value_feedback=(
|
||||
value_feedback
|
||||
if value_feedback is not None
|
||||
else "No Ground Truth Value provided, so no evaluation on value is performed."
|
||||
),
|
||||
)
|
||||
)
|
||||
if (
|
||||
APIBackend().build_messages_and_calculate_token(
|
||||
user_prompt=user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
)
|
||||
> RD_AGENT_SETTINGS.chat_token_limit
|
||||
):
|
||||
execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :]
|
||||
else:
|
||||
break
|
||||
|
||||
final_evaluation_dict = json.loads(
|
||||
APIBackend().build_messages_and_create_chat_completion(
|
||||
user_prompt=user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
json_mode=True,
|
||||
),
|
||||
)
|
||||
if isinstance(final_evaluation_dict["final_decision"], str) and final_evaluation_dict[
|
||||
"final_decision"
|
||||
].lower() in ("true", "false"):
|
||||
final_evaluation_dict["final_decision"] = bool(final_evaluation_dict["final_decision"])
|
||||
return (
|
||||
final_evaluation_dict["final_decision"],
|
||||
final_evaluation_dict["final_feedback"],
|
||||
)
|
||||
|
||||
|
||||
class FactorSingleFeedback:
|
||||
"""This class is a feedback to single implementation which is generated from an evaluator."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
execution_feedback: str = None,
|
||||
value_generated_flag: bool = False,
|
||||
code_feedback: str = None,
|
||||
factor_value_feedback: str = None,
|
||||
final_decision: bool = None,
|
||||
final_feedback: str = None,
|
||||
final_decision_based_on_gt: bool = None,
|
||||
) -> None:
|
||||
self.execution_feedback = execution_feedback
|
||||
self.value_generated_flag = value_generated_flag
|
||||
self.code_feedback = code_feedback
|
||||
self.factor_value_feedback = factor_value_feedback
|
||||
self.final_decision = final_decision
|
||||
self.final_feedback = final_feedback
|
||||
self.final_decision_based_on_gt = final_decision_based_on_gt
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"""------------------Factor Execution Feedback------------------
|
||||
{self.execution_feedback}
|
||||
------------------Factor Code Feedback------------------
|
||||
{self.code_feedback}
|
||||
------------------Factor Value Feedback------------------
|
||||
{self.factor_value_feedback}
|
||||
------------------Factor Final Feedback------------------
|
||||
{self.final_feedback}
|
||||
------------------Factor Final Decision------------------
|
||||
This implementation is {'SUCCESS' if self.final_decision else 'FAIL'}.
|
||||
"""
|
||||
|
||||
|
||||
class FactorMultiFeedback(
|
||||
Feedback,
|
||||
List[FactorSingleFeedback],
|
||||
):
|
||||
"""Feedback contains a list, each element is the corresponding feedback for each factor implementation."""
|
||||
|
||||
|
||||
class FactorEvaluatorForCoder(FactorEvaluator):
|
||||
"""This class is the v1 version of evaluator for a single factor implementation.
|
||||
It calls several evaluators in share modules to evaluate the factor implementation.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.value_evaluator = FactorValueEvaluator(self.scen)
|
||||
self.code_evaluator = FactorCodeEvaluator(self.scen)
|
||||
self.final_decision_evaluator = FactorFinalDecisionEvaluator(self.scen)
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
target_task: FactorTask,
|
||||
implementation: Implementation,
|
||||
gt_implementation: Implementation = None,
|
||||
queried_knowledge: QueriedKnowledge = None,
|
||||
**kwargs,
|
||||
) -> FactorSingleFeedback:
|
||||
if implementation is None:
|
||||
return None
|
||||
|
||||
target_task_information = target_task.get_factor_information()
|
||||
if (
|
||||
queried_knowledge is not None
|
||||
and target_task_information in queried_knowledge.success_task_to_knowledge_dict
|
||||
):
|
||||
return queried_knowledge.success_task_to_knowledge_dict[target_task_information].feedback
|
||||
elif queried_knowledge is not None and target_task_information in queried_knowledge.failed_task_info_set:
|
||||
return FactorSingleFeedback(
|
||||
execution_feedback="This task has failed too many times, skip implementation.",
|
||||
value_generated_flag=False,
|
||||
code_feedback="This task has failed too many times, skip code evaluation.",
|
||||
factor_value_feedback="This task has failed too many times, skip value evaluation.",
|
||||
final_decision=False,
|
||||
final_feedback="This task has failed too many times, skip final decision evaluation.",
|
||||
final_decision_based_on_gt=False,
|
||||
)
|
||||
else:
|
||||
factor_feedback = FactorSingleFeedback()
|
||||
|
||||
# 1. Get factor execution feedback to generated implementation and remove the long list of numbers in execution feedback
|
||||
(
|
||||
execution_feedback,
|
||||
gen_df,
|
||||
) = implementation.execute()
|
||||
|
||||
execution_feedback = re.sub(r"(?<=\D)(,\s+-?\d+\.\d+){50,}(?=\D)", ", ", execution_feedback)
|
||||
factor_feedback.execution_feedback = "\n".join(
|
||||
[line for line in execution_feedback.split("\n") if "warning" not in line.lower()]
|
||||
)
|
||||
|
||||
# 2. Get factor value feedback
|
||||
if gen_df is None:
|
||||
factor_feedback.factor_value_feedback = "No factor value generated, skip value evaluation."
|
||||
factor_feedback.value_generated_flag = False
|
||||
same_value_or_high_correlation = None
|
||||
else:
|
||||
factor_feedback.value_generated_flag = True
|
||||
(
|
||||
factor_feedback.factor_value_feedback,
|
||||
same_value_or_high_correlation,
|
||||
) = self.value_evaluator.evaluate(implementation=implementation, gt_implementation=gt_implementation)
|
||||
|
||||
factor_feedback.final_decision_based_on_gt = gt_implementation is not None
|
||||
|
||||
if same_value_or_high_correlation is not None and same_value_or_high_correlation is True:
|
||||
# To avoid confusion, when same_value_or_high_correlation is True, we do not need code feedback
|
||||
factor_feedback.code_feedback = "Final decision is True and there are no code critics."
|
||||
factor_feedback.final_decision = same_value_or_high_correlation
|
||||
factor_feedback.final_feedback = "Value evaluation passed, skip final decision evaluation."
|
||||
else:
|
||||
factor_feedback.code_feedback, _ = self.code_evaluator.evaluate(
|
||||
target_task=target_task,
|
||||
implementation=implementation,
|
||||
execution_feedback=factor_feedback.execution_feedback,
|
||||
value_feedback=factor_feedback.factor_value_feedback,
|
||||
gt_implementation=gt_implementation,
|
||||
)
|
||||
(
|
||||
factor_feedback.final_decision,
|
||||
factor_feedback.final_feedback,
|
||||
) = self.final_decision_evaluator.evaluate(
|
||||
target_task=target_task,
|
||||
execution_feedback=factor_feedback.execution_feedback,
|
||||
value_feedback=factor_feedback.factor_value_feedback,
|
||||
code_feedback=factor_feedback.code_feedback,
|
||||
)
|
||||
return factor_feedback
|
||||
|
||||
|
||||
class FactorMultiEvaluator(Evaluator):
|
||||
def __init__(self, single_evaluator, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.single_factor_implementation_evaluator = single_evaluator
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
evo: FactorEvolvingItem,
|
||||
queried_knowledge: QueriedKnowledge = None,
|
||||
**kwargs,
|
||||
) -> FactorMultiFeedback:
|
||||
multi_implementation_feedback = FactorMultiFeedback()
|
||||
|
||||
# for index in range(len(evo.sub_tasks)):
|
||||
# corresponding_implementation = evo.sub_implementations[index]
|
||||
# corresponding_gt_implementation = (
|
||||
# evo.sub_gt_implementations[index] if evo.sub_gt_implementations is not None else None
|
||||
# )
|
||||
|
||||
# multi_implementation_feedback.append(
|
||||
# self.single_factor_implementation_evaluator.evaluate(
|
||||
# target_task=evo.sub_tasks[index],
|
||||
# implementation=corresponding_implementation,
|
||||
# gt_implementation=corresponding_gt_implementation,
|
||||
# queried_knowledge=queried_knowledge,
|
||||
# )
|
||||
# )
|
||||
|
||||
calls = []
|
||||
for index in range(len(evo.sub_tasks)):
|
||||
corresponding_implementation = evo.sub_implementations[index]
|
||||
corresponding_gt_implementation = (
|
||||
evo.sub_gt_implementations[index] if evo.sub_gt_implementations is not None else None
|
||||
)
|
||||
calls.append(
|
||||
(
|
||||
self.single_factor_implementation_evaluator.evaluate,
|
||||
(
|
||||
evo.sub_tasks[index],
|
||||
corresponding_implementation,
|
||||
corresponding_gt_implementation,
|
||||
queried_knowledge,
|
||||
),
|
||||
),
|
||||
)
|
||||
multi_implementation_feedback = multiprocessing_wrapper(calls, n=FACTOR_IMPLEMENT_SETTINGS.evo_multi_proc_n)
|
||||
|
||||
final_decision = [
|
||||
None if single_feedback is None else single_feedback.final_decision
|
||||
for single_feedback in multi_implementation_feedback
|
||||
]
|
||||
RDAgentLog().info(f"Final decisions: {final_decision} True count: {final_decision.count(True)}")
|
||||
|
||||
return multi_implementation_feedback
|
||||
|
||||
|
||||
# TODO:
|
||||
def shorten_prompt(tpl: str, render_kwargs: dict, shorten_key: str, max_trail: int = 10) -> str:
|
||||
"""When the prompt is too long. We have to shorten it.
|
||||
But we should not truncate the prompt directly, so we should find the key we want to shorten and then shorten it.
|
||||
"""
|
||||
# TODO: this should replace most of code in
|
||||
# - FactorFinalDecisionEvaluator.evaluate
|
||||
# - FactorCodeEvaluator.evaluate
|
||||
@@ -0,0 +1,30 @@
|
||||
from rdagent.components.coder.factor_coder.factor import (
|
||||
FactorExperiment,
|
||||
FactorTask,
|
||||
FileBasedFactorImplementation,
|
||||
)
|
||||
from rdagent.core.evolving_framework import EvolvableSubjects
|
||||
from rdagent.core.log import RDAgentLog
|
||||
|
||||
|
||||
class FactorEvolvingItem(FactorExperiment, EvolvableSubjects):
|
||||
"""
|
||||
Intermediate item of factor implementation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sub_tasks: list[FactorTask],
|
||||
sub_gt_implementations: list[FileBasedFactorImplementation] = None,
|
||||
):
|
||||
FactorExperiment.__init__(self, sub_tasks=sub_tasks)
|
||||
self.corresponding_selection: list = None
|
||||
if sub_gt_implementations is not None and len(
|
||||
sub_gt_implementations,
|
||||
) != len(self.sub_tasks):
|
||||
self.sub_gt_implementations = None
|
||||
RDAgentLog().warning(
|
||||
"The length of sub_gt_implementations is not equal to the length of sub_tasks, set sub_gt_implementations to None",
|
||||
)
|
||||
else:
|
||||
self.sub_gt_implementations = sub_gt_implementations
|
||||
@@ -0,0 +1,336 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from abc import abstractmethod
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from jinja2 import Environment, StrictUndefined
|
||||
|
||||
from rdagent.components.coder.factor_coder.config import FACTOR_IMPLEMENT_SETTINGS
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.evolvable_subjects import (
|
||||
FactorEvolvingItem,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.scheduler import (
|
||||
LLMSelect,
|
||||
RandomSelect,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.factor import (
|
||||
FactorTask,
|
||||
FileBasedFactorImplementation,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.utils import get_data_folder_intro
|
||||
from rdagent.core.conf import RD_AGENT_SETTINGS
|
||||
from rdagent.core.evolving_framework import EvolvingStrategy, QueriedKnowledge
|
||||
from rdagent.core.experiment import Implementation
|
||||
from rdagent.core.prompts import Prompts
|
||||
from rdagent.core.utils import multiprocessing_wrapper
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.knowledge_management import (
|
||||
FactorQueriedKnowledge,
|
||||
FactorQueriedKnowledgeV1,
|
||||
)
|
||||
|
||||
implement_prompts = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml")
|
||||
|
||||
|
||||
class MultiProcessEvolvingStrategy(EvolvingStrategy):
|
||||
@abstractmethod
|
||||
def implement_one_factor(
|
||||
self,
|
||||
target_task: FactorTask,
|
||||
queried_knowledge: QueriedKnowledge = None,
|
||||
) -> Implementation:
|
||||
raise NotImplementedError
|
||||
|
||||
def evolve(
|
||||
self,
|
||||
*,
|
||||
evo: FactorEvolvingItem,
|
||||
queried_knowledge: FactorQueriedKnowledge | None = None,
|
||||
**kwargs,
|
||||
) -> FactorEvolvingItem:
|
||||
self.num_loop += 1
|
||||
new_evo = deepcopy(evo)
|
||||
|
||||
# 1.找出需要evolve的factor
|
||||
to_be_finished_task_index = []
|
||||
for index, target_factor_task in enumerate(new_evo.sub_tasks):
|
||||
target_factor_task_desc = target_factor_task.get_factor_information()
|
||||
if target_factor_task_desc in queried_knowledge.success_task_to_knowledge_dict:
|
||||
new_evo.sub_implementations[index] = queried_knowledge.success_task_to_knowledge_dict[
|
||||
target_factor_task_desc
|
||||
].implementation
|
||||
elif (
|
||||
target_factor_task_desc not in queried_knowledge.success_task_to_knowledge_dict
|
||||
and target_factor_task_desc not in queried_knowledge.failed_task_info_set
|
||||
):
|
||||
to_be_finished_task_index.append(index)
|
||||
|
||||
# 2. 选择selection方法
|
||||
# if the number of factors to be implemented is larger than the limit, we need to select some of them
|
||||
if FACTOR_IMPLEMENT_SETTINGS.select_ratio < 1:
|
||||
# if the number of loops is equal to the select_loop, we need to select some of them
|
||||
implementation_factors_per_round = int(
|
||||
FACTOR_IMPLEMENT_SETTINGS.select_ratio * len(to_be_finished_task_index)
|
||||
)
|
||||
if FACTOR_IMPLEMENT_SETTINGS.select_method == "random":
|
||||
to_be_finished_task_index = RandomSelect(
|
||||
to_be_finished_task_index,
|
||||
implementation_factors_per_round,
|
||||
)
|
||||
|
||||
if FACTOR_IMPLEMENT_SETTINGS.select_method == "scheduler":
|
||||
to_be_finished_task_index = LLMSelect(
|
||||
to_be_finished_task_index,
|
||||
implementation_factors_per_round,
|
||||
new_evo,
|
||||
queried_knowledge.former_traces,
|
||||
self.scen,
|
||||
)
|
||||
|
||||
result = multiprocessing_wrapper(
|
||||
[
|
||||
(self.implement_one_factor, (new_evo.sub_tasks[target_index], queried_knowledge))
|
||||
for target_index in to_be_finished_task_index
|
||||
],
|
||||
n=FACTOR_IMPLEMENT_SETTINGS.evo_multi_proc_n,
|
||||
)
|
||||
|
||||
for index, target_index in enumerate(to_be_finished_task_index):
|
||||
new_evo.sub_implementations[target_index] = result[index]
|
||||
|
||||
# for target_index in to_be_finished_task_index:
|
||||
# new_evo.sub_implementations[target_index] = self.implement_one_factor(
|
||||
# new_evo.sub_tasks[target_index], queried_knowledge
|
||||
# )
|
||||
|
||||
new_evo.corresponding_selection = to_be_finished_task_index
|
||||
|
||||
return new_evo
|
||||
|
||||
|
||||
class FactorEvolvingStrategy(MultiProcessEvolvingStrategy):
|
||||
def implement_one_factor(
|
||||
self,
|
||||
target_task: FactorTask,
|
||||
queried_knowledge: FactorQueriedKnowledgeV1 = None,
|
||||
) -> Implementation:
|
||||
factor_information_str = target_task.get_factor_information()
|
||||
|
||||
if queried_knowledge is not None and factor_information_str in queried_knowledge.success_task_to_knowledge_dict:
|
||||
return queried_knowledge.success_task_to_knowledge_dict[factor_information_str].implementation
|
||||
elif queried_knowledge is not None and factor_information_str in queried_knowledge.failed_task_info_set:
|
||||
return None
|
||||
else:
|
||||
queried_similar_successful_knowledge = (
|
||||
queried_knowledge.working_task_to_similar_successful_knowledge_dict[factor_information_str]
|
||||
if queried_knowledge is not None
|
||||
else []
|
||||
)
|
||||
queried_former_failed_knowledge = (
|
||||
queried_knowledge.working_task_to_former_failed_knowledge_dict[factor_information_str]
|
||||
if queried_knowledge is not None
|
||||
else []
|
||||
)
|
||||
|
||||
queried_former_failed_knowledge_to_render = queried_former_failed_knowledge
|
||||
|
||||
system_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
implement_prompts["evolving_strategy_factor_implementation_v1_system"],
|
||||
)
|
||||
.render(
|
||||
data_info=get_data_folder_intro(),
|
||||
queried_former_failed_knowledge=queried_former_failed_knowledge_to_render,
|
||||
)
|
||||
)
|
||||
session = APIBackend(use_chat_cache=False).build_chat_session(
|
||||
session_system_prompt=system_prompt,
|
||||
)
|
||||
|
||||
queried_similar_successful_knowledge_to_render = queried_similar_successful_knowledge
|
||||
while True:
|
||||
user_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
implement_prompts["evolving_strategy_factor_implementation_v1_user"],
|
||||
)
|
||||
.render(
|
||||
factor_information_str=factor_information_str,
|
||||
queried_similar_successful_knowledge=queried_similar_successful_knowledge_to_render,
|
||||
)
|
||||
.strip("\n")
|
||||
)
|
||||
if (
|
||||
session.build_chat_completion_message_and_calculate_token(
|
||||
user_prompt,
|
||||
)
|
||||
< RD_AGENT_SETTINGS.chat_token_limit
|
||||
):
|
||||
break
|
||||
elif len(queried_former_failed_knowledge_to_render) > 1:
|
||||
queried_former_failed_knowledge_to_render = queried_former_failed_knowledge_to_render[1:]
|
||||
elif len(queried_similar_successful_knowledge_to_render) > 1:
|
||||
queried_similar_successful_knowledge_to_render = queried_similar_successful_knowledge_to_render[1:]
|
||||
|
||||
code = json.loads(
|
||||
session.build_chat_completion(
|
||||
user_prompt=user_prompt,
|
||||
json_mode=True,
|
||||
),
|
||||
)["code"]
|
||||
# ast.parse(code)
|
||||
factor_implementation = FileBasedFactorImplementation(
|
||||
target_task,
|
||||
code,
|
||||
)
|
||||
|
||||
return factor_implementation
|
||||
|
||||
|
||||
class FactorEvolvingStrategyWithGraph(MultiProcessEvolvingStrategy):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.num_loop = 0
|
||||
self.haveSelected = False
|
||||
|
||||
def implement_one_factor(
|
||||
self,
|
||||
target_task: FactorTask,
|
||||
queried_knowledge,
|
||||
) -> Implementation:
|
||||
error_summary = FACTOR_IMPLEMENT_SETTINGS.v2_error_summary
|
||||
# 1. 提取因子的背景信息
|
||||
target_factor_task_information = target_task.get_factor_information()
|
||||
|
||||
# 2. 检查该因子是否需要继续做（是否已经作对，是否做错太多）
|
||||
if (
|
||||
queried_knowledge is not None
|
||||
and target_factor_task_information in queried_knowledge.success_task_to_knowledge_dict
|
||||
):
|
||||
return queried_knowledge.success_task_to_knowledge_dict[target_factor_task_information].implementation
|
||||
elif queried_knowledge is not None and target_factor_task_information in queried_knowledge.failed_task_info_set:
|
||||
return None
|
||||
else:
|
||||
# 3. 取出knowledge里面的经验数据（similar success、similar error、former_trace）
|
||||
queried_similar_component_knowledge = (
|
||||
queried_knowledge.component_with_success_task[target_factor_task_information]
|
||||
if queried_knowledge is not None
|
||||
else []
|
||||
) # A list, [success task implement knowledge]
|
||||
|
||||
queried_similar_error_knowledge = (
|
||||
queried_knowledge.error_with_success_task[target_factor_task_information]
|
||||
if queried_knowledge is not None
|
||||
else {}
|
||||
) # A dict, {{error_type:[[error_imp_knowledge, success_imp_knowledge],...]},...}
|
||||
|
||||
queried_former_failed_knowledge = (
|
||||
queried_knowledge.former_traces[target_factor_task_information] if queried_knowledge is not None else []
|
||||
)
|
||||
|
||||
queried_former_failed_knowledge_to_render = queried_former_failed_knowledge
|
||||
|
||||
system_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
implement_prompts["evolving_strategy_factor_implementation_v1_system"],
|
||||
)
|
||||
.render(
|
||||
scenario=self.scen.get_scenario_all_desc(),
|
||||
queried_former_failed_knowledge=queried_former_failed_knowledge_to_render,
|
||||
)
|
||||
)
|
||||
|
||||
session = APIBackend(use_chat_cache=False).build_chat_session(
|
||||
session_system_prompt=system_prompt,
|
||||
)
|
||||
|
||||
queried_similar_component_knowledge_to_render = queried_similar_component_knowledge
|
||||
queried_similar_error_knowledge_to_render = queried_similar_error_knowledge
|
||||
error_summary_critics = ""
|
||||
# 动态地防止prompt超长
|
||||
while True:
|
||||
# 总结error（可选）
|
||||
if (
|
||||
error_summary
|
||||
and len(queried_similar_error_knowledge_to_render) != 0
|
||||
and len(queried_former_failed_knowledge_to_render) != 0
|
||||
):
|
||||
error_summary_system_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(implement_prompts["evolving_strategy_error_summary_v2_system"])
|
||||
.render(
|
||||
factor_information_str=target_factor_task_information,
|
||||
code_and_feedback=queried_former_failed_knowledge_to_render[
|
||||
-1
|
||||
].get_implementation_and_feedback_str(),
|
||||
)
|
||||
.strip("\n")
|
||||
)
|
||||
session_summary = APIBackend(use_chat_cache=False).build_chat_session(
|
||||
session_system_prompt=error_summary_system_prompt,
|
||||
)
|
||||
while True:
|
||||
error_summary_user_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(implement_prompts["evolving_strategy_error_summary_v2_user"])
|
||||
.render(
|
||||
queried_similar_component_knowledge=queried_similar_component_knowledge_to_render,
|
||||
)
|
||||
.strip("\n")
|
||||
)
|
||||
if (
|
||||
session_summary.build_chat_completion_message_and_calculate_token(error_summary_user_prompt)
|
||||
< RD_AGENT_SETTINGS.chat_token_limit
|
||||
):
|
||||
break
|
||||
elif len(queried_similar_error_knowledge_to_render) > 0:
|
||||
queried_similar_error_knowledge_to_render = queried_similar_error_knowledge_to_render[:-1]
|
||||
error_summary_critics = session_summary.build_chat_completion(
|
||||
user_prompt=error_summary_user_prompt,
|
||||
json_mode=False,
|
||||
)
|
||||
# 构建user_prompt。开始写代码
|
||||
user_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
implement_prompts["evolving_strategy_factor_implementation_v2_user"],
|
||||
)
|
||||
.render(
|
||||
factor_information_str=target_factor_task_information,
|
||||
queried_similar_component_knowledge=queried_similar_component_knowledge_to_render,
|
||||
queried_similar_error_knowledge=queried_similar_error_knowledge_to_render,
|
||||
error_summary=error_summary,
|
||||
error_summary_critics=error_summary_critics,
|
||||
)
|
||||
.strip("\n")
|
||||
)
|
||||
if (
|
||||
session.build_chat_completion_message_and_calculate_token(
|
||||
user_prompt,
|
||||
)
|
||||
< RD_AGENT_SETTINGS.chat_token_limit
|
||||
):
|
||||
break
|
||||
elif len(queried_former_failed_knowledge_to_render) > 1:
|
||||
queried_former_failed_knowledge_to_render = queried_former_failed_knowledge_to_render[1:]
|
||||
elif len(queried_similar_component_knowledge_to_render) > len(
|
||||
queried_similar_error_knowledge_to_render,
|
||||
):
|
||||
queried_similar_component_knowledge_to_render = queried_similar_component_knowledge_to_render[:-1]
|
||||
elif len(queried_similar_error_knowledge_to_render) > 0:
|
||||
queried_similar_error_knowledge_to_render = queried_similar_error_knowledge_to_render[:-1]
|
||||
|
||||
response = session.build_chat_completion(
|
||||
user_prompt=user_prompt,
|
||||
json_mode=True,
|
||||
)
|
||||
code = json.loads(response)["code"]
|
||||
factor_implementation = FileBasedFactorImplementation(target_task, code)
|
||||
return factor_implementation
|
||||
@@ -0,0 +1,912 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import random
|
||||
import re
|
||||
from itertools import combinations
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
from jinja2 import Environment, StrictUndefined
|
||||
|
||||
from rdagent.components.coder.factor_coder.config import FACTOR_IMPLEMENT_SETTINGS
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.evaluators import (
|
||||
FactorSingleFeedback,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.factor import FactorTask
|
||||
from rdagent.components.knowledge_management.graph import (
|
||||
UndirectedGraph,
|
||||
UndirectedNode,
|
||||
)
|
||||
from rdagent.core.evolving_framework import (
|
||||
EvolvableSubjects,
|
||||
EvoStep,
|
||||
Knowledge,
|
||||
KnowledgeBase,
|
||||
QueriedKnowledge,
|
||||
RAGStrategy,
|
||||
)
|
||||
from rdagent.core.experiment import Implementation
|
||||
from rdagent.core.log import RDAgentLog
|
||||
from rdagent.core.prompts import Prompts
|
||||
from rdagent.oai.llm_utils import (
|
||||
APIBackend,
|
||||
calculate_embedding_distance_between_str_list,
|
||||
)
|
||||
|
||||
|
||||
class FactorKnowledge(Knowledge):
|
||||
def __init__(
|
||||
self,
|
||||
target_task: FactorTask,
|
||||
implementation: Implementation,
|
||||
feedback: FactorSingleFeedback,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize a FactorKnowledge object. The FactorKnowledge object is used to store a factor implementation without the ground truth code and value.
|
||||
|
||||
Args:
|
||||
factor (Factor): The factor object associated with the KnowledgeManagement.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
self.target_task = target_task
|
||||
self.implementation = implementation
|
||||
self.feedback = feedback
|
||||
|
||||
def get_implementation_and_feedback_str(self) -> str:
|
||||
return f"""------------------Factor implementation code:------------------
|
||||
{self.implementation.code}
|
||||
------------------Factor implementation feedback:------------------
|
||||
{self.feedback!s}
|
||||
"""
|
||||
|
||||
|
||||
class FactorQueriedKnowledge(QueriedKnowledge):
|
||||
def __init__(self, success_task_to_knowledge_dict: dict = {}, failed_task_info_set: set = set()) -> None:
|
||||
self.success_task_to_knowledge_dict = success_task_to_knowledge_dict
|
||||
self.failed_task_info_set = failed_task_info_set
|
||||
|
||||
|
||||
class FactorKnowledgeBaseV1(KnowledgeBase):
|
||||
def __init__(self) -> None:
|
||||
self.implementation_trace: dict[str, FactorKnowledge] = dict()
|
||||
self.success_task_info_set: set[str] = set()
|
||||
|
||||
self.task_to_embedding = dict()
|
||||
|
||||
def query(self) -> QueriedKnowledge | None:
|
||||
"""
|
||||
Query the knowledge base to get the queried knowledge. So far is handled in RAG strategy.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class FactorQueriedKnowledgeV1(FactorQueriedKnowledge):
|
||||
def __init__(self) -> None:
|
||||
self.working_task_to_former_failed_knowledge_dict = dict()
|
||||
self.working_task_to_similar_successful_knowledge_dict = dict()
|
||||
super().__init__()
|
||||
|
||||
|
||||
class FactorRAGStrategyV1(RAGStrategy):
|
||||
def __init__(self, knowledgebase: FactorKnowledgeBaseV1) -> None:
|
||||
super().__init__(knowledgebase)
|
||||
self.current_generated_trace_count = 0
|
||||
|
||||
def generate_knowledge(
|
||||
self,
|
||||
evolving_trace: list[EvoStep],
|
||||
*,
|
||||
return_knowledge: bool = False,
|
||||
) -> Knowledge | None:
|
||||
if len(evolving_trace) == self.current_generated_trace_count:
|
||||
return
|
||||
else:
|
||||
for trace_index in range(
|
||||
self.current_generated_trace_count,
|
||||
len(evolving_trace),
|
||||
):
|
||||
evo_step = evolving_trace[trace_index]
|
||||
implementations = evo_step.evolvable_subjects
|
||||
feedback = evo_step.feedback
|
||||
for task_index in range(len(implementations.sub_tasks)):
|
||||
target_task = implementations.sub_tasks[task_index]
|
||||
target_task_information = target_task.get_factor_information()
|
||||
implementation = implementations.sub_implementations[task_index]
|
||||
single_feedback = feedback[task_index]
|
||||
if single_feedback is None:
|
||||
continue
|
||||
single_knowledge = FactorKnowledge(
|
||||
target_task=target_task,
|
||||
implementation=implementation,
|
||||
feedback=single_feedback,
|
||||
)
|
||||
if target_task_information not in self.knowledgebase.success_task_info_set:
|
||||
self.knowledgebase.implementation_trace.setdefault(
|
||||
target_task_information,
|
||||
[],
|
||||
).append(single_knowledge)
|
||||
|
||||
if single_feedback.final_decision == True:
|
||||
self.knowledgebase.success_task_info_set.add(
|
||||
target_task_information,
|
||||
)
|
||||
self.current_generated_trace_count = len(evolving_trace)
|
||||
|
||||
def query(
|
||||
self,
|
||||
evo: EvolvableSubjects,
|
||||
evolving_trace: list[EvoStep],
|
||||
) -> QueriedKnowledge | None:
|
||||
v1_query_former_trace_limit = FACTOR_IMPLEMENT_SETTINGS.v1_query_former_trace_limit
|
||||
v1_query_similar_success_limit = FACTOR_IMPLEMENT_SETTINGS.v1_query_similar_success_limit
|
||||
fail_task_trial_limit = FACTOR_IMPLEMENT_SETTINGS.fail_task_trial_limit
|
||||
|
||||
queried_knowledge = FactorQueriedKnowledgeV1()
|
||||
for target_factor_task in evo.sub_tasks:
|
||||
target_factor_task_information = target_factor_task.get_factor_information()
|
||||
if target_factor_task_information in self.knowledgebase.success_task_info_set:
|
||||
queried_knowledge.success_task_to_knowledge_dict[target_factor_task_information] = (
|
||||
self.knowledgebase.implementation_trace[target_factor_task_information][-1]
|
||||
)
|
||||
elif (
|
||||
len(
|
||||
self.knowledgebase.implementation_trace.setdefault(
|
||||
target_factor_task_information,
|
||||
[],
|
||||
),
|
||||
)
|
||||
>= fail_task_trial_limit
|
||||
):
|
||||
queried_knowledge.failed_task_info_set.add(target_factor_task_information)
|
||||
else:
|
||||
queried_knowledge.working_task_to_former_failed_knowledge_dict[target_factor_task_information] = (
|
||||
self.knowledgebase.implementation_trace.setdefault(
|
||||
target_factor_task_information,
|
||||
[],
|
||||
)[-v1_query_former_trace_limit:]
|
||||
)
|
||||
|
||||
knowledge_base_success_task_list = list(
|
||||
self.knowledgebase.success_task_info_set,
|
||||
)
|
||||
similarity = calculate_embedding_distance_between_str_list(
|
||||
[target_factor_task_information],
|
||||
knowledge_base_success_task_list,
|
||||
)[0]
|
||||
similar_indexes = sorted(
|
||||
range(len(similarity)),
|
||||
key=lambda i: similarity[i],
|
||||
reverse=True,
|
||||
)[:v1_query_similar_success_limit]
|
||||
similar_successful_knowledge = [
|
||||
self.knowledgebase.implementation_trace.setdefault(
|
||||
knowledge_base_success_task_list[index],
|
||||
[],
|
||||
)[-1]
|
||||
for index in similar_indexes
|
||||
]
|
||||
queried_knowledge.working_task_to_similar_successful_knowledge_dict[target_factor_task_information] = (
|
||||
similar_successful_knowledge
|
||||
)
|
||||
return queried_knowledge
|
||||
|
||||
|
||||
class FactorQueriedGraphKnowledge(FactorQueriedKnowledge):
|
||||
# Aggregation of knowledge
|
||||
def __init__(
|
||||
self,
|
||||
former_traces: dict = {},
|
||||
component_with_success_task: dict = {},
|
||||
error_with_success_task: dict = {},
|
||||
**kwargs,
|
||||
) -> None:
|
||||
self.former_traces = former_traces
|
||||
self.component_with_success_task = component_with_success_task
|
||||
self.error_with_success_task = error_with_success_task
|
||||
super().__init__(**kwargs)
|
||||
|
||||
|
||||
class FactorGraphRAGStrategy(RAGStrategy):
|
||||
def __init__(self, knowledgebase: FactorGraphKnowledgeBase) -> None:
|
||||
super().__init__(knowledgebase)
|
||||
self.current_generated_trace_count = 0
|
||||
self.prompt = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml")
|
||||
|
||||
def generate_knowledge(
|
||||
self,
|
||||
evolving_trace: list[EvoStep],
|
||||
*,
|
||||
return_knowledge: bool = False,
|
||||
) -> Knowledge | None:
|
||||
if len(evolving_trace) == self.current_generated_trace_count:
|
||||
return None
|
||||
|
||||
else:
|
||||
for trace_index in range(self.current_generated_trace_count, len(evolving_trace)):
|
||||
evo_step = evolving_trace[trace_index]
|
||||
implementations = evo_step.evolvable_subjects
|
||||
feedback = evo_step.feedback
|
||||
for task_index in range(len(implementations.sub_tasks)):
|
||||
single_feedback = feedback[task_index]
|
||||
target_task = implementations.sub_tasks[task_index]
|
||||
target_task_information = target_task.get_factor_information()
|
||||
implementation = implementations.sub_implementations[task_index]
|
||||
single_feedback = feedback[task_index]
|
||||
if single_feedback is None:
|
||||
continue
|
||||
single_knowledge = FactorKnowledge(
|
||||
target_task=target_task,
|
||||
implementation=implementation,
|
||||
feedback=single_feedback,
|
||||
)
|
||||
if (
|
||||
target_task_information not in self.knowledgebase.success_task_to_knowledge_dict
|
||||
and implementation is not None
|
||||
):
|
||||
self.knowledgebase.working_trace_knowledge.setdefault(target_task_information, []).append(
|
||||
single_knowledge,
|
||||
) # save to working trace
|
||||
if single_feedback.final_decision == True:
|
||||
self.knowledgebase.success_task_to_knowledge_dict.setdefault(
|
||||
target_task_information,
|
||||
single_knowledge,
|
||||
)
|
||||
# Do summary for the last step and update the knowledge graph
|
||||
self.knowledgebase.update_success_task(
|
||||
target_task_information,
|
||||
)
|
||||
else:
|
||||
# generate error node and store into knowledge base
|
||||
error_analysis_result = []
|
||||
if not single_feedback.value_generated_flag:
|
||||
error_analysis_result = self.analyze_error(
|
||||
single_feedback.execution_feedback,
|
||||
feedback_type="execution",
|
||||
)
|
||||
else:
|
||||
error_analysis_result = self.analyze_error(
|
||||
single_feedback.factor_value_feedback,
|
||||
feedback_type="value",
|
||||
)
|
||||
self.knowledgebase.working_trace_error_analysis.setdefault(
|
||||
target_task_information,
|
||||
[],
|
||||
).append(
|
||||
error_analysis_result,
|
||||
) # save to working trace error record, for graph update
|
||||
|
||||
self.current_generated_trace_count = len(evolving_trace)
|
||||
return None
|
||||
|
||||
def query(self, evo: EvolvableSubjects, evolving_trace: list[EvoStep]) -> QueriedKnowledge | None:
|
||||
conf_knowledge_sampler = FACTOR_IMPLEMENT_SETTINGS.v2_knowledge_sampler
|
||||
factor_implementation_queried_graph_knowledge = FactorQueriedGraphKnowledge(
|
||||
success_task_to_knowledge_dict=self.knowledgebase.success_task_to_knowledge_dict,
|
||||
)
|
||||
|
||||
factor_implementation_queried_graph_knowledge = self.former_trace_query(
|
||||
evo,
|
||||
factor_implementation_queried_graph_knowledge,
|
||||
FACTOR_IMPLEMENT_SETTINGS.v2_query_former_trace_limit,
|
||||
)
|
||||
factor_implementation_queried_graph_knowledge = self.component_query(
|
||||
evo,
|
||||
factor_implementation_queried_graph_knowledge,
|
||||
FACTOR_IMPLEMENT_SETTINGS.v2_query_component_limit,
|
||||
knowledge_sampler=conf_knowledge_sampler,
|
||||
)
|
||||
factor_implementation_queried_graph_knowledge = self.error_query(
|
||||
evo,
|
||||
factor_implementation_queried_graph_knowledge,
|
||||
FACTOR_IMPLEMENT_SETTINGS.v2_query_error_limit,
|
||||
knowledge_sampler=conf_knowledge_sampler,
|
||||
)
|
||||
return factor_implementation_queried_graph_knowledge
|
||||
|
||||
def analyze_component(
|
||||
self,
|
||||
target_factor_task_information,
|
||||
) -> list[UndirectedNode]: # Hardcode: certain component nodes
|
||||
all_component_nodes = self.knowledgebase.graph.get_all_nodes_by_label_list(["component"])
|
||||
if not len(all_component_nodes):
|
||||
return []
|
||||
all_component_content = ""
|
||||
for _, component_node in enumerate(all_component_nodes):
|
||||
all_component_content += f"{component_node.content}, \n"
|
||||
analyze_component_system_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(self.prompt["analyze_component_prompt_v1_system"])
|
||||
.render(
|
||||
all_component_content=all_component_content,
|
||||
)
|
||||
)
|
||||
|
||||
analyze_component_user_prompt = target_factor_task_information
|
||||
try:
|
||||
component_no_list = json.loads(
|
||||
APIBackend().build_messages_and_create_chat_completion(
|
||||
system_prompt=analyze_component_system_prompt,
|
||||
user_prompt=analyze_component_user_prompt,
|
||||
json_mode=True,
|
||||
),
|
||||
)["component_no_list"]
|
||||
return [all_component_nodes[index - 1] for index in sorted(list(set(component_no_list)))]
|
||||
except:
|
||||
RDAgentLog().warning("Error when analyzing components.")
|
||||
analyze_component_user_prompt = "Your response is not a valid component index list."
|
||||
|
||||
return []
|
||||
|
||||
def analyze_error(
|
||||
self,
|
||||
single_feedback,
|
||||
feedback_type="execution",
|
||||
) -> list[
|
||||
UndirectedNode | str
|
||||
]: # Hardcode: Raised errors, existed error nodes + not existed error nodes(here, they are strs)
|
||||
if feedback_type == "execution":
|
||||
match = re.search(
|
||||
r'File "(?P<file>.+)", line (?P<line>\d+), in (?P<function>.+)\n\s+(?P<error_line>.+)\n(?P<error_type>\w+): (?P<error_message>.+)',
|
||||
single_feedback,
|
||||
)
|
||||
if match:
|
||||
error_details = match.groupdict()
|
||||
# last_traceback = f'File "{error_details["file"]}", line {error_details["line"]}, in {error_details["function"]}\n {error_details["error_line"]}'
|
||||
error_type = error_details["error_type"]
|
||||
error_line = error_details["error_line"]
|
||||
error_contents = [f"ErrorType: {error_type}" + "\n" + f"Error line: {error_line}"]
|
||||
else:
|
||||
error_contents = ["Undefined Error"]
|
||||
elif feedback_type == "value": # value check error
|
||||
value_check_types = r"The source dataframe and the ground truth dataframe have different rows count.|The source dataframe and the ground truth dataframe have different index.|Some values differ by more than the tolerance of 1e-6.|No sufficient correlation found when shifting up|Something wrong happens when naming the multi indices of the dataframe."
|
||||
error_contents = re.findall(value_check_types, single_feedback)
|
||||
else:
|
||||
error_contents = ["Undefined Error"]
|
||||
|
||||
all_error_nodes = self.knowledgebase.graph.get_all_nodes_by_label_list(["error"])
|
||||
if not len(all_error_nodes):
|
||||
return error_contents
|
||||
else:
|
||||
error_list = []
|
||||
for error_content in error_contents:
|
||||
for error_node in all_error_nodes:
|
||||
if error_content == error_node.content:
|
||||
error_list.append(error_node)
|
||||
else:
|
||||
error_list.append(error_content)
|
||||
if error_list[-1] in error_list[:-1]:
|
||||
error_list.pop()
|
||||
|
||||
return error_list
|
||||
|
||||
def former_trace_query(
|
||||
self,
|
||||
evo: EvolvableSubjects,
|
||||
factor_implementation_queried_graph_knowledge: FactorQueriedGraphKnowledge,
|
||||
v2_query_former_trace_limit: int = 5,
|
||||
) -> Union[QueriedKnowledge, set]:
|
||||
"""
|
||||
Query the former trace knowledge of the working trace, and find all the failed task information which tried more than fail_task_trial_limit times
|
||||
"""
|
||||
fail_task_trial_limit = FACTOR_IMPLEMENT_SETTINGS.fail_task_trial_limit
|
||||
|
||||
for target_factor_task in evo.sub_tasks:
|
||||
target_factor_task_information = target_factor_task.get_factor_information()
|
||||
if (
|
||||
target_factor_task_information not in self.knowledgebase.success_task_to_knowledge_dict
|
||||
and target_factor_task_information in self.knowledgebase.working_trace_knowledge
|
||||
and len(self.knowledgebase.working_trace_knowledge[target_factor_task_information])
|
||||
>= fail_task_trial_limit
|
||||
):
|
||||
factor_implementation_queried_graph_knowledge.failed_task_info_set.add(target_factor_task_information)
|
||||
|
||||
if (
|
||||
target_factor_task_information not in self.knowledgebase.success_task_to_knowledge_dict
|
||||
and target_factor_task_information
|
||||
not in factor_implementation_queried_graph_knowledge.failed_task_info_set
|
||||
and target_factor_task_information in self.knowledgebase.working_trace_knowledge
|
||||
):
|
||||
former_trace_knowledge = copy.copy(
|
||||
self.knowledgebase.working_trace_knowledge[target_factor_task_information],
|
||||
)
|
||||
# in former trace query we will delete the right trace in the following order:[..., value_generated_flag is True, value_generated_flag is False, ...]
|
||||
# because we think this order means a deterioration of the trial (like a wrong gradient descent)
|
||||
current_index = 1
|
||||
while current_index < len(former_trace_knowledge):
|
||||
if (
|
||||
not former_trace_knowledge[current_index].feedback.value_generated_flag
|
||||
and former_trace_knowledge[current_index - 1].feedback.value_generated_flag
|
||||
):
|
||||
former_trace_knowledge.pop(current_index)
|
||||
else:
|
||||
current_index += 1
|
||||
|
||||
factor_implementation_queried_graph_knowledge.former_traces[target_factor_task_information] = (
|
||||
former_trace_knowledge[-v2_query_former_trace_limit:]
|
||||
)
|
||||
else:
|
||||
factor_implementation_queried_graph_knowledge.former_traces[target_factor_task_information] = []
|
||||
|
||||
return factor_implementation_queried_graph_knowledge
|
||||
|
||||
def component_query(
|
||||
self,
|
||||
evo: EvolvableSubjects,
|
||||
factor_implementation_queried_graph_knowledge: FactorQueriedGraphKnowledge,
|
||||
v2_query_component_limit: int = 5,
|
||||
knowledge_sampler: float = 1.0,
|
||||
) -> QueriedKnowledge | None:
|
||||
# queried_component_knowledge = FactorQueriedGraphComponentKnowledge()
|
||||
for target_factor_task in evo.sub_tasks:
|
||||
target_factor_task_information = target_factor_task.get_factor_information()
|
||||
if (
|
||||
target_factor_task_information in self.knowledgebase.success_task_to_knowledge_dict
|
||||
or target_factor_task_information in factor_implementation_queried_graph_knowledge.failed_task_info_set
|
||||
):
|
||||
factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
] = []
|
||||
else:
|
||||
if target_factor_task_information not in self.knowledgebase.task_to_component_nodes:
|
||||
self.knowledgebase.task_to_component_nodes[target_factor_task_information] = self.analyze_component(
|
||||
target_factor_task_information,
|
||||
)
|
||||
|
||||
component_analysis_result = self.knowledgebase.task_to_component_nodes[target_factor_task_information]
|
||||
|
||||
if len(component_analysis_result) > 1:
|
||||
task_des_node_list = self.knowledgebase.graph_query_by_intersection(
|
||||
component_analysis_result,
|
||||
constraint_labels=["task_description"],
|
||||
)
|
||||
single_component_constraint = (v2_query_component_limit // len(component_analysis_result)) + 1
|
||||
else:
|
||||
task_des_node_list = []
|
||||
single_component_constraint = v2_query_component_limit
|
||||
factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
] = []
|
||||
for component_node in component_analysis_result:
|
||||
# Reverse iterate, a trade-off with intersection search
|
||||
count = 0
|
||||
for task_des_node in self.knowledgebase.graph_query_by_node(
|
||||
node=component_node,
|
||||
step=1,
|
||||
constraint_labels=["task_description"],
|
||||
block=True,
|
||||
)[::-1]:
|
||||
if task_des_node not in task_des_node_list:
|
||||
task_des_node_list.append(task_des_node)
|
||||
count += 1
|
||||
if count >= single_component_constraint:
|
||||
break
|
||||
|
||||
for node in task_des_node_list:
|
||||
for searched_node in self.knowledgebase.graph_query_by_node(
|
||||
node=node,
|
||||
step=50,
|
||||
constraint_labels=[
|
||||
"task_success_implement",
|
||||
],
|
||||
block=True,
|
||||
):
|
||||
if searched_node.label == "task_success_implement":
|
||||
target_knowledge = self.knowledgebase.node_to_implementation_knowledge_dict[
|
||||
searched_node.id
|
||||
]
|
||||
if (
|
||||
target_knowledge
|
||||
not in factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
]
|
||||
):
|
||||
factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
].append(target_knowledge)
|
||||
|
||||
# finally add embedding related knowledge
|
||||
knowledge_base_success_task_list = list(self.knowledgebase.success_task_to_knowledge_dict)
|
||||
|
||||
similarity = calculate_embedding_distance_between_str_list(
|
||||
[target_factor_task_information],
|
||||
knowledge_base_success_task_list,
|
||||
)[0]
|
||||
similar_indexes = sorted(
|
||||
range(len(similarity)),
|
||||
key=lambda i: similarity[i],
|
||||
reverse=True,
|
||||
)
|
||||
embedding_similar_successful_knowledge = [
|
||||
self.knowledgebase.success_task_to_knowledge_dict[knowledge_base_success_task_list[index]]
|
||||
for index in similar_indexes
|
||||
]
|
||||
for knowledge in embedding_similar_successful_knowledge:
|
||||
if (
|
||||
knowledge
|
||||
not in factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
]
|
||||
):
|
||||
factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
].append(knowledge)
|
||||
|
||||
if knowledge_sampler > 0:
|
||||
factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
] = [
|
||||
knowledge
|
||||
for knowledge in factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
]
|
||||
if random.uniform(0, 1) <= knowledge_sampler
|
||||
]
|
||||
|
||||
# Make sure no less than half of the knowledge are from GT
|
||||
queried_knowledge_list = factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
]
|
||||
queried_from_gt_knowledge_list = [
|
||||
knowledge
|
||||
for knowledge in queried_knowledge_list
|
||||
if knowledge.feedback is not None and knowledge.feedback.final_decision_based_on_gt == True
|
||||
]
|
||||
queried_without_gt_knowledge_list = [
|
||||
knowledge
|
||||
for knowledge in queried_knowledge_list
|
||||
if knowledge.feedback is not None and knowledge.feedback.final_decision_based_on_gt == False
|
||||
]
|
||||
queried_from_gt_knowledge_count = max(
|
||||
min(v2_query_component_limit // 2, len(queried_from_gt_knowledge_list)),
|
||||
v2_query_component_limit - len(queried_without_gt_knowledge_list),
|
||||
)
|
||||
factor_implementation_queried_graph_knowledge.component_with_success_task[
|
||||
target_factor_task_information
|
||||
] = (
|
||||
queried_from_gt_knowledge_list[:queried_from_gt_knowledge_count]
|
||||
+ queried_without_gt_knowledge_list[: v2_query_component_limit - queried_from_gt_knowledge_count]
|
||||
)
|
||||
|
||||
return factor_implementation_queried_graph_knowledge
|
||||
|
||||
def error_query(
|
||||
self,
|
||||
evo: EvolvableSubjects,
|
||||
factor_implementation_queried_graph_knowledge: FactorQueriedGraphKnowledge,
|
||||
v2_query_error_limit: int = 5,
|
||||
knowledge_sampler: float = 1.0,
|
||||
) -> QueriedKnowledge | None:
|
||||
# queried_error_knowledge = FactorQueriedGraphErrorKnowledge()
|
||||
for task_index, target_factor_task in enumerate(evo.sub_tasks):
|
||||
target_factor_task_information = target_factor_task.get_factor_information()
|
||||
factor_implementation_queried_graph_knowledge.error_with_success_task[target_factor_task_information] = {}
|
||||
if (
|
||||
target_factor_task_information in self.knowledgebase.success_task_to_knowledge_dict
|
||||
or target_factor_task_information in factor_implementation_queried_graph_knowledge.failed_task_info_set
|
||||
):
|
||||
factor_implementation_queried_graph_knowledge.error_with_success_task[
|
||||
target_factor_task_information
|
||||
] = []
|
||||
else:
|
||||
factor_implementation_queried_graph_knowledge.error_with_success_task[
|
||||
target_factor_task_information
|
||||
] = []
|
||||
if (
|
||||
target_factor_task_information in self.knowledgebase.working_trace_error_analysis
|
||||
and len(self.knowledgebase.working_trace_error_analysis[target_factor_task_information]) > 0
|
||||
and len(factor_implementation_queried_graph_knowledge.former_traces[target_factor_task_information])
|
||||
> 0
|
||||
):
|
||||
queried_last_trace = factor_implementation_queried_graph_knowledge.former_traces[
|
||||
target_factor_task_information
|
||||
][-1]
|
||||
target_index = self.knowledgebase.working_trace_knowledge[target_factor_task_information].index(
|
||||
queried_last_trace,
|
||||
)
|
||||
last_knowledge_error_analysis_result = self.knowledgebase.working_trace_error_analysis[
|
||||
target_factor_task_information
|
||||
][target_index]
|
||||
else:
|
||||
last_knowledge_error_analysis_result = []
|
||||
|
||||
error_nodes = []
|
||||
for error_node in last_knowledge_error_analysis_result:
|
||||
if not isinstance(error_node, UndirectedNode):
|
||||
error_node = self.knowledgebase.graph_get_node_by_content(content=error_node)
|
||||
if error_node is None:
|
||||
continue
|
||||
error_nodes.append(error_node)
|
||||
|
||||
if len(error_nodes) > 1:
|
||||
task_trace_node_list = self.knowledgebase.graph_query_by_intersection(
|
||||
error_nodes,
|
||||
constraint_labels=["task_trace"],
|
||||
output_intersection_origin=True,
|
||||
)
|
||||
single_error_constraint = (v2_query_error_limit // len(error_nodes)) + 1
|
||||
else:
|
||||
task_trace_node_list = []
|
||||
single_error_constraint = v2_query_error_limit
|
||||
for error_node in error_nodes:
|
||||
# Reverse iterate, a trade-off with intersection search
|
||||
count = 0
|
||||
for task_trace_node in self.knowledgebase.graph_query_by_node(
|
||||
node=error_node,
|
||||
step=1,
|
||||
constraint_labels=["task_trace"],
|
||||
block=True,
|
||||
)[::-1]:
|
||||
if task_trace_node not in task_trace_node_list:
|
||||
task_trace_node_list.append([[error_node], task_trace_node])
|
||||
count += 1
|
||||
if count >= single_error_constraint:
|
||||
break
|
||||
|
||||
# for error_node in last_knowledge_error_analysis_result:
|
||||
# if not isinstance(error_node, UndirectedNode):
|
||||
# error_node = self.knowledgebase.graph_get_node_by_content(content=error_node)
|
||||
# if error_node is None:
|
||||
# continue
|
||||
# for searched_node in self.knowledgebase.graph_query_by_node(
|
||||
# node=error_node,
|
||||
# step=1,
|
||||
# constraint_labels=["task_trace"],
|
||||
# block=True,
|
||||
# ):
|
||||
# if searched_node not in [node[0] for node in task_trace_node_list]:
|
||||
# task_trace_node_list.append((searched_node, error_node.content))
|
||||
|
||||
same_error_success_knowledge_pair_list = []
|
||||
same_error_success_node_set = set()
|
||||
for error_node_list, trace_node in task_trace_node_list:
|
||||
for searched_trace_success_node in self.knowledgebase.graph_query_by_node(
|
||||
node=trace_node,
|
||||
step=50,
|
||||
constraint_labels=[
|
||||
"task_trace",
|
||||
"task_success_implement",
|
||||
"task_description",
|
||||
],
|
||||
block=True,
|
||||
):
|
||||
if (
|
||||
searched_trace_success_node not in same_error_success_node_set
|
||||
and searched_trace_success_node.label == "task_success_implement"
|
||||
):
|
||||
same_error_success_node_set.add(searched_trace_success_node)
|
||||
|
||||
trace_knowledge = self.knowledgebase.node_to_implementation_knowledge_dict[trace_node.id]
|
||||
success_knowledge = self.knowledgebase.node_to_implementation_knowledge_dict[
|
||||
searched_trace_success_node.id
|
||||
]
|
||||
error_content = ""
|
||||
for index, error_node in enumerate(error_node_list):
|
||||
error_content += f"{index+1}. {error_node.content}; "
|
||||
same_error_success_knowledge_pair_list.append(
|
||||
(
|
||||
error_content,
|
||||
(trace_knowledge, success_knowledge),
|
||||
),
|
||||
)
|
||||
|
||||
if knowledge_sampler > 0:
|
||||
same_error_success_knowledge_pair_list = [
|
||||
knowledge
|
||||
for knowledge in same_error_success_knowledge_pair_list
|
||||
if random.uniform(0, 1) <= knowledge_sampler
|
||||
]
|
||||
|
||||
same_error_success_knowledge_pair_list = same_error_success_knowledge_pair_list[:v2_query_error_limit]
|
||||
factor_implementation_queried_graph_knowledge.error_with_success_task[
|
||||
target_factor_task_information
|
||||
] = same_error_success_knowledge_pair_list
|
||||
|
||||
return factor_implementation_queried_graph_knowledge
|
||||
|
||||
|
||||
class FactorGraphKnowledgeBase(KnowledgeBase):
|
||||
def __init__(self, init_component_list=None) -> None:
|
||||
"""
|
||||
Load knowledge, offer brief information of knowledge and common handle interfaces
|
||||
"""
|
||||
self.graph: UndirectedGraph = UndirectedGraph.load(Path.cwd() / "graph.pkl")
|
||||
RDAgentLog().info(f"Knowledge Graph loaded, size={self.graph.size()}")
|
||||
|
||||
if init_component_list:
|
||||
for component in init_component_list:
|
||||
exist_node = self.graph.get_node_by_content(content=component)
|
||||
node = exist_node if exist_node else UndirectedNode(content=component, label="component")
|
||||
self.graph.add_nodes(node=node, neighbors=[])
|
||||
|
||||
# A dict containing all working trace until they fail or succeed
|
||||
self.working_trace_knowledge = {}
|
||||
|
||||
# A dict containing error analysis each step aligned with working trace
|
||||
self.working_trace_error_analysis = {}
|
||||
|
||||
# Add already success task
|
||||
self.success_task_to_knowledge_dict = {}
|
||||
|
||||
# key:node_id(for task trace and success implement), value:knowledge instance(aka 'FactorKnowledge')
|
||||
self.node_to_implementation_knowledge_dict = {}
|
||||
|
||||
# store the task description to component nodes
|
||||
self.task_to_component_nodes = {}
|
||||
|
||||
def get_all_nodes_by_label(self, label: str) -> list[UndirectedNode]:
|
||||
return self.graph.get_all_nodes_by_label(label)
|
||||
|
||||
def update_success_task(
|
||||
self,
|
||||
success_task_info: str,
|
||||
): # Transfer the success tasks' working trace to knowledge storage & graph
|
||||
success_task_trace = self.working_trace_knowledge[success_task_info]
|
||||
success_task_error_analysis_record = (
|
||||
self.working_trace_error_analysis[success_task_info]
|
||||
if success_task_info in self.working_trace_error_analysis
|
||||
else []
|
||||
)
|
||||
task_des_node = UndirectedNode(content=success_task_info, label="task_description")
|
||||
self.graph.add_nodes(
|
||||
node=task_des_node,
|
||||
neighbors=self.task_to_component_nodes[success_task_info],
|
||||
) # 1st version, we assume that all component nodes are given
|
||||
for index, trace_unit in enumerate(success_task_trace): # every unit: single_knowledge
|
||||
neighbor_nodes = [task_des_node]
|
||||
if index != len(success_task_trace) - 1:
|
||||
trace_node = UndirectedNode(
|
||||
content=trace_unit.get_implementation_and_feedback_str(),
|
||||
label="task_trace",
|
||||
)
|
||||
self.node_to_implementation_knowledge_dict[trace_node.id] = trace_unit
|
||||
for node_index, error_node in enumerate(success_task_error_analysis_record[index]):
|
||||
if type(error_node).__name__ == "str":
|
||||
queried_node = self.graph.get_node_by_content(content=error_node)
|
||||
if queried_node is None:
|
||||
new_error_node = UndirectedNode(content=error_node, label="error")
|
||||
self.graph.add_node(node=new_error_node)
|
||||
success_task_error_analysis_record[index][node_index] = new_error_node
|
||||
else:
|
||||
success_task_error_analysis_record[index][node_index] = queried_node
|
||||
neighbor_nodes.extend(success_task_error_analysis_record[index])
|
||||
self.graph.add_nodes(node=trace_node, neighbors=neighbor_nodes)
|
||||
else:
|
||||
success_node = UndirectedNode(
|
||||
content=trace_unit.get_implementation_and_feedback_str(),
|
||||
label="task_success_implement",
|
||||
)
|
||||
self.graph.add_nodes(node=success_node, neighbors=neighbor_nodes)
|
||||
self.node_to_implementation_knowledge_dict[success_node.id] = trace_unit
|
||||
|
||||
def query(self):
|
||||
pass
|
||||
|
||||
def graph_get_node_by_content(self, content: str) -> UndirectedNode:
|
||||
return self.graph.get_node_by_content(content=content)
|
||||
|
||||
def graph_query_by_content(
|
||||
self,
|
||||
content: Union[str, list[str]],
|
||||
topk_k: int = 5,
|
||||
step: int = 1,
|
||||
constraint_labels: list[str] = None,
|
||||
constraint_node: UndirectedNode = None,
|
||||
similarity_threshold: float = 0.0,
|
||||
constraint_distance: float = 0,
|
||||
block: bool = False,
|
||||
) -> list[UndirectedNode]:
|
||||
"""
|
||||
search graph by content similarity and connection relationship, return empty list if nodes' chain without node
|
||||
near to constraint_node
|
||||
|
||||
Parameters
|
||||
----------
|
||||
constraint_distance
|
||||
content
|
||||
topk_k: the upper number of output for each query, if the number of fit nodes is less than topk_k, return all fit nodes's content
|
||||
step
|
||||
constraint_labels
|
||||
constraint_node
|
||||
similarity_threshold
|
||||
block: despite the start node, the search can only flow through the constraint_label type nodes
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
"""
|
||||
|
||||
return self.graph.query_by_content(
|
||||
content=content,
|
||||
topk_k=topk_k,
|
||||
step=step,
|
||||
constraint_labels=constraint_labels,
|
||||
constraint_node=constraint_node,
|
||||
similarity_threshold=similarity_threshold,
|
||||
constraint_distance=constraint_distance,
|
||||
block=block,
|
||||
)
|
||||
|
||||
def graph_query_by_node(
|
||||
self,
|
||||
node: UndirectedNode,
|
||||
step: int = 1,
|
||||
constraint_labels: list[str] = None,
|
||||
constraint_node: UndirectedNode = None,
|
||||
constraint_distance: float = 0,
|
||||
block: bool = False,
|
||||
) -> list[UndirectedNode]:
|
||||
"""
|
||||
search graph by connection, return empty list if nodes' chain without node near to constraint_node
|
||||
Parameters
|
||||
----------
|
||||
node : start node
|
||||
step : the max steps will be searched
|
||||
constraint_labels : the labels of output nodes
|
||||
constraint_node : the node that the output nodes must connect to
|
||||
constraint_distance : the max distance between output nodes and constraint_node
|
||||
block: despite the start node, the search can only flow through the constraint_label type nodes
|
||||
|
||||
Returns
|
||||
-------
|
||||
A list of nodes
|
||||
|
||||
"""
|
||||
nodes = self.graph.query_by_node(
|
||||
node=node,
|
||||
step=step,
|
||||
constraint_labels=constraint_labels,
|
||||
constraint_node=constraint_node,
|
||||
constraint_distance=constraint_distance,
|
||||
block=block,
|
||||
)
|
||||
return nodes
|
||||
|
||||
def graph_query_by_intersection(
|
||||
self,
|
||||
nodes: list[UndirectedNode],
|
||||
steps: int = 1,
|
||||
constraint_labels: list[str] = None,
|
||||
output_intersection_origin: bool = False,
|
||||
) -> list[UndirectedNode] | list[list[list[UndirectedNode], UndirectedNode]]:
|
||||
"""
|
||||
search graph by node intersection, node intersected by a higher frequency has a prior order in the list
|
||||
Parameters
|
||||
----------
|
||||
nodes : node list
|
||||
step : the max steps will be searched
|
||||
constraint_labels : the labels of output nodes
|
||||
output_intersection_origin: output the list that contains the node which form this intersection node
|
||||
|
||||
Returns
|
||||
-------
|
||||
A list of nodes
|
||||
|
||||
"""
|
||||
node_count = len(nodes)
|
||||
assert node_count >= 2, "nodes length must >=2"
|
||||
intersection_node_list = []
|
||||
if output_intersection_origin:
|
||||
origin_list = []
|
||||
for k in range(node_count, 1, -1):
|
||||
possible_combinations = combinations(nodes, k)
|
||||
for possible_combination in possible_combinations:
|
||||
node_list = list(possible_combination)
|
||||
intersection_node_list.extend(
|
||||
self.graph.get_nodes_intersection(node_list, steps=steps, constraint_labels=constraint_labels),
|
||||
)
|
||||
if output_intersection_origin:
|
||||
for _ in range(len(intersection_node_list)):
|
||||
origin_list.append(node_list)
|
||||
intersection_node_list_sort_by_freq = []
|
||||
for index, node in enumerate(intersection_node_list):
|
||||
if node not in intersection_node_list_sort_by_freq:
|
||||
if output_intersection_origin:
|
||||
intersection_node_list_sort_by_freq.append([origin_list[index], node])
|
||||
else:
|
||||
intersection_node_list_sort_by_freq.append(node)
|
||||
|
||||
return intersection_node_list_sort_by_freq
|
||||
@@ -0,0 +1,89 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
from jinja2 import Environment, StrictUndefined
|
||||
|
||||
from rdagent.components.coder.factor_coder.CoSTEER.evolvable_subjects import (
|
||||
FactorEvolvingItem,
|
||||
)
|
||||
from rdagent.components.coder.factor_coder.utils import get_data_folder_intro
|
||||
from rdagent.core.conf import RD_AGENT_SETTINGS
|
||||
from rdagent.core.log import RDAgentLog
|
||||
from rdagent.core.prompts import Prompts
|
||||
from rdagent.core.scenario import Scenario
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
scheduler_prompts = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml")
|
||||
|
||||
|
||||
def RandomSelect(to_be_finished_task_index, implementation_factors_per_round):
|
||||
import random
|
||||
|
||||
to_be_finished_task_index = random.sample(
|
||||
to_be_finished_task_index,
|
||||
implementation_factors_per_round,
|
||||
)
|
||||
|
||||
RDAgentLog().info(f"The random selection is: {to_be_finished_task_index}")
|
||||
return to_be_finished_task_index
|
||||
|
||||
|
||||
def LLMSelect(
|
||||
to_be_finished_task_index,
|
||||
implementation_factors_per_round,
|
||||
evo: FactorEvolvingItem,
|
||||
former_trace: Dict,
|
||||
scen: Scenario,
|
||||
):
|
||||
tasks = []
|
||||
for i in to_be_finished_task_index:
|
||||
# find corresponding former trace for each task
|
||||
target_factor_task_information = evo.sub_tasks[i].get_factor_information()
|
||||
if target_factor_task_information in former_trace:
|
||||
tasks.append((i, evo.sub_tasks[i], former_trace[target_factor_task_information]))
|
||||
|
||||
system_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
scheduler_prompts["select_implementable_factor_system"],
|
||||
)
|
||||
.render(
|
||||
scenario=scen.get_scenario_all_desc(),
|
||||
)
|
||||
)
|
||||
|
||||
while True:
|
||||
user_prompt = (
|
||||
Environment(undefined=StrictUndefined)
|
||||
.from_string(
|
||||
scheduler_prompts["select_implementable_factor_user"],
|
||||
)
|
||||
.render(
|
||||
factor_num=implementation_factors_per_round,
|
||||
sub_tasks=tasks,
|
||||
)
|
||||
)
|
||||
if (
|
||||
APIBackend().build_messages_and_calculate_token(
|
||||
user_prompt=user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
)
|
||||
< RD_AGENT_SETTINGS.chat_token_limit
|
||||
):
|
||||
break
|
||||
|
||||
response = APIBackend().build_messages_and_create_chat_completion(
|
||||
user_prompt=user_prompt,
|
||||
system_prompt=system_prompt,
|
||||
json_mode=True,
|
||||
)
|
||||
try:
|
||||
selection = json.loads(response)["selected_factor"]
|
||||
if not isinstance(selection, list):
|
||||
return to_be_finished_task_index
|
||||
selection_index = [x for x in selection if isinstance(x, int)]
|
||||
except:
|
||||
return to_be_finished_task_index
|
||||
|
||||
return selection_index
|
||||
@@ -0,0 +1,48 @@
|
||||
from pathlib import Path
|
||||
from typing import Literal, Union
|
||||
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
SELECT_METHOD = Literal["random", "scheduler"]
|
||||
|
||||
|
||||
class FactorImplementSettings(BaseSettings):
|
||||
file_based_execution_data_folder: str = str(
|
||||
(Path().cwd() / "git_ignore_folder" / "factor_implementation_source_data").absolute(),
|
||||
)
|
||||
file_based_execution_workspace: str = str(
|
||||
(Path().cwd() / "git_ignore_folder" / "factor_implementation_workspace").absolute(),
|
||||
)
|
||||
implementation_execution_cache_location: str = str(
|
||||
(Path().cwd() / "git_ignore_folder" / "factor_implementation_execution_cache").absolute(),
|
||||
)
|
||||
enable_execution_cache: bool = True # whether to enable the execution cache
|
||||
|
||||
# TODO: the factor implement specific settings should not appear in this settings
|
||||
# Evolving should have a method specific settings
|
||||
# evolving related config
|
||||
fail_task_trial_limit: int = 20
|
||||
|
||||
v1_query_former_trace_limit: int = 5
|
||||
v1_query_similar_success_limit: int = 5
|
||||
|
||||
v2_query_component_limit: int = 1
|
||||
v2_query_error_limit: int = 1
|
||||
v2_query_former_trace_limit: int = 1
|
||||
v2_error_summary: bool = False
|
||||
v2_knowledge_sampler: float = 1.0
|
||||
|
||||
evo_multi_proc_n: int = 16 # how many processes to use for evolving (including eval & generation)
|
||||
|
||||
file_based_execution_timeout: int = 120 # seconds for each factor implementation execution
|
||||
|
||||
select_method: SELECT_METHOD = "random"
|
||||
select_ratio: float = 0.5
|
||||
|
||||
max_loop: int = 10
|
||||
|
||||
knowledge_base_path: Union[str, None] = None
|
||||
new_knowledge_base_path: Union[str, None] = None
|
||||
|
||||
|
||||
FACTOR_IMPLEMENT_SETTINGS = FactorImplementSettings()
|
||||
@@ -0,0 +1,221 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pickle
|
||||
import subprocess
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Tuple, Union
|
||||
|
||||
import pandas as pd
|
||||
from filelock import FileLock
|
||||
|
||||
from rdagent.components.coder.factor_coder.config import FACTOR_IMPLEMENT_SETTINGS
|
||||
from rdagent.core.exception import (
|
||||
CodeFormatException,
|
||||
NoOutputException,
|
||||
RuntimeErrorException,
|
||||
)
|
||||
from rdagent.core.experiment import Experiment, FBImplementation, Task
|
||||
from rdagent.core.log import RDAgentLog
|
||||
from rdagent.oai.llm_utils import md5_hash
|
||||
|
||||
|
||||
class FactorTask(Task):
|
||||
# TODO: generalized the attributes into the Task
|
||||
# - factor_* -> *
|
||||
def __init__(
|
||||
self,
|
||||
factor_name,
|
||||
factor_description,
|
||||
factor_formulation,
|
||||
variables: dict = {},
|
||||
resource: str = None,
|
||||
) -> None:
|
||||
self.factor_name = factor_name
|
||||
self.factor_description = factor_description
|
||||
self.factor_formulation = factor_formulation
|
||||
self.variables = variables
|
||||
self.factor_resources = resource
|
||||
|
||||
def get_factor_information(self):
|
||||
return f"""factor_name: {self.factor_name}
|
||||
factor_description: {self.factor_description}
|
||||
factor_formulation: {self.factor_formulation}
|
||||
variables: {str(self.variables)}"""
|
||||
|
||||
@staticmethod
|
||||
def from_dict(dict):
|
||||
return FactorTask(**dict)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}[{self.factor_name}]>"
|
||||
|
||||
|
||||
class FileBasedFactorImplementation(FBImplementation):
|
||||
"""
|
||||
This class is used to implement a factor by writing the code to a file.
|
||||
Input data and output factor value are also written to files.
|
||||
"""
|
||||
|
||||
# TODO: (Xiao) think raising errors may get better information for processing
|
||||
FB_FROM_CACHE = "The factor value has been executed and stored in the instance variable."
|
||||
FB_EXEC_SUCCESS = "Execution succeeded without error."
|
||||
FB_CODE_NOT_SET = "code is not set."
|
||||
FB_EXECUTION_SUCCEEDED = "Execution succeeded without error."
|
||||
FB_OUTPUT_FILE_NOT_FOUND = "\nExpected output file not found."
|
||||
FB_OUTPUT_FILE_FOUND = "\nExpected output file found."
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
target_task: FactorTask,
|
||||
code,
|
||||
executed_factor_value_dataframe=None,
|
||||
raise_exception=False,
|
||||
) -> None:
|
||||
super().__init__(target_task)
|
||||
self.code = code
|
||||
self.executed_factor_value_dataframe = executed_factor_value_dataframe
|
||||
self.logger = RDAgentLog()
|
||||
self.raise_exception = raise_exception
|
||||
self.workspace_path = Path(
|
||||
FACTOR_IMPLEMENT_SETTINGS.file_based_execution_workspace,
|
||||
) / str(uuid.uuid4())
|
||||
|
||||
@staticmethod
|
||||
def link_data_to_workspace(data_path: Path, workspace_path: Path):
|
||||
data_path = Path(data_path)
|
||||
workspace_path = Path(workspace_path)
|
||||
for data_file_path in data_path.iterdir():
|
||||
workspace_data_file_path = workspace_path / data_file_path.name
|
||||
if workspace_data_file_path.exists():
|
||||
workspace_data_file_path.unlink()
|
||||
subprocess.run(
|
||||
["ln", "-s", data_file_path, workspace_data_file_path],
|
||||
check=False,
|
||||
)
|
||||
|
||||
def execute_desc(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def prepare(self, *args, **kwargs):
|
||||
# TODO move the prepare part code in execute into here
|
||||
return super().prepare(*args, **kwargs)
|
||||
|
||||
def execute(self, store_result: bool = False) -> Tuple[str, pd.DataFrame]:
|
||||
"""
|
||||
execute the implementation and get the factor value by the following steps:
|
||||
1. make the directory in workspace path
|
||||
2. write the code to the file in the workspace path
|
||||
3. link all the source data to the workspace path folder
|
||||
4. execute the code
|
||||
5. read the factor value from the output file in the workspace path folder
|
||||
returns the execution feedback as a string and the factor value as a pandas dataframe
|
||||
|
||||
parameters:
|
||||
store_result: if True, store the factor value in the instance variable, this feature is to be used in the gt implementation to avoid multiple execution on the same gt implementation
|
||||
"""
|
||||
if self.code is None:
|
||||
if self.raise_exception:
|
||||
raise CodeFormatException(self.FB_CODE_NOT_SET)
|
||||
else:
|
||||
# TODO: to make the interface compatible with previous code. I kept the original behavior.
|
||||
raise ValueError(self.FB_CODE_NOT_SET)
|
||||
with FileLock(self.workspace_path / "execution.lock"):
|
||||
if FACTOR_IMPLEMENT_SETTINGS.enable_execution_cache:
|
||||
# NOTE: cache the result for the same code
|
||||
target_file_name = md5_hash(self.code)
|
||||
cache_file_path = (
|
||||
Path(FACTOR_IMPLEMENT_SETTINGS.implementation_execution_cache_location) / f"{target_file_name}.pkl"
|
||||
)
|
||||
Path(FACTOR_IMPLEMENT_SETTINGS.implementation_execution_cache_location).mkdir(
|
||||
exist_ok=True, parents=True
|
||||
)
|
||||
if cache_file_path.exists() and not self.raise_exception:
|
||||
cached_res = pickle.load(open(cache_file_path, "rb"))
|
||||
if store_result and cached_res[1] is not None:
|
||||
self.executed_factor_value_dataframe = cached_res[1]
|
||||
return cached_res
|
||||
|
||||
if self.executed_factor_value_dataframe is not None:
|
||||
return self.FB_FROM_CACHE, self.executed_factor_value_dataframe
|
||||
|
||||
source_data_path = Path(
|
||||
FACTOR_IMPLEMENT_SETTINGS.file_based_execution_data_folder,
|
||||
)
|
||||
self.workspace_path.mkdir(exist_ok=True, parents=True)
|
||||
source_data_path.mkdir(exist_ok=True, parents=True)
|
||||
code_path = self.workspace_path / f"{self.target_task.factor_name}.py"
|
||||
code_path.write_text(self.code)
|
||||
|
||||
self.link_data_to_workspace(source_data_path, self.workspace_path)
|
||||
|
||||
execution_feedback = self.FB_EXECUTION_SUCCEEDED
|
||||
try:
|
||||
subprocess.check_output(
|
||||
f"python {code_path}",
|
||||
shell=True,
|
||||
cwd=self.workspace_path,
|
||||
stderr=subprocess.STDOUT,
|
||||
timeout=FACTOR_IMPLEMENT_SETTINGS.file_based_execution_timeout,
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
import site
|
||||
|
||||
execution_feedback = (
|
||||
e.output.decode()
|
||||
.replace(str(code_path.parent.absolute()), r"/path/to")
|
||||
.replace(str(site.getsitepackages()[0]), r"/path/to/site-packages")
|
||||
)
|
||||
if len(execution_feedback) > 2000:
|
||||
execution_feedback = (
|
||||
execution_feedback[:1000] + "....hidden long error message...." + execution_feedback[-1000:]
|
||||
)
|
||||
if self.raise_exception:
|
||||
raise RuntimeErrorException(execution_feedback)
|
||||
except subprocess.TimeoutExpired:
|
||||
execution_feedback += f"Execution timeout error and the timeout is set to {FACTOR_IMPLEMENT_SETTINGS.file_based_execution_timeout} seconds."
|
||||
if self.raise_exception:
|
||||
raise RuntimeErrorException(execution_feedback)
|
||||
|
||||
workspace_output_file_path = self.workspace_path / "result.h5"
|
||||
if not workspace_output_file_path.exists():
|
||||
execution_feedback += self.FB_OUTPUT_FILE_NOT_FOUND
|
||||
executed_factor_value_dataframe = None
|
||||
if self.raise_exception:
|
||||
raise NoOutputException(execution_feedback)
|
||||
else:
|
||||
try:
|
||||
executed_factor_value_dataframe = pd.read_hdf(workspace_output_file_path)
|
||||
execution_feedback += self.FB_OUTPUT_FILE_FOUND
|
||||
except Exception as e:
|
||||
execution_feedback += f"Error found when reading hdf file: {e}"[:1000]
|
||||
executed_factor_value_dataframe = None
|
||||
|
||||
if store_result and executed_factor_value_dataframe is not None:
|
||||
self.executed_factor_value_dataframe = executed_factor_value_dataframe
|
||||
|
||||
if FACTOR_IMPLEMENT_SETTINGS.enable_execution_cache:
|
||||
pickle.dump(
|
||||
(execution_feedback, executed_factor_value_dataframe),
|
||||
open(cache_file_path, "wb"),
|
||||
)
|
||||
return execution_feedback, executed_factor_value_dataframe
|
||||
|
||||
def __str__(self) -> str:
|
||||
# NOTE:
|
||||
# If the code cache works, the workspace will be None.
|
||||
return f"File Factor[{self.target_task.factor_name}]: {self.workspace_path}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__str__()
|
||||
|
||||
@staticmethod
|
||||
def from_folder(task: FactorTask, path: Union[str, Path], **kwargs):
|
||||
path = Path(path)
|
||||
factor_path = (path / task.factor_name).with_suffix(".py")
|
||||
with factor_path.open("r") as f:
|
||||
code = f.read()
|
||||
return FileBasedFactorImplementation(task, code=code, **kwargs)
|
||||
|
||||
|
||||
class FactorExperiment(Experiment[FactorTask, FileBasedFactorImplementation]): ...
|
||||
@@ -0,0 +1,225 @@
|
||||
|
||||
evaluator_code_feedback_v1_system: |-
|
||||
User is trying to implement some factors in the following scenario:
|
||||
{{ scenario }}
|
||||
User will provide you the information of the factor.
|
||||
|
||||
Your job is to check whether user's code is align with the factor and the scenario.
|
||||
The user will provide the source python code and the execution error message if execution failed.
|
||||
The user might provide you the ground truth code for you to provide the critic. You should not leak the ground truth code to the user in any form but you can use it to provide the critic.
|
||||
|
||||
User has also compared the factor values calculated by the user's code and the ground truth code. The user will provide you some analyze result comparing two output. You may find some error in the code which caused the difference between the two output.
|
||||
|
||||
If the ground truth code is provided, your critic should only consider checking whether the user's code is align with the ground truth code since the ground truth is definitely correct.
|
||||
If the ground truth code is not provided, your critic should consider checking whether the user's code is reasonable and correct.
|
||||
|
||||
You should provide the suggestion to each of your critic to help the user improve the code. Please response the critic in the following format. Here is an example structure for the output:
|
||||
critic 1: The critic message to critic 1
|
||||
critic 2: The critic message to critic 2
|
||||
|
||||
evaluator_code_feedback_v1_user: |-
|
||||
--------------Factor information:---------------
|
||||
{{ factor_information }}
|
||||
--------------Python code:---------------
|
||||
{{ code }}
|
||||
--------------Execution feedback:---------------
|
||||
{{ execution_feedback }}
|
||||
{% if factor_value_feedback is not none %}
|
||||
--------------Factor value feedback:---------------
|
||||
{{ factor_value_feedback }}
|
||||
{% endif %}
|
||||
{% if gt_code is not none %}
|
||||
--------------Ground truth Python code:---------------
|
||||
{{ gt_code }}
|
||||
{% endif %}
|
||||
|
||||
evolving_strategy_factor_implementation_v1_system: |-
|
||||
User is trying to implement some factors in the following scenario:
|
||||
{{ scenario }}
|
||||
Your code is expected to align the scenario in any form which means The user needs to get the exact factor values with your code as expected.
|
||||
|
||||
To help you write the correct code, the user might provide multiple information that helps you write the correct code:
|
||||
1. The user might provide you the correct code to similar factors. Your should learn from these code to write the correct code.
|
||||
2. The user might provide you the failed former code and the corresponding feedback to the code. The feedback contains to the execution, the code and the factor value. You should analyze the feedback and try to correct the latest code.
|
||||
3. The user might provide you the suggestion to the latest fail code and some similar fail to correct pairs. Each pair contains the fail code with similar error and the corresponding corrected version code. You should learn from these suggestion to write the correct code.
|
||||
|
||||
Your must write your code based on your former latest attempt below which consists of your former code and code feedback, you should read the former attempt carefully and must not modify the right part of your former code.
|
||||
|
||||
{% if queried_former_failed_knowledge|length != 0 %}
|
||||
--------------Your former latest attempt:---------------
|
||||
{% for former_failed_knowledge in queried_former_failed_knowledge %}
|
||||
=====Code to implementation {{ loop.index }}=====
|
||||
{{ former_failed_knowledge.implementation.code }}
|
||||
=====Feedback to implementation {{ loop.index }}=====
|
||||
{{ former_failed_knowledge.feedback }}
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
Please response the code in the following json format. Here is an example structure for the JSON output:
|
||||
{
|
||||
"code": "The Python code as a string."
|
||||
}
|
||||
|
||||
evolving_strategy_factor_implementation_v1_user: |-
|
||||
--------------Target factor information:---------------
|
||||
{{ factor_information_str }}
|
||||
|
||||
{% if queried_similar_successful_knowledge|length != 0 %}
|
||||
--------------Correct code to similar factors:---------------
|
||||
{% for similar_successful_knowledge in queried_similar_successful_knowledge %}
|
||||
=====Factor {{loop.index}}:=====
|
||||
{{ similar_successful_knowledge.target_task.get_factor_information() }}
|
||||
=====Code:=====
|
||||
{{ similar_successful_knowledge.implementation.code }}
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
{% if queried_former_failed_knowledge|length != 0 %}
|
||||
--------------Former failed code:---------------
|
||||
{% for former_failed_knowledge in queried_former_failed_knowledge %}
|
||||
=====Code to implementation {{ loop.index }}=====
|
||||
{{ former_failed_knowledge.implementation.code }}
|
||||
=====Feedback to implementation {{ loop.index }}=====
|
||||
{{ former_failed_knowledge.feedback }}
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
evolving_strategy_factor_implementation_v2_user: |-
|
||||
--------------Target factor information:---------------
|
||||
{{ factor_information_str }}
|
||||
|
||||
{% if queried_similar_error_knowledge|length != 0 %}
|
||||
{% if not error_summary %}
|
||||
Recall your last failure, your implementation met some errors.
|
||||
When doing other tasks, you met some similar errors but you finally solve them. Here are some examples:
|
||||
{% for error_content, similar_error_knowledge in queried_similar_error_knowledge %}
|
||||
--------------Factor information to similar error ({{error_content}}):---------------
|
||||
{{ similar_error_knowledge[0].target_task.get_factor_information() }}
|
||||
=====Code with similar error ({{error_content}}):=====
|
||||
{{ similar_error_knowledge[0].implementation.code }}
|
||||
=====Success code to former code with similar error ({{error_content}}):=====
|
||||
{{ similar_error_knowledge[1].implementation.code }}
|
||||
{% endfor %}
|
||||
{% else %}
|
||||
Recall your last failure, your implementation met some errors.
|
||||
After reviewing some similar errors and their solutions, here are some suggestions for you to correct your code:
|
||||
{{error_summary_critics}}
|
||||
{% endif %}
|
||||
{% endif %}
|
||||
{% if queried_similar_component_knowledge|length != 0 %}
|
||||
Here are some success implements of similar component tasks, take them as references:
|
||||
--------------Correct code to similar factors:---------------
|
||||
{% for similar_component_knowledge in queried_similar_component_knowledge %}
|
||||
=====Factor {{loop.index}}:=====
|
||||
{{ similar_component_knowledge.target_task.get_factor_information() }}
|
||||
=====Code:=====
|
||||
{{ similar_component_knowledge.implementation.code }}
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
evolving_strategy_error_summary_v2_system: |-
|
||||
You are doing the following task:
|
||||
{{factor_information_str}}
|
||||
|
||||
You have written some code but it meets errors like the following:
|
||||
{{code_and_feedback}}
|
||||
|
||||
The user has found some tasks that met similar errors, and their final correct solutions.
|
||||
Please refer to these similar errors and their solutions, provide some clear, short and accurate critics that might help you solve the issues in your code.
|
||||
|
||||
Please response the critic in the following format. Here is an example structure for the output:
|
||||
critic 1: The critic message to critic 1
|
||||
critic 2: The critic message to critic 2
|
||||
|
||||
evolving_strategy_error_summary_v2_user: |-
|
||||
{% if queried_similar_error_knowledge|length != 0 %}
|
||||
{% for error_content, similar_error_knowledge in queried_similar_error_knowledge %}
|
||||
--------------Factor information to similar error ({{error_content}}):---------------
|
||||
{{ similar_error_knowledge[0].target_task.get_factor_information() }}
|
||||
=====Code with similar error ({{error_content}}):=====
|
||||
{{ similar_error_knowledge[0].implementation.code }}
|
||||
=====Success code to former code with similar error ({{error_content}}):=====
|
||||
{{ similar_error_knowledge[1].implementation.code }}
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
|
||||
select_implementable_factor_system: |-
|
||||
User is trying to implement some factors in the following scenario:
|
||||
{{ scenario }}
|
||||
Your job is to help the user select the easiest-to-implement factors. Some factors may be difficult to implement due to a lack of information or excessive complexity. The user will provide the number of factors you should pick and information about the factors, including their descriptions, formulas, and variable explanations.
|
||||
User will provide you the former attempt to implement the factor and the feedback to the implementation. You need to carefully review your previous attempts. Some factors have been repeatedly tried without success. You should consider discarding these factors.
|
||||
Please analyze the difficulties of the each factors and provide the reason and response the indices of selected implementable factor in the json format. Here is an example structure for the JSON output:
|
||||
{
|
||||
"Analysis": "Analyze the difficulties of the each factors and provide the reason why the factor can be implemented or not."
|
||||
"selected_factor": "The indices of selected factor index in the list, like [0, 2, 3].The length should be the number of factor left after filtering.",
|
||||
}
|
||||
|
||||
select_implementable_factor_user: |-
|
||||
Number of factor you should pick: {{ factor_num }}
|
||||
{% for factor_info in sub_tasks %}
|
||||
=============Factor index:{{factor_info[0]}}:=============
|
||||
=====Factor name:=====
|
||||
{{ factor_info[1].factor_name }}
|
||||
=====Factor description:=====
|
||||
{{ factor_info[1].factor_description }}
|
||||
=====Factor formulation:=====
|
||||
{{ factor_info[1].factor_formulation }}
|
||||
{% if factor_info[2]|length != 0 %}
|
||||
--------------Your former attempt:---------------
|
||||
{% for former_attempt in factor_info[2] %}
|
||||
=====Code to attempt {{ loop.index }}=====
|
||||
{{ former_attempt.implementation.code }}
|
||||
=====Feedback to attempt {{ loop.index }}=====
|
||||
{{ former_attempt.feedback }}
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
|
||||
analyze_component_prompt_v1_system: |-
|
||||
User is getting a new task that might consist of the components below (given in component_index: component_description):
|
||||
{{all_component_content}}
|
||||
|
||||
You should find out what components does the new task have, and put their indices in a list.
|
||||
Please response the critic in the json format. Here is an example structure for the JSON output, please strictly follow the format:
|
||||
{
|
||||
"component_no_list": the list containing indices of components.
|
||||
}
|
||||
|
||||
evaluator_output_format_system: |-
|
||||
User is trying to implement some factors in the following scenario:
|
||||
{{ scenario }}
|
||||
User will provide you the format of the output. Please help to check whether the output is align with the format.
|
||||
Please respond in the JSON format. Here is an example structure for the JSON output:
|
||||
{
|
||||
"output_format_decision": True,
|
||||
"output_format_feedback": "The output format is correct."
|
||||
}
|
||||
|
||||
|
||||
evaluator_final_decision_v1_system: |-
|
||||
User is trying to implement some factors in the following scenario:
|
||||
{{ scenario }}
|
||||
User has finished evaluation and got some feedback from the evaluator.
|
||||
The evaluator run the code and get the factor value dataframe and provide several feedback regarding user's code and code output. You should analyze the feedback and considering the scenario and factor description to give a final decision about the evaluation result. The final decision concludes whether the factor is implemented correctly and if not, detail feedback containing reason and suggestion if the final decision is False.
|
||||
|
||||
The implementation final decision is considered in the following logic:
|
||||
1. If the value and the ground truth value are exactly the same under a small tolerance, the implementation is considered correct.
|
||||
2. If the value and the ground truth value have a high correlation on ic or rank ic, the implementation is considered correct.
|
||||
3. If no ground truth value is not provided, the implementation is considered correct if the code execution is successful and the code feedback is align with the scenario and factor description.
|
||||
|
||||
Please response the critic in the json format. Here is an example structure for the JSON output, please strictly follow the format:
|
||||
{
|
||||
"final_decision": True,
|
||||
"final_feedback": "The final feedback message",
|
||||
}
|
||||
|
||||
evaluator_final_decision_v1_user: |-
|
||||
--------------Factor information:---------------
|
||||
{{ factor_information }}
|
||||
--------------Execution feedback:---------------
|
||||
{{ execution_feedback }}
|
||||
--------------Code feedback:---------------
|
||||
{{ code_feedback }}
|
||||
--------------Factor value feedback:---------------
|
||||
{{ factor_value_feedback }}
|
||||
@@ -0,0 +1,51 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
|
||||
# render it with jinja
|
||||
from jinja2 import Environment, StrictUndefined
|
||||
|
||||
from rdagent.components.coder.factor_coder.config import FACTOR_IMPLEMENT_SETTINGS
|
||||
|
||||
TPL = """
|
||||
{{file_name}}
|
||||
```{{type_desc}}
|
||||
{{content}}
|
||||
````
|
||||
"""
|
||||
# Create a Jinja template from the string
|
||||
JJ_TPL = Environment(undefined=StrictUndefined).from_string(TPL)
|
||||
|
||||
|
||||
def get_data_folder_intro():
|
||||
"""Directly get the info of the data folder.
|
||||
It is for preparing prompting message.
|
||||
"""
|
||||
content_l = []
|
||||
for p in Path(FACTOR_IMPLEMENT_SETTINGS.file_based_execution_data_folder).iterdir():
|
||||
if p.name.endswith(".h5"):
|
||||
df = pd.read_hdf(p)
|
||||
# get df.head() as string with full width
|
||||
pd.set_option("display.max_columns", None) # or 1000
|
||||
pd.set_option("display.max_rows", None) # or 1000
|
||||
pd.set_option("display.max_colwidth", None) # or 199
|
||||
rendered = JJ_TPL.render(
|
||||
file_name=p.name,
|
||||
type_desc="generated by `pd.read_hdf(filename).head()`",
|
||||
content=df.head().to_string(),
|
||||
)
|
||||
content_l.append(rendered)
|
||||
elif p.name.endswith(".md"):
|
||||
with open(p) as f:
|
||||
content = f.read()
|
||||
rendered = JJ_TPL.render(
|
||||
file_name=p.name,
|
||||
type_desc="markdown",
|
||||
content=content,
|
||||
)
|
||||
content_l.append(rendered)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"file type {p.name} is not supported. Please implement its description function.",
|
||||
)
|
||||
return "\n ----------------- file splitter -------------\n".join(content_l)
|
||||
@@ -0,0 +1,72 @@
|
||||
# TODO: inherent from the benchmark base class
|
||||
import torch
|
||||
|
||||
from rdagent.components.coder.model_coder.model import ModelImplementation
|
||||
|
||||
|
||||
def get_data_conf(init_val):
|
||||
# TODO: design this step in the workflow
|
||||
in_dim = 1000
|
||||
in_channels = 128
|
||||
exec_config = {"model_eval_param_init": init_val}
|
||||
node_feature = torch.randn(in_dim, in_channels)
|
||||
edge_index = torch.randint(0, in_dim, (2, 2000))
|
||||
return (node_feature, edge_index), exec_config
|
||||
|
||||
|
||||
class ModelImpValEval:
|
||||
"""
|
||||
Evaluate the similarity of the model structure by changing the input and observe the output.
|
||||
|
||||
Assumption:
|
||||
- If the model structure is similar, the output will change in similar way when we change the input.
|
||||
|
||||
Challenge:
|
||||
- The key difference between it and implementing factors is that we have parameters in the layers (Factor operators often have no parameters or are given parameters).
|
||||
- we try to initialize the model param in similar value. So only the model structure is different.
|
||||
|
||||
Comparing the correlation of following sequences
|
||||
- modelA[init1](input1).hidden_out1, modelA[init1](input2).hidden_out1, ...
|
||||
- modelB[init1](input1).hidden_out1, modelB[init1](input2).hidden_out1, ...
|
||||
|
||||
For each hidden output, we can calculate a correlation. The average correlation will be the metrics.
|
||||
"""
|
||||
|
||||
def evaluate(self, gt: ModelImplementation, gen: ModelImplementation):
|
||||
round_n = 10
|
||||
|
||||
eval_pairs: list[tuple] = []
|
||||
|
||||
# run different input value
|
||||
for _ in range(round_n):
|
||||
# run different model initial parameters.
|
||||
for init_val in [-0.2, -0.1, 0.1, 0.2]:
|
||||
data, exec_config = get_data_conf(init_val)
|
||||
gt_res = gt.execute(data=data, config=exec_config)
|
||||
res = gen.execute(data=data, config=exec_config)
|
||||
eval_pairs.append((res, gt_res))
|
||||
|
||||
# flat and concat the output
|
||||
res_batch, gt_res_batch = [], []
|
||||
for res, gt_res in eval_pairs:
|
||||
res_batch.append(res.reshape(-1))
|
||||
gt_res_batch.append(gt_res.reshape(-1))
|
||||
res_batch = torch.stack(res_batch)
|
||||
gt_res_batch = torch.stack(gt_res_batch)
|
||||
|
||||
res_batch = res_batch.detach().numpy()
|
||||
gt_res_batch = gt_res_batch.detach().numpy()
|
||||
|
||||
# pearson correlation of each hidden output
|
||||
def norm(x):
|
||||
return (x - x.mean(axis=0)) / x.std(axis=0)
|
||||
|
||||
dim_corr = (norm(res_batch) * norm(gt_res_batch)).mean(axis=0) # the correlation of each hidden output
|
||||
|
||||
# aggregate all the correlation
|
||||
avr_corr = dim_corr.mean()
|
||||
# FIXME:
|
||||
# It is too high(e.g. 0.944) .
|
||||
# Check if it is not a good evaluation!!
|
||||
# Maybe all the same initial params will results in extreamly high correlation without regard to the model structure.
|
||||
return avr_corr
|
||||
@@ -0,0 +1,134 @@
|
||||
import math
|
||||
from typing import Any, Callable, Dict, Optional, Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import Parameter
|
||||
from torch_geometric.nn.conv import GCNConv, MessagePassing
|
||||
from torch_geometric.nn.inits import zeros
|
||||
from torch_geometric.nn.resolver import activation_resolver
|
||||
from torch_geometric.typing import Adj
|
||||
|
||||
|
||||
class AntiSymmetricConv(torch.nn.Module):
|
||||
r"""The anti-symmetric graph convolutional operator from the
|
||||
`"Anti-Symmetric DGN: a stable architecture for Deep Graph Networks"
|
||||
<https://openreview.net/forum?id=J3Y7cgZOOS>`_ paper.
|
||||
|
||||
.. math::
|
||||
\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left(
|
||||
(\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i +
|
||||
\Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),
|
||||
|
||||
where :math:`\Phi(\mathbf{X}, \mathcal{N}_i)` denotes a
|
||||
:class:`~torch.nn.conv.MessagePassing` layer.
|
||||
|
||||
Args:
|
||||
in_channels (int): Size of each input sample.
|
||||
phi (MessagePassing, optional): The message passing module
|
||||
:math:`\Phi`. If set to :obj:`None`, will use a
|
||||
:class:`~torch_geometric.nn.conv.GCNConv` layer as default.
|
||||
(default: :obj:`None`)
|
||||
num_iters (int, optional): The number of times the anti-symmetric deep
|
||||
graph network operator is called. (default: :obj:`1`)
|
||||
epsilon (float, optional): The discretization step size
|
||||
:math:`\epsilon`. (default: :obj:`0.1`)
|
||||
gamma (float, optional): The strength of the diffusion :math:`\gamma`.
|
||||
It regulates the stability of the method. (default: :obj:`0.1`)
|
||||
act (str, optional): The non-linear activation function :math:`\sigma`,
|
||||
*e.g.*, :obj:`"tanh"` or :obj:`"relu"`. (default: :class:`"tanh"`)
|
||||
act_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
respective activation function defined by :obj:`act`.
|
||||
(default: :obj:`None`)
|
||||
bias (bool, optional): If set to :obj:`False`, the layer will not learn
|
||||
an additive bias. (default: :obj:`True`)
|
||||
|
||||
Shapes:
|
||||
- **input:**
|
||||
node features :math:`(|\mathcal{V}|, F_{in})`,
|
||||
edge indices :math:`(2, |\mathcal{E}|)`,
|
||||
edge weights :math:`(|\mathcal{E}|)` *(optional)*
|
||||
- **output:** node features :math:`(|\mathcal{V}|, F_{in})`
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
phi: Optional[MessagePassing] = None,
|
||||
num_iters: int = 1,
|
||||
epsilon: float = 0.1,
|
||||
gamma: float = 0.1,
|
||||
act: Union[str, Callable, None] = "tanh",
|
||||
act_kwargs: Optional[Dict[str, Any]] = None,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.num_iters = num_iters
|
||||
self.gamma = gamma
|
||||
self.epsilon = epsilon
|
||||
self.act = activation_resolver(act, **(act_kwargs or {}))
|
||||
|
||||
if phi is None:
|
||||
phi = GCNConv(in_channels, in_channels, bias=False)
|
||||
|
||||
self.W = Parameter(torch.empty(in_channels, in_channels))
|
||||
self.register_buffer("eye", torch.eye(in_channels))
|
||||
self.phi = phi
|
||||
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(in_channels))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
torch.nn.init.kaiming_uniform_(self.W, a=math.sqrt(5))
|
||||
self.phi.reset_parameters()
|
||||
zeros(self.bias)
|
||||
|
||||
def forward(self, x: Tensor, edge_index: Adj, *args, **kwargs) -> Tensor:
|
||||
r"""Runs the forward pass of the module."""
|
||||
antisymmetric_W = self.W - self.W.t() - self.gamma * self.eye
|
||||
|
||||
for _ in range(self.num_iters):
|
||||
h = self.phi(x, edge_index, *args, **kwargs)
|
||||
h = x @ antisymmetric_W.t() + h
|
||||
|
||||
if self.bias is not None:
|
||||
h += self.bias
|
||||
|
||||
if self.act is not None:
|
||||
h = self.act(h)
|
||||
|
||||
x = x + self.epsilon * h
|
||||
|
||||
return x
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"{self.in_channels}, "
|
||||
f"phi={self.phi}, "
|
||||
f"num_iters={self.num_iters}, "
|
||||
f"epsilon={self.epsilon}, "
|
||||
f"gamma={self.gamma})"
|
||||
)
|
||||
|
||||
|
||||
model_cls = AntiSymmetricConv
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = AntiSymmetricConv(in_channels=node_features.size(-1))
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,89 @@
|
||||
import copy
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch_geometric.nn.conv import MessagePassing
|
||||
|
||||
|
||||
class DirGNNConv(torch.nn.Module):
|
||||
r"""A generic wrapper for computing graph convolution on directed
|
||||
graphs as described in the `"Edge Directionality Improves Learning on
|
||||
Heterophilic Graphs" <https://arxiv.org/abs/2305.10498>`_ paper.
|
||||
:class:`DirGNNConv` will pass messages both from source nodes to target
|
||||
nodes and from target nodes to source nodes.
|
||||
|
||||
Args:
|
||||
conv (MessagePassing): The underlying
|
||||
:class:`~torch_geometric.nn.conv.MessagePassing` layer to use.
|
||||
alpha (float, optional): The alpha coefficient used to weight the
|
||||
aggregations of in- and out-edges as part of a convex combination.
|
||||
(default: :obj:`0.5`)
|
||||
root_weight (bool, optional): If set to :obj:`True`, the layer will add
|
||||
transformed root node features to the output.
|
||||
(default: :obj:`True`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conv: MessagePassing,
|
||||
alpha: float = 0.5,
|
||||
root_weight: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.alpha = alpha
|
||||
self.root_weight = root_weight
|
||||
|
||||
self.conv_in = copy.deepcopy(conv)
|
||||
self.conv_out = copy.deepcopy(conv)
|
||||
|
||||
if hasattr(conv, "add_self_loops"):
|
||||
self.conv_in.add_self_loops = False
|
||||
self.conv_out.add_self_loops = False
|
||||
if hasattr(conv, "root_weight"):
|
||||
self.conv_in.root_weight = False
|
||||
self.conv_out.root_weight = False
|
||||
|
||||
if root_weight:
|
||||
self.lin = torch.nn.Linear(conv.in_channels, conv.out_channels)
|
||||
else:
|
||||
self.lin = None
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
self.conv_in.reset_parameters()
|
||||
self.conv_out.reset_parameters()
|
||||
if self.lin is not None:
|
||||
self.lin.reset_parameters()
|
||||
|
||||
def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
|
||||
"""""" # noqa: D419
|
||||
x_in = self.conv_in(x, edge_index)
|
||||
x_out = self.conv_out(x, edge_index.flip([0]))
|
||||
|
||||
out = self.alpha * x_out + (1 - self.alpha) * x_in
|
||||
|
||||
if self.root_weight:
|
||||
out = out + self.lin(x)
|
||||
|
||||
return out
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}({self.conv_in}, alpha={self.alpha})"
|
||||
|
||||
|
||||
model_cls = DirGNNConv
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = DirGNNConv(MessagePassing())
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,198 @@
|
||||
import inspect
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch.nn import Dropout, Linear, Sequential
|
||||
from torch_geometric.nn.attention import PerformerAttention
|
||||
from torch_geometric.nn.conv import MessagePassing
|
||||
from torch_geometric.nn.inits import reset
|
||||
from torch_geometric.nn.resolver import activation_resolver, normalization_resolver
|
||||
from torch_geometric.typing import Adj
|
||||
from torch_geometric.utils import to_dense_batch
|
||||
|
||||
|
||||
class GPSConv(torch.nn.Module):
|
||||
r"""The general, powerful, scalable (GPS) graph transformer layer from the
|
||||
`"Recipe for a General, Powerful, Scalable Graph Transformer"
|
||||
<https://arxiv.org/abs/2205.12454>`_ paper.
|
||||
|
||||
The GPS layer is based on a 3-part recipe:
|
||||
|
||||
1. Inclusion of positional (PE) and structural encodings (SE) to the input
|
||||
features (done in a pre-processing step via
|
||||
:class:`torch_geometric.transforms`).
|
||||
2. A local message passing layer (MPNN) that operates on the input graph.
|
||||
3. A global attention layer that operates on the entire graph.
|
||||
|
||||
.. note::
|
||||
|
||||
For an example of using :class:`GPSConv`, see
|
||||
`examples/graph_gps.py
|
||||
<https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
|
||||
graph_gps.py>`_.
|
||||
|
||||
Args:
|
||||
channels (int): Size of each input sample.
|
||||
conv (MessagePassing, optional): The local message passing layer.
|
||||
heads (int, optional): Number of multi-head-attentions.
|
||||
(default: :obj:`1`)
|
||||
dropout (float, optional): Dropout probability of intermediate
|
||||
embeddings. (default: :obj:`0.`)
|
||||
act (str or Callable, optional): The non-linear activation function to
|
||||
use. (default: :obj:`"relu"`)
|
||||
act_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
respective activation function defined by :obj:`act`.
|
||||
(default: :obj:`None`)
|
||||
norm (str or Callable, optional): The normalization function to
|
||||
use. (default: :obj:`"batch_norm"`)
|
||||
norm_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
respective normalization function defined by :obj:`norm`.
|
||||
(default: :obj:`None`)
|
||||
attn_type (str): Global attention type, :obj:`multihead` or
|
||||
:obj:`performer`. (default: :obj:`multihead`)
|
||||
attn_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
attention layer. (default: :obj:`None`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
conv: Optional[MessagePassing],
|
||||
heads: int = 1,
|
||||
dropout: float = 0.0,
|
||||
act: str = "relu",
|
||||
act_kwargs: Optional[Dict[str, Any]] = None,
|
||||
norm: Optional[str] = "batch_norm",
|
||||
norm_kwargs: Optional[Dict[str, Any]] = None,
|
||||
attn_type: str = "multihead",
|
||||
attn_kwargs: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.channels = channels
|
||||
self.conv = conv
|
||||
self.heads = heads
|
||||
self.dropout = dropout
|
||||
self.attn_type = attn_type
|
||||
|
||||
attn_kwargs = attn_kwargs or {}
|
||||
if attn_type == "multihead":
|
||||
self.attn = torch.nn.MultiheadAttention(
|
||||
channels,
|
||||
heads,
|
||||
batch_first=True,
|
||||
**attn_kwargs,
|
||||
)
|
||||
elif attn_type == "performer":
|
||||
self.attn = PerformerAttention(
|
||||
channels=channels,
|
||||
heads=heads,
|
||||
**attn_kwargs,
|
||||
)
|
||||
else:
|
||||
# TODO: Support BigBird
|
||||
raise ValueError(f"{attn_type} is not supported")
|
||||
|
||||
self.mlp = Sequential(
|
||||
Linear(channels, channels * 2),
|
||||
activation_resolver(act, **(act_kwargs or {})),
|
||||
Dropout(dropout),
|
||||
Linear(channels * 2, channels),
|
||||
Dropout(dropout),
|
||||
)
|
||||
|
||||
norm_kwargs = norm_kwargs or {}
|
||||
self.norm1 = normalization_resolver(norm, channels, **norm_kwargs)
|
||||
self.norm2 = normalization_resolver(norm, channels, **norm_kwargs)
|
||||
self.norm3 = normalization_resolver(norm, channels, **norm_kwargs)
|
||||
|
||||
self.norm_with_batch = False
|
||||
if self.norm1 is not None:
|
||||
signature = inspect.signature(self.norm1.forward)
|
||||
self.norm_with_batch = "batch" in signature.parameters
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
if self.conv is not None:
|
||||
self.conv.reset_parameters()
|
||||
self.attn._reset_parameters()
|
||||
reset(self.mlp)
|
||||
if self.norm1 is not None:
|
||||
self.norm1.reset_parameters()
|
||||
if self.norm2 is not None:
|
||||
self.norm2.reset_parameters()
|
||||
if self.norm3 is not None:
|
||||
self.norm3.reset_parameters()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
edge_index: Adj,
|
||||
batch: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> Tensor:
|
||||
r"""Runs the forward pass of the module."""
|
||||
hs = []
|
||||
if self.conv is not None: # Local MPNN.
|
||||
h = self.conv(x, edge_index, **kwargs)
|
||||
h = F.dropout(h, p=self.dropout, training=self.training)
|
||||
h = h + x
|
||||
if self.norm1 is not None:
|
||||
if self.norm_with_batch:
|
||||
h = self.norm1(h, batch=batch)
|
||||
else:
|
||||
h = self.norm1(h)
|
||||
hs.append(h)
|
||||
|
||||
# Global attention transformer-style model.
|
||||
h, mask = to_dense_batch(x, batch)
|
||||
|
||||
if isinstance(self.attn, torch.nn.MultiheadAttention):
|
||||
h, _ = self.attn(h, h, h, key_padding_mask=~mask, need_weights=False)
|
||||
elif isinstance(self.attn, PerformerAttention):
|
||||
h = self.attn(h, mask=mask)
|
||||
|
||||
h = h[mask]
|
||||
h = F.dropout(h, p=self.dropout, training=self.training)
|
||||
h = h + x # Residual connection.
|
||||
if self.norm2 is not None:
|
||||
if self.norm_with_batch:
|
||||
h = self.norm2(h, batch=batch)
|
||||
else:
|
||||
h = self.norm2(h)
|
||||
hs.append(h)
|
||||
|
||||
out = sum(hs) # Combine local and global outputs.
|
||||
|
||||
out = out + self.mlp(out)
|
||||
if self.norm3 is not None:
|
||||
if self.norm_with_batch:
|
||||
out = self.norm3(out, batch=batch)
|
||||
else:
|
||||
out = self.norm3(out)
|
||||
|
||||
return out
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}({self.channels}, "
|
||||
f"conv={self.conv}, heads={self.heads}, "
|
||||
f"attn_type={self.attn_type})"
|
||||
)
|
||||
|
||||
|
||||
model_cls = GPSConv
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = GPSConv(channels=node_features.size(-1), conv=MessagePassing())
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,187 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import BatchNorm1d, Parameter
|
||||
from torch_geometric.nn import inits
|
||||
from torch_geometric.nn.conv import MessagePassing
|
||||
from torch_geometric.nn.models import MLP
|
||||
from torch_geometric.typing import Adj, OptTensor
|
||||
from torch_geometric.utils import spmm
|
||||
|
||||
|
||||
class SparseLinear(MessagePassing):
|
||||
def __init__(self, in_channels: int, out_channels: int, bias: bool = True):
|
||||
super().__init__(aggr="add")
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.weight = Parameter(torch.empty(in_channels, out_channels))
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(out_channels))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
inits.kaiming_uniform(self.weight, fan=self.in_channels, a=math.sqrt(5))
|
||||
inits.uniform(self.in_channels, self.bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
edge_index: Adj,
|
||||
edge_weight: OptTensor = None,
|
||||
) -> Tensor:
|
||||
# propagate_type: (weight: Tensor, edge_weight: OptTensor)
|
||||
out = self.propagate(edge_index, weight=self.weight, edge_weight=edge_weight)
|
||||
|
||||
if self.bias is not None:
|
||||
out = out + self.bias
|
||||
|
||||
return out
|
||||
|
||||
def message(self, weight_j: Tensor, edge_weight: OptTensor) -> Tensor:
|
||||
if edge_weight is None:
|
||||
return weight_j
|
||||
else:
|
||||
return edge_weight.view(-1, 1) * weight_j
|
||||
|
||||
def message_and_aggregate(self, adj_t: Adj, weight: Tensor) -> Tensor:
|
||||
return spmm(adj_t, weight, reduce=self.aggr)
|
||||
|
||||
|
||||
class LINKX(torch.nn.Module):
|
||||
r"""The LINKX model from the `"Large Scale Learning on Non-Homophilous
|
||||
Graphs: New Benchmarks and Strong Simple Methods"
|
||||
<https://arxiv.org/abs/2110.14446>`_ paper.
|
||||
|
||||
.. math::
|
||||
\mathbf{H}_{\mathbf{A}} &= \textrm{MLP}_{\mathbf{A}}(\mathbf{A})
|
||||
|
||||
\mathbf{H}_{\mathbf{X}} &= \textrm{MLP}_{\mathbf{X}}(\mathbf{X})
|
||||
|
||||
\mathbf{Y} &= \textrm{MLP}_{f} \left( \sigma \left( \mathbf{W}
|
||||
[\mathbf{H}_{\mathbf{A}}, \mathbf{H}_{\mathbf{X}}] +
|
||||
\mathbf{H}_{\mathbf{A}} + \mathbf{H}_{\mathbf{X}} \right) \right)
|
||||
|
||||
.. note::
|
||||
|
||||
For an example of using LINKX, see `examples/linkx.py <https://
|
||||
github.com/pyg-team/pytorch_geometric/blob/master/examples/linkx.py>`_.
|
||||
|
||||
Args:
|
||||
num_nodes (int): The number of nodes in the graph.
|
||||
in_channels (int): Size of each input sample, or :obj:`-1` to derive
|
||||
the size from the first input(s) to the forward method.
|
||||
hidden_channels (int): Size of each hidden sample.
|
||||
out_channels (int): Size of each output sample.
|
||||
num_layers (int): Number of layers of :math:`\textrm{MLP}_{f}`.
|
||||
num_edge_layers (int, optional): Number of layers of
|
||||
:math:`\textrm{MLP}_{\mathbf{A}}`. (default: :obj:`1`)
|
||||
num_node_layers (int, optional): Number of layers of
|
||||
:math:`\textrm{MLP}_{\mathbf{X}}`. (default: :obj:`1`)
|
||||
dropout (float, optional): Dropout probability of each hidden
|
||||
embedding. (default: :obj:`0.0`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_nodes: int,
|
||||
in_channels: int,
|
||||
hidden_channels: int,
|
||||
out_channels: int,
|
||||
num_layers: int,
|
||||
num_edge_layers: int = 1,
|
||||
num_node_layers: int = 1,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.num_nodes = num_nodes
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_edge_layers = num_edge_layers
|
||||
|
||||
self.edge_lin = SparseLinear(num_nodes, hidden_channels)
|
||||
|
||||
if self.num_edge_layers > 1:
|
||||
self.edge_norm = BatchNorm1d(hidden_channels)
|
||||
channels = [hidden_channels] * num_edge_layers
|
||||
self.edge_mlp = MLP(channels, dropout=0.0, act_first=True)
|
||||
else:
|
||||
self.edge_norm = None
|
||||
self.edge_mlp = None
|
||||
|
||||
channels = [in_channels] + [hidden_channels] * num_node_layers
|
||||
self.node_mlp = MLP(channels, dropout=0.0, act_first=True)
|
||||
|
||||
self.cat_lin1 = torch.nn.Linear(hidden_channels, hidden_channels)
|
||||
self.cat_lin2 = torch.nn.Linear(hidden_channels, hidden_channels)
|
||||
|
||||
channels = [hidden_channels] * num_layers + [out_channels]
|
||||
self.final_mlp = MLP(channels, dropout=dropout, act_first=True)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
self.edge_lin.reset_parameters()
|
||||
if self.edge_norm is not None:
|
||||
self.edge_norm.reset_parameters()
|
||||
if self.edge_mlp is not None:
|
||||
self.edge_mlp.reset_parameters()
|
||||
self.node_mlp.reset_parameters()
|
||||
self.cat_lin1.reset_parameters()
|
||||
self.cat_lin2.reset_parameters()
|
||||
self.final_mlp.reset_parameters()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: OptTensor,
|
||||
edge_index: Adj,
|
||||
edge_weight: OptTensor = None,
|
||||
) -> Tensor:
|
||||
"""""" # noqa: D419
|
||||
out = self.edge_lin(edge_index, edge_weight)
|
||||
|
||||
if self.edge_norm is not None and self.edge_mlp is not None:
|
||||
out = out.relu_()
|
||||
out = self.edge_norm(out)
|
||||
out = self.edge_mlp(out)
|
||||
|
||||
out = out + self.cat_lin1(out)
|
||||
|
||||
if x is not None:
|
||||
x = self.node_mlp(x)
|
||||
out = out + x
|
||||
out = out + self.cat_lin2(x)
|
||||
|
||||
return self.final_mlp(out.relu_())
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}(num_nodes={self.num_nodes}, "
|
||||
f"in_channels={self.in_channels}, "
|
||||
f"out_channels={self.out_channels})"
|
||||
)
|
||||
|
||||
|
||||
model_cls = LINKX
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = LINKX(
|
||||
num_nodes=node_features.size(0),
|
||||
in_channels=node_features.size(1),
|
||||
hidden_channels=node_features.size(1),
|
||||
out_channels=node_features.size(1),
|
||||
num_layers=1,
|
||||
)
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,118 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch_geometric.nn import SimpleConv
|
||||
from torch_geometric.nn.dense.linear import Linear
|
||||
|
||||
|
||||
class PMLP(torch.nn.Module):
|
||||
r"""The P(ropagational)MLP model from the `"Graph Neural Networks are
|
||||
Inherently Good Generalizers: Insights by Bridging GNNs and MLPs"
|
||||
<https://arxiv.org/abs/2212.09034>`_ paper.
|
||||
:class:`PMLP` is identical to a standard MLP during training, but then
|
||||
adopts a GNN architecture during testing.
|
||||
|
||||
Args:
|
||||
in_channels (int): Size of each input sample.
|
||||
hidden_channels (int): Size of each hidden sample.
|
||||
out_channels (int): Size of each output sample.
|
||||
num_layers (int): The number of layers.
|
||||
dropout (float, optional): Dropout probability of each hidden
|
||||
embedding. (default: :obj:`0.`)
|
||||
norm (bool, optional): If set to :obj:`False`, will not apply batch
|
||||
normalization. (default: :obj:`True`)
|
||||
bias (bool, optional): If set to :obj:`False`, the module
|
||||
will not learn additive biases. (default: :obj:`True`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
hidden_channels: int,
|
||||
out_channels: int,
|
||||
num_layers: int,
|
||||
dropout: float = 0.0,
|
||||
norm: bool = True,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_layers = num_layers
|
||||
self.dropout = dropout
|
||||
self.bias = bias
|
||||
|
||||
self.lins = torch.nn.ModuleList()
|
||||
self.lins.append(Linear(in_channels, hidden_channels, self.bias))
|
||||
for _ in range(self.num_layers - 2):
|
||||
lin = Linear(hidden_channels, hidden_channels, self.bias)
|
||||
self.lins.append(lin)
|
||||
self.lins.append(Linear(hidden_channels, out_channels, self.bias))
|
||||
|
||||
self.norm = None
|
||||
if norm:
|
||||
self.norm = torch.nn.BatchNorm1d(
|
||||
hidden_channels,
|
||||
affine=False,
|
||||
track_running_stats=False,
|
||||
)
|
||||
|
||||
self.conv = SimpleConv(aggr="mean", combine_root="self_loop")
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
for lin in self.lins:
|
||||
torch.nn.init.xavier_uniform_(lin.weight, gain=1.414)
|
||||
if self.bias:
|
||||
torch.nn.init.zeros_(lin.bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
edge_index: Optional[Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""""" # noqa: D419
|
||||
if not self.training and edge_index is None:
|
||||
raise ValueError(f"'edge_index' needs to be present during " f"inference in '{self.__class__.__name__}'")
|
||||
|
||||
for i in range(self.num_layers):
|
||||
x = x @ self.lins[i].weight.t()
|
||||
if not self.training:
|
||||
x = self.conv(x, edge_index)
|
||||
if self.bias:
|
||||
x = x + self.lins[i].bias
|
||||
if i != self.num_layers - 1:
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
x = x.relu()
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
|
||||
return x
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}({self.in_channels}, " f"{self.out_channels}, num_layers={self.num_layers})"
|
||||
|
||||
|
||||
model_cls = PMLP
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = PMLP(
|
||||
in_channels=node_features.size(-1),
|
||||
hidden_channels=node_features.size(-1),
|
||||
out_channels=node_features.size(-1),
|
||||
num_layers=1,
|
||||
)
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,13 @@
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class ModelImplSettings(BaseSettings):
|
||||
workspace_path: Path = Path("./git_ignore_folder/model_imp_workspace/") # Added type annotation for work_space
|
||||
|
||||
class Config:
|
||||
env_prefix = "MODEL_IMPL_" # Use MODEL_IMPL_ as prefix for environment variables
|
||||
|
||||
|
||||
MODEL_IMPL_SETTINGS = ModelImplSettings()
|
||||
@@ -0,0 +1,59 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def shape_evaluator(target, prediction):
|
||||
if target is None or prediction is None:
|
||||
return None, 0
|
||||
tar_shape = target.shape
|
||||
pre_shape = prediction.shape
|
||||
|
||||
diff = []
|
||||
for i in range(max(len(tar_shape), len(pre_shape))):
|
||||
dim_tar = tar_shape[i] if i < len(tar_shape) else 0
|
||||
dim_pre = pre_shape[i] if i < len(pre_shape) else 0
|
||||
diff.append(abs(dim_tar - dim_pre))
|
||||
|
||||
metric = 1 / (np.exp(np.mean(diff)) + 1)
|
||||
return diff, metric
|
||||
|
||||
|
||||
def reshape_tensor(original_tensor, target_shape):
|
||||
new_tensor = torch.zeros(target_shape)
|
||||
for i, dim in enumerate(original_tensor.shape):
|
||||
new_tensor = new_tensor.narrow(i, 0, dim).copy_(original_tensor)
|
||||
|
||||
return new_tensor
|
||||
|
||||
|
||||
def value_evaluator(target, prediction):
|
||||
if target is None or prediction is None:
|
||||
return None, 0
|
||||
tar_shape = target.shape
|
||||
pre_shape = prediction.shape
|
||||
|
||||
# Determine the shape of the padded tensors
|
||||
dims = [
|
||||
max(s1, s2)
|
||||
for s1, s2 in zip(
|
||||
tar_shape + (1,) * (len(pre_shape) - len(tar_shape)),
|
||||
pre_shape + (1,) * (len(tar_shape) - len(pre_shape)),
|
||||
)
|
||||
]
|
||||
# Reshape both tensors to the determined shape
|
||||
target = target.reshape(*tar_shape, *(1,) * (max(len(tar_shape), len(pre_shape)) - len(tar_shape)))
|
||||
prediction = prediction.reshape(*pre_shape, *(1,) * (max(len(tar_shape), len(pre_shape)) - len(pre_shape)))
|
||||
target_padded = reshape_tensor(target, dims)
|
||||
prediction_padded = reshape_tensor(prediction, dims)
|
||||
|
||||
# Calculate the mean absolute difference
|
||||
diff = torch.abs(target_padded - prediction_padded)
|
||||
metric = 1 / (1 + np.exp(torch.mean(diff).item()))
|
||||
return diff, metric
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tar = torch.rand(4, 5, 5)
|
||||
pre = torch.rand(4, 1)
|
||||
print(shape_evaluator(tar, pre))
|
||||
print(value_evaluator(tar, pre)[1])
|
||||
@@ -0,0 +1,136 @@
|
||||
"""
|
||||
This is just an exmaple.
|
||||
It will be replaced wtih a list of ground truth tasks.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Any, Callable, Dict, Optional, Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import Parameter
|
||||
from torch_geometric.nn.conv import GCNConv, MessagePassing
|
||||
from torch_geometric.nn.inits import zeros
|
||||
from torch_geometric.nn.resolver import activation_resolver
|
||||
from torch_geometric.typing import Adj
|
||||
|
||||
|
||||
class AntiSymmetricConv(torch.nn.Module):
|
||||
r"""The anti-symmetric graph convolutional operator from the
|
||||
`"Anti-Symmetric DGN: a stable architecture for Deep Graph Networks"
|
||||
<https://openreview.net/forum?id=J3Y7cgZOOS>`_ paper.
|
||||
|
||||
.. math::
|
||||
\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left(
|
||||
(\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i +
|
||||
\Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),
|
||||
|
||||
where :math:`\Phi(\mathbf{X}, \mathcal{N}_i)` denotes a
|
||||
:class:`~torch.nn.conv.MessagePassing` layer.
|
||||
|
||||
Args:
|
||||
in_channels (int): Size of each input sample.
|
||||
phi (MessagePassing, optional): The message passing module
|
||||
:math:`\Phi`. If set to :obj:`None`, will use a
|
||||
:class:`~torch_geometric.nn.conv.GCNConv` layer as default.
|
||||
(default: :obj:`None`)
|
||||
num_iters (int, optional): The number of times the anti-symmetric deep
|
||||
graph network operator is called. (default: :obj:`1`)
|
||||
epsilon (float, optional): The discretization step size
|
||||
:math:`\epsilon`. (default: :obj:`0.1`)
|
||||
gamma (float, optional): The strength of the diffusion :math:`\gamma`.
|
||||
It regulates the stability of the method. (default: :obj:`0.1`)
|
||||
act (str, optional): The non-linear activation function :math:`\sigma`,
|
||||
*e.g.*, :obj:`"tanh"` or :obj:`"relu"`. (default: :class:`"tanh"`)
|
||||
act_kwargs (Dict[str, Any], optional): Arguments passed to the
|
||||
respective activation function defined by :obj:`act`.
|
||||
(default: :obj:`None`)
|
||||
bias (bool, optional): If set to :obj:`False`, the layer will not learn
|
||||
an additive bias. (default: :obj:`True`)
|
||||
|
||||
Shapes:
|
||||
- **input:**
|
||||
node features :math:`(|\mathcal{V}|, F_{in})`,
|
||||
edge indices :math:`(2, |\mathcal{E}|)`,
|
||||
edge weights :math:`(|\mathcal{E}|)` *(optional)*
|
||||
- **output:** node features :math:`(|\mathcal{V}|, F_{in})`
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
phi: Optional[MessagePassing] = None,
|
||||
num_iters: int = 1,
|
||||
epsilon: float = 0.1,
|
||||
gamma: float = 0.1,
|
||||
act: Union[str, Callable, None] = "tanh",
|
||||
act_kwargs: Optional[Dict[str, Any]] = None,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.num_iters = num_iters
|
||||
self.gamma = gamma
|
||||
self.epsilon = epsilon
|
||||
self.act = activation_resolver(act, **(act_kwargs or {}))
|
||||
|
||||
if phi is None:
|
||||
phi = GCNConv(in_channels, in_channels, bias=False)
|
||||
|
||||
self.W = Parameter(torch.empty(in_channels, in_channels))
|
||||
self.register_buffer("eye", torch.eye(in_channels))
|
||||
self.phi = phi
|
||||
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(in_channels))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
r"""Resets all learnable parameters of the module."""
|
||||
torch.nn.init.kaiming_uniform_(self.W, a=math.sqrt(5))
|
||||
self.phi.reset_parameters()
|
||||
zeros(self.bias)
|
||||
|
||||
def forward(self, x: Tensor, edge_index: Adj, *args, **kwargs) -> Tensor:
|
||||
r"""Runs the forward pass of the module."""
|
||||
antisymmetric_W = self.W - self.W.t() - self.gamma * self.eye
|
||||
|
||||
for _ in range(self.num_iters):
|
||||
h = self.phi(x, edge_index, *args, **kwargs)
|
||||
h = x @ antisymmetric_W.t() + h
|
||||
|
||||
if self.bias is not None:
|
||||
h += self.bias
|
||||
|
||||
if self.act is not None:
|
||||
h = self.act(h)
|
||||
|
||||
x = x + self.epsilon * h
|
||||
|
||||
return x
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"{self.in_channels}, "
|
||||
f"phi={self.phi}, "
|
||||
f"num_iters={self.num_iters}, "
|
||||
f"epsilon={self.epsilon}, "
|
||||
f"gamma={self.gamma})"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
node_features = torch.load("node_features.pt")
|
||||
edge_index = torch.load("edge_index.pt")
|
||||
|
||||
# Model instantiation and forward pass
|
||||
model = AntiSymmetricConv(in_channels=node_features.size(-1))
|
||||
output = model(node_features, edge_index)
|
||||
|
||||
# Save output to a file
|
||||
torch.save(output, "gt_output.pt")
|
||||
@@ -0,0 +1,90 @@
|
||||
"""
|
||||
This file will be removed in the future and replaced by
|
||||
- rdagent/app/model_implementation/eval.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# randomly generate a input graph, node_feature and edge_index
|
||||
# 1000 nodes, 128 dim node feature, 2000 edges
|
||||
import torch
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
assert load_dotenv()
|
||||
formula_info = {
|
||||
"name": "Anti-Symmetric Deep Graph Network (A-DGN)",
|
||||
"description": "A framework for stable and non-dissipative DGN design. It ensures long-range information preservation between nodes and prevents gradient vanishing or explosion during training.",
|
||||
"formulation": "x_u^{(l)} = x_u^{(l-1)} + \\epsilon \\sigma \\left( W^T x_u^{(l-1)} + \\Phi(X^{(l-1)}, N_u) + b \\right)",
|
||||
"variables": {
|
||||
"x_u^{(l)}": "The state of node u at layer l",
|
||||
"\\epsilon": "The step size in the Euler discretization",
|
||||
"\\sigma": "A monotonically non-decreasing activation function",
|
||||
"W": "An anti-symmetric weight matrix",
|
||||
"X^{(l-1)}": "The node feature matrix at layer l-1",
|
||||
"N_u": "The set of neighbors of node u",
|
||||
"b": "A bias vector",
|
||||
},
|
||||
}
|
||||
|
||||
system_prompt = "You are an assistant whose job is to answer user's question."
|
||||
user_prompt = "With the following given information, write a python code using pytorch and torch_geometric to implement the model. This model is in the graph learning field, only have one layer. The input will be node_feature [num_nodes, dim_feature] and edge_index [2, num_edges], and they should be loaded from the files 'node_features.pt' and 'edge_index.pt'. There is not edge attribute or edge weight as input. The model should detect the node_feature and edge_index shape, if there is Linear transformation layer in the model, the input and output shape should be consistent. The in_channels is the dimension of the node features. You code should contain additional 'if __name__ == '__main__', where you should load the node_feature and edge_index from the files and run the model, and save the output to a file 'llm_output.pt'. Implement the model forward function based on the following information: model formula information. 1. model name: {}, 2. model description: {}, 3. model formulation: {}, 4. model variables: {}. You must complete the forward function as far as you can do.".format(
|
||||
formula_info["name"],
|
||||
formula_info["description"],
|
||||
formula_info["formulation"],
|
||||
formula_info["variables"],
|
||||
)
|
||||
|
||||
resp = APIBackend(use_chat_cache=False).build_messages_and_create_chat_completion(user_prompt, system_prompt)
|
||||
|
||||
print(resp)
|
||||
|
||||
# take the code part from the response and save it to a file, the code is covered in the ```python``` block
|
||||
code = resp.split("```python")[1].split("```")[0]
|
||||
with open("llm_code.py", "w") as f:
|
||||
f.write(code)
|
||||
|
||||
average_shape_eval = []
|
||||
average_value_eval = []
|
||||
for test_mode in ["zeros", "ones", "randn"]:
|
||||
if test_mode == "zeros":
|
||||
node_feature = torch.zeros(1000, 128)
|
||||
elif test_mode == "ones":
|
||||
node_feature = torch.ones(1000, 128)
|
||||
elif test_mode == "randn":
|
||||
node_feature = torch.randn(1000, 128)
|
||||
edge_index = torch.randint(0, 1000, (2, 2000))
|
||||
|
||||
torch.save(node_feature, "node_features.pt")
|
||||
torch.save(edge_index, "edge_index.pt")
|
||||
|
||||
try:
|
||||
os.system("python llm_code.py")
|
||||
except:
|
||||
print("Error in running the LLM code")
|
||||
os.system("python gt_code.py")
|
||||
os.system("rm edge_index.pt")
|
||||
os.system("rm node_features.pt")
|
||||
# load the output and print the shape
|
||||
|
||||
from evaluator import shape_evaluator, value_evaluator
|
||||
|
||||
try:
|
||||
llm_output = torch.load("llm_output.pt")
|
||||
except:
|
||||
llm_output = None
|
||||
gt_output = torch.load("gt_output.pt")
|
||||
|
||||
average_shape_eval.append(shape_evaluator(llm_output, gt_output)[1])
|
||||
average_value_eval.append(value_evaluator(llm_output, gt_output)[1])
|
||||
|
||||
print("Shape evaluation: ", average_shape_eval[-1])
|
||||
print("Value evaluation:super().generate(task_l) ", average_value_eval[-1])
|
||||
|
||||
os.system("rm llm_output.pt")
|
||||
os.system("rm gt_output.pt")
|
||||
os.system("rm llm_code.py")
|
||||
|
||||
print("Average shape evaluation: ", sum(average_shape_eval) / len(average_shape_eval))
|
||||
print("Average value evaluation: ", sum(average_value_eval) / len(average_value_eval))
|
||||
@@ -0,0 +1,210 @@
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Sequence
|
||||
|
||||
import torch
|
||||
|
||||
from rdagent.components.coder.model_coder.conf import MODEL_IMPL_SETTINGS
|
||||
from rdagent.components.loader.task_loader import ModelTaskLoader
|
||||
from rdagent.core.exception import CodeFormatException
|
||||
from rdagent.core.experiment import Experiment, FBImplementation, ImpLoader, Task
|
||||
from rdagent.utils import get_module_by_module_path
|
||||
|
||||
|
||||
class ModelTask(Task):
|
||||
# TODO: it should change when the Task changes.
|
||||
name: str
|
||||
description: str
|
||||
formulation: str
|
||||
variables: Dict[str, str] # map the variable name to the variable description
|
||||
|
||||
def __init__(
|
||||
self, name: str, description: str, formulation: str, variables: Dict[str, str], key: Optional[str] = None
|
||||
) -> None:
|
||||
"""
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
key : Optional[str]
|
||||
Key is a string to identify the task.
|
||||
It will be used to connect to other information(e.g. ground truth).
|
||||
"""
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.formulation = formulation
|
||||
self.variables = variables
|
||||
self.key = key
|
||||
|
||||
def get_information(self):
|
||||
return f"""name: {self.name}
|
||||
description: {self.description}
|
||||
formulation: {self.formulation}
|
||||
variables: {self.variables}
|
||||
key: {self.key}
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def from_dict(dict):
|
||||
return ModelTask(**dict)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__} {self.name}>"
|
||||
|
||||
|
||||
class ModelImplementation(FBImplementation):
|
||||
"""
|
||||
It is a Pytorch model implementation task;
|
||||
All the things are placed in a folder.
|
||||
|
||||
Folder
|
||||
- data source and documents prepared by `prepare`
|
||||
- Please note that new data may be passed in dynamically in `execute`
|
||||
- code (file `model.py` ) injected by `inject_code`
|
||||
- the `model.py` that contains a variable named `model_cls` which indicates the implemented model structure
|
||||
- `model_cls` is a instance of `torch.nn.Module`;
|
||||
|
||||
|
||||
We'll import the model in the implementation in file `model.py` after setting the cwd into the directory
|
||||
- from model import model_cls
|
||||
- initialize the model by initializing it `model_cls(input_dim=INPUT_DIM)`
|
||||
- And then verify the model.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, target_task: Task) -> None:
|
||||
super().__init__(target_task)
|
||||
self.path = None
|
||||
|
||||
def prepare(self) -> None:
|
||||
"""
|
||||
Prepare for the workspace;
|
||||
"""
|
||||
unique_id = uuid.uuid4()
|
||||
self.path = MODEL_IMPL_SETTINGS.workspace_path / f"M{unique_id}"
|
||||
# start with `M` so that it can be imported via python
|
||||
self.path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def execute(self, data=None, config: dict = {}):
|
||||
mod = get_module_by_module_path(str(self.path / "model.py"))
|
||||
try:
|
||||
model_cls = mod.model_cls
|
||||
except AttributeError:
|
||||
raise CodeFormatException("The model_cls is not implemented in the model.py")
|
||||
# model_init =
|
||||
|
||||
assert isinstance(data, tuple)
|
||||
node_feature, _ = data
|
||||
in_channels = node_feature.size(-1)
|
||||
m = model_cls(in_channels)
|
||||
|
||||
# TODO: initialize all the parameters of `m` to `model_eval_param_init`
|
||||
model_eval_param_init: float = config["model_eval_param_init"]
|
||||
|
||||
# initialize all parameters of `m` to `model_eval_param_init`
|
||||
for _, param in m.named_parameters():
|
||||
param.data.fill_(model_eval_param_init)
|
||||
|
||||
assert isinstance(data, tuple)
|
||||
return m(*data)
|
||||
|
||||
def execute_desc(self) -> str:
|
||||
return """
|
||||
The the implemented code will be placed in a file like <uuid>/model.py
|
||||
|
||||
We'll import the model in the implementation in file `model.py` after setting the cwd into the directory
|
||||
- from model import model_cls (So you must have a variable named `model_cls` in the file)
|
||||
- So your implemented code could follow the following pattern
|
||||
```Python
|
||||
class XXXLayer(torch.nn.Module):
|
||||
...
|
||||
model_cls = XXXLayer
|
||||
```
|
||||
- initialize the model by initializing it `model_cls(input_dim=INPUT_DIM)`
|
||||
- And then verify the model by comparing the output tensors by feeding specific input tensor.
|
||||
"""
|
||||
|
||||
|
||||
class ModelExperiment(Experiment[ModelTask, ModelImplementation]):
|
||||
...
|
||||
|
||||
|
||||
class ModelTaskLoaderJson(ModelTaskLoader):
|
||||
# def __init__(self, json_uri: str, select_model: Optional[str] = None) -> None:
|
||||
# super().__init__()
|
||||
# self.json_uri = json_uri
|
||||
# self.select_model = 'A-DGN'
|
||||
|
||||
# def load(self, *argT, **kwargs) -> Sequence[ModelImplTask]:
|
||||
# # json is supposed to be in the format of {model_name: dict{model_data}}
|
||||
# model_dict = json.load(open(self.json_uri, "r"))
|
||||
# if self.select_model is not None:
|
||||
# assert self.select_model in model_dict
|
||||
# model_name = self.select_model
|
||||
# model_data = model_dict[self.select_model]
|
||||
# else:
|
||||
# model_name, model_data = list(model_dict.items())[0]
|
||||
|
||||
# model_impl_task = ModelImplTask(
|
||||
# name=model_name,
|
||||
# description=model_data["description"],
|
||||
# formulation=model_data["formulation"],
|
||||
# variables=model_data["variables"],
|
||||
# key=model_name
|
||||
# )
|
||||
|
||||
# return [model_impl_task]
|
||||
|
||||
def __init__(self, json_uri: str) -> None:
|
||||
super().__init__()
|
||||
self.json_uri = json_uri
|
||||
|
||||
def load(self, *argT, **kwargs) -> Sequence[ModelTask]:
|
||||
# json is supposed to be in the format of {model_name: dict{model_data}}
|
||||
model_dict = json.load(open(self.json_uri, "r"))
|
||||
|
||||
# FIXME: the model in the json file is not right due to extraction error
|
||||
# We should fix them case by case in the future
|
||||
#
|
||||
# formula_info = {
|
||||
# "name": "Anti-Symmetric Deep Graph Network (A-DGN)",
|
||||
# "description": "A framework for stable and non-dissipative DGN design. It ensures long-range information preservation between nodes and prevents gradient vanishing or explosion during training.",
|
||||
# "formulation": r"\mathbf{x}^{\prime}_i = \mathbf{x}_i + \epsilon \cdot \sigma \left( (\mathbf{W}-\mathbf{W}^T-\gamma \mathbf{I}) \mathbf{x}_i + \Phi(\mathbf{X}, \mathcal{N}_i) + \mathbf{b}\right),",
|
||||
# "variables": {
|
||||
# r"\mathbf{x}_i": "The state of node i at previous layer",
|
||||
# r"\epsilon": "The step size in the Euler discretization",
|
||||
# r"\sigma": "A monotonically non-decreasing activation function",
|
||||
# r"\Phi": "A graph convolutional operator",
|
||||
# r"W": "An anti-symmetric weight matrix",
|
||||
# r"\mathbf{x}^{\prime}_i": "The node feature matrix at layer l-1",
|
||||
# r"\mathcal{N}_i": "The set of neighbors of node u",
|
||||
# r"\mathbf{b}": "A bias vector",
|
||||
# },
|
||||
# "key": "A-DGN",
|
||||
# }
|
||||
model_impl_task_list = []
|
||||
for model_name, model_data in model_dict.items():
|
||||
model_impl_task = ModelTask(
|
||||
name=model_name,
|
||||
description=model_data["description"],
|
||||
formulation=model_data["formulation"],
|
||||
variables=model_data["variables"],
|
||||
key=model_data["key"],
|
||||
)
|
||||
model_impl_task_list.append(model_impl_task)
|
||||
return model_impl_task_list
|
||||
|
||||
|
||||
class ModelImpLoader(ImpLoader[ModelTask, ModelImplementation]):
|
||||
def __init__(self, path: Path) -> None:
|
||||
self.path = Path(path)
|
||||
|
||||
def load(self, task: ModelTask) -> ModelImplementation:
|
||||
assert task.key is not None
|
||||
mti = ModelImplementation(task)
|
||||
mti.prepare()
|
||||
with open(self.path / f"{task.key}.py", "r") as f:
|
||||
code = f.read()
|
||||
mti.inject_code(**{"model.py": code})
|
||||
return mti
|
||||
@@ -0,0 +1,45 @@
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from jinja2 import Environment, StrictUndefined
|
||||
|
||||
from rdagent.components.coder.model_coder.model import (
|
||||
ModelExperiment,
|
||||
ModelImplementation,
|
||||
)
|
||||
from rdagent.core.prompts import Prompts
|
||||
from rdagent.core.task_generator import TaskGenerator
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
DIRNAME = Path(__file__).absolute().resolve().parent
|
||||
|
||||
|
||||
class ModelCodeWriter(TaskGenerator[ModelExperiment]):
|
||||
def generate(self, exp: ModelExperiment) -> ModelExperiment:
|
||||
mti_l = []
|
||||
for t in exp.sub_tasks:
|
||||
mti = ModelImplementation(t)
|
||||
mti.prepare()
|
||||
pr = Prompts(file_path=DIRNAME / "prompt.yaml")
|
||||
|
||||
user_prompt_tpl = Environment(undefined=StrictUndefined).from_string(pr["code_implement_user"])
|
||||
sys_prompt_tpl = Environment(undefined=StrictUndefined).from_string(pr["code_implement_sys"])
|
||||
|
||||
user_prompt = user_prompt_tpl.render(
|
||||
name=t.name,
|
||||
description=t.description,
|
||||
formulation=t.formulation,
|
||||
variables=t.variables,
|
||||
execute_desc=mti.execute_desc(),
|
||||
)
|
||||
system_prompt = sys_prompt_tpl.render()
|
||||
|
||||
resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt)
|
||||
|
||||
# Extract the code part from the response
|
||||
match = re.search(r".*```[Pp]ython\n(.*)\n```.*", resp, re.DOTALL)
|
||||
code = match.group(1)
|
||||
mti.inject_code(**{"model.py": code})
|
||||
mti_l.append(mti)
|
||||
exp.sub_implementations = mti_l
|
||||
return exp
|
||||
@@ -0,0 +1,18 @@
|
||||
|
||||
|
||||
code_implement_sys: -|
|
||||
You are an assistant whose job is to answer user's question."
|
||||
code_implement_user: -|
|
||||
With the following given information, write a python code using pytorch and torch_geometric to implement the model.
|
||||
This model is in the graph learning field, only have one layer.
|
||||
The input will be node_feature [num_nodes, dim_feature] and edge_index [2, num_edges] (It would be the input of the forward model)
|
||||
There is not edge attribute or edge weight as input. The model should detect the node_feature and edge_index shape, if there is Linear transformation layer in the model, the input and output shape should be consistent. The in_channels is the dimension of the node features.
|
||||
Implement the model forward function based on the following information:model formula information.
|
||||
1. model name:{{name}}
|
||||
2. model description:{{description}}
|
||||
3. model formulation:{{formulation}}
|
||||
4. model variables:{{variables}}.
|
||||
You must complete the forward function as far as you can do.
|
||||
\# Execution
|
||||
Your implemented code will be exectued in the follow way
|
||||
{{execute_desc}}
|
||||
@@ -0,0 +1,8 @@
|
||||
extract_model_formulation_system: |-
|
||||
offer description of the proposed model in this paper, write a latex formula with variable of the model. the format should be like " "Model Name": {
|
||||
"description": "",
|
||||
"formulation": "",
|
||||
"variables": {
|
||||
"\\hat{y}_u": "The predicted output for node u",
|
||||
}"
|
||||
such format content should be begin with ```json and end with ``` and the content should be in json format.
|
||||
@@ -0,0 +1,134 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from rdagent.components.coder.model_coder.model import ModelExperiment, ModelTask
|
||||
from rdagent.components.document_reader.document_reader import (
|
||||
load_and_process_pdfs_by_langchain,
|
||||
)
|
||||
from rdagent.components.loader.task_loader import ModelTaskLoader
|
||||
from rdagent.core.log import RDAgentLog
|
||||
from rdagent.core.prompts import Prompts
|
||||
from rdagent.oai.llm_utils import APIBackend
|
||||
|
||||
document_process_prompts = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
|
||||
|
||||
|
||||
def extract_model_from_doc(doc_content: str) -> dict:
|
||||
"""
|
||||
Extract model information from document content.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
doc_content : str
|
||||
Document content.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
{factor_name: dict{description, formulation, variables}}
|
||||
"""
|
||||
session = APIBackend().build_chat_session(
|
||||
session_system_prompt=document_process_prompts["extract_model_formulation_system"],
|
||||
)
|
||||
current_user_prompt = doc_content
|
||||
|
||||
# Extract model information from document content.
|
||||
model_dict = {}
|
||||
|
||||
for _ in range(10):
|
||||
# try to extract model information from the document content, retry at most 10 times.
|
||||
extract_result_resp = session.build_chat_completion(
|
||||
user_prompt=current_user_prompt,
|
||||
json_mode=False,
|
||||
)
|
||||
re_search_res = re.search(r"```json(.*)```", extract_result_resp, re.S)
|
||||
ret_json_str = re_search_res.group(1) if re_search_res is not None else ""
|
||||
try:
|
||||
ret_dict = json.loads(ret_json_str)
|
||||
parse_success = bool(isinstance(ret_dict, dict))
|
||||
except json.JSONDecodeError:
|
||||
parse_success = False
|
||||
if ret_json_str is None or not parse_success:
|
||||
current_user_prompt = "Your response didn't follow the instruction might be wrong json format. Try again."
|
||||
else:
|
||||
for name, formulation_and_description in ret_dict.items():
|
||||
if name not in model_dict:
|
||||
model_dict[name] = formulation_and_description
|
||||
if len(model_dict) == 0:
|
||||
current_user_prompt = "No model extracted. Please try again."
|
||||
else:
|
||||
break
|
||||
|
||||
RDAgentLog().info(f"已经完成{len(model_dict)}个模型的提取")
|
||||
|
||||
return model_dict
|
||||
|
||||
|
||||
def merge_file_to_model_dict_to_model_dict(
|
||||
file_to_model_dict: dict[str, dict],
|
||||
) -> dict:
|
||||
model_dict = {}
|
||||
for file_name in file_to_model_dict:
|
||||
for model_name in file_to_model_dict[file_name]:
|
||||
model_dict.setdefault(model_name, [])
|
||||
model_dict[model_name].append(file_to_model_dict[file_name][model_name])
|
||||
|
||||
model_dict_simple_deduplication = {}
|
||||
for model_name in model_dict:
|
||||
if len(model_dict[model_name]) > 1:
|
||||
model_dict_simple_deduplication[model_name] = max(
|
||||
model_dict[model_name],
|
||||
key=lambda x: len(x["formulation"]),
|
||||
)
|
||||
else:
|
||||
model_dict_simple_deduplication[model_name] = model_dict[model_name][0]
|
||||
return model_dict_simple_deduplication
|
||||
|
||||
|
||||
def extract_model_from_docs(docs_dict):
|
||||
model_dict = {}
|
||||
for doc_name, doc_content in docs_dict.items():
|
||||
model_dict[doc_name] = extract_model_from_doc(doc_content)
|
||||
return model_dict
|
||||
|
||||
|
||||
class ModelExperimentLoaderFromDict(ModelTaskLoader):
|
||||
def load(self, model_dict: dict) -> list:
|
||||
"""Load data from a dict."""
|
||||
task_l = []
|
||||
for model_name, model_data in model_dict.items():
|
||||
task = ModelTask(
|
||||
name=model_name,
|
||||
description=model_data["description"],
|
||||
formulation=model_data["formulation"],
|
||||
variables=model_data["variables"],
|
||||
key=model_name,
|
||||
)
|
||||
task_l.append(task)
|
||||
return ModelExperiment(sub_tasks=task_l)
|
||||
|
||||
|
||||
class ModelExperimentLoaderFromPDFfiles(ModelTaskLoader):
|
||||
def load(self, file_or_folder_path: Path) -> dict:
|
||||
docs_dict = load_and_process_pdfs_by_langchain(Path(file_or_folder_path)) # dict{file_path:content}
|
||||
model_dict = extract_model_from_docs(
|
||||
docs_dict
|
||||
) # dict{file_name: dict{model_name: dict{description, formulation, variables}}}
|
||||
model_dict = merge_file_to_model_dict_to_model_dict(
|
||||
model_dict
|
||||
) # dict {model_name: dict{description, formulation, variables}}
|
||||
return ModelExperimentLoaderFromDict().load(model_dict)
|
||||
|
||||
|
||||
def main(path="../test_doc"):
|
||||
doc_dict = load_and_process_pdfs_by_langchain(Path(path))
|
||||
print(doc_dict.keys()) # if you run code like "python -u", the print content will be truncated
|
||||
|
||||
|
||||
import fire
|
||||
|
||||
if __name__ == "__main__":
|
||||
fire.Fire(main)
|
||||
Reference in New Issue
Block a user