wehub-resource-sync
472 lines
16 KiB
Python
472 lines
16 KiB
Python
# TODO: remove `self.scen` if traces will be passed into the instance.
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from abc import ABC, abstractmethod
|
|
from typing import TYPE_CHECKING, Generic, TypeVar
|
|
|
|
from rdagent.core.conf import RD_AGENT_SETTINGS
|
|
from rdagent.core.evaluation import Feedback
|
|
from rdagent.core.experiment import (
|
|
ASpecificExp,
|
|
ASpecificPlan,
|
|
Experiment,
|
|
ExperimentPlan,
|
|
)
|
|
from rdagent.core.knowledge_base import KnowledgeBase
|
|
from rdagent.core.scenario import Scenario
|
|
|
|
if TYPE_CHECKING:
|
|
from rdagent.utils.workflow.loop import LoopBase
|
|
|
|
|
|
class Hypothesis:
|
|
"""
|
|
TODO: We may have better name for it.
|
|
|
|
Name Candidates:
|
|
- Belief
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
hypothesis: str,
|
|
reason: str,
|
|
concise_reason: str,
|
|
concise_observation: str,
|
|
concise_justification: str,
|
|
concise_knowledge: str,
|
|
) -> None:
|
|
self.hypothesis: str = hypothesis
|
|
self.reason: str = reason
|
|
self.concise_reason: str = concise_reason
|
|
self.concise_observation: str = concise_observation
|
|
self.concise_justification: str = concise_justification
|
|
self.concise_knowledge: str = concise_knowledge
|
|
|
|
def __str__(self) -> str:
|
|
return f"""Hypothesis: {self.hypothesis}
|
|
Reason: {self.reason}"""
|
|
|
|
# source: data_ana | model_nan = None
|
|
|
|
|
|
# Origin(path of repo/data/feedback) => view/summarization => generated Hypothesis
|
|
|
|
|
|
class ExperimentFeedback(Feedback):
|
|
def __init__(
|
|
self,
|
|
reason: str,
|
|
*,
|
|
code_change_summary: str | None = None,
|
|
decision: bool,
|
|
eda_improvement: str | None = None,
|
|
exception: Exception | None = None,
|
|
) -> None:
|
|
self.decision = decision
|
|
self.eda_improvement = eda_improvement
|
|
self.reason = reason
|
|
# Exception is not None means failing to generate runnable experiments due to exception.
|
|
# Runable reuslts are not always good.
|
|
self.exception: Exception | None = (
|
|
exception # if the experiment raises exception, it will be integrated into part of the feedback.
|
|
)
|
|
self.code_change_summary = code_change_summary
|
|
|
|
def __bool__(self) -> bool:
|
|
return self.decision
|
|
|
|
def __str__(self) -> str:
|
|
res = f"Decision: {self.decision}\nReason: {self.reason}"
|
|
code_change_summary = getattr(self, "code_change_summary", None)
|
|
if code_change_summary is not None:
|
|
res += "\nCode Change Summary: " + code_change_summary
|
|
return res
|
|
|
|
@classmethod
|
|
def from_exception(cls, e: Exception) -> ExperimentFeedback:
|
|
"""
|
|
A convenient method to create Feedback from an exception.
|
|
"""
|
|
return cls(decision=False, reason=f"The experiment fails due to {e!s}", exception=e)
|
|
|
|
|
|
class HypothesisFeedback(ExperimentFeedback):
|
|
def __init__(
|
|
self,
|
|
reason: str,
|
|
decision: bool,
|
|
code_change_summary: str = "",
|
|
*,
|
|
observations: str | None = None,
|
|
hypothesis_evaluation: str | None = None,
|
|
new_hypothesis: str | None = None,
|
|
eda_improvement: str | None = None,
|
|
acceptable: bool | None = None,
|
|
exception: Exception | None = None,
|
|
) -> None:
|
|
super().__init__(
|
|
reason,
|
|
decision=decision,
|
|
code_change_summary=code_change_summary,
|
|
eda_improvement=eda_improvement,
|
|
exception=exception,
|
|
)
|
|
self.observations = observations
|
|
self.hypothesis_evaluation = hypothesis_evaluation
|
|
self.new_hypothesis = new_hypothesis
|
|
self.acceptable = acceptable
|
|
|
|
def __str__(self) -> str:
|
|
upper_str = f"""{super().__str__()}"""
|
|
if self.observations is not None:
|
|
upper_str += f"\nObservations: {self.observations}"
|
|
if self.hypothesis_evaluation is not None:
|
|
upper_str += f"\nHypothesis Evaluation: {self.hypothesis_evaluation}"
|
|
if self.new_hypothesis is not None:
|
|
upper_str += f"\nNew Hypothesis: {self.new_hypothesis}"
|
|
if self.eda_improvement is not None:
|
|
upper_str += f"\nEDA Improvement: {self.eda_improvement}"
|
|
if self.acceptable is not None:
|
|
upper_str += f"\nOverall Acceptable: {self.acceptable}"
|
|
return upper_str
|
|
|
|
|
|
ASpecificScen = TypeVar("ASpecificScen", bound=Scenario)
|
|
ASpecificKB = TypeVar("ASpecificKB", bound=KnowledgeBase)
|
|
|
|
|
|
class Trace(Generic[ASpecificScen, ASpecificKB]):
|
|
NodeType = tuple[Experiment, ExperimentFeedback] # Define NodeType as a new type representing the tuple
|
|
NEW_ROOT: tuple = ()
|
|
SEL_LATEST_SOTA: tuple = (-1,) # select the SOTA experiment in latest node
|
|
|
|
def __init__(self, scen: ASpecificScen, knowledge_base: ASpecificKB | None = None) -> None:
|
|
self.scen: ASpecificScen = scen
|
|
|
|
# BEGIN: graph structure -------------------------
|
|
self.hist: list[Trace.NodeType] = (
|
|
[]
|
|
) # List of tuples containing experiments and their feedback, organized over time.
|
|
self.dag_parent: list[tuple[int, ...]] = [] # List of tuples representing parent indices in the DAG structure.
|
|
# Definition:
|
|
# - (,) represents no parent (root node in one tree);
|
|
# - (1,) presents one parent;
|
|
# - (1, 2) represents two parents (Multiple parent is not implemented yet).
|
|
# Syntax sugar for the parent relationship:
|
|
# - Only for selection:
|
|
# - (-1,) indicates that select the last record node as parent.
|
|
|
|
# NOTE: the sequence of hist and dag_parent is organized by the order to record the experiment.
|
|
# So it may be different from the order of the loop_id.
|
|
# So we need an extra mapping to map the enqueue id back to the loop id.
|
|
self.idx2loop_id: dict[int, int] = {}
|
|
|
|
# Design discussion:
|
|
# - If we unifiy the loop_id and the enqueue id, we will have less recognition burden.
|
|
# - If we use different id for loop and enqueue, we don't have to handle the placeholder logic.
|
|
# END: graph structure -------------------------
|
|
|
|
# TODO: self.hist is 2-tuple now, remove hypothesis from it, change old code for this later.
|
|
self.knowledge_base: ASpecificKB | None = knowledge_base
|
|
|
|
# The next expending point of the selection. Set it as a state of the trace will make
|
|
self.current_selection: tuple[int, ...] = self.SEL_LATEST_SOTA
|
|
|
|
def get_sota_hypothesis_and_experiment(self) -> tuple[Hypothesis | None, Experiment | None]:
|
|
"""Access the last experiment result, sub-task, and the corresponding hypothesis."""
|
|
# TODO: The return value does not align with the signature.
|
|
for experiment, feedback in self.hist[::-1]:
|
|
if feedback.decision:
|
|
return experiment.hypothesis, experiment
|
|
|
|
return None, None
|
|
|
|
def is_selection_new_tree(self, selection: tuple[int, ...] | None = None) -> bool:
|
|
"""
|
|
Check if the current trace is a new tree.
|
|
- selection maybe (-1,) when the dag_parent is empty.
|
|
"""
|
|
if selection is None:
|
|
selection = self.get_current_selection()
|
|
|
|
return selection == self.NEW_ROOT or len(self.dag_parent) == 0
|
|
|
|
def get_current_selection(self) -> tuple[int, ...]:
|
|
return self.current_selection
|
|
|
|
def set_current_selection(self, selection: tuple[int, ...]) -> None:
|
|
self.current_selection = selection
|
|
|
|
def get_parent_exps(
|
|
self,
|
|
selection: tuple[int, ...] | None = None,
|
|
) -> list[Trace.NodeType]:
|
|
"""
|
|
Collect all ancestors of the given selection.
|
|
The return list follows the order of [root->...->parent->current_node].
|
|
"""
|
|
if selection is None:
|
|
selection = self.get_current_selection()
|
|
|
|
if self.is_selection_new_tree(selection):
|
|
return []
|
|
|
|
return [self.hist[i] for i in self.get_parents(selection[0])]
|
|
|
|
def exp2idx(self, exp: Experiment | list[Experiment]) -> int | list[int] | None:
|
|
if isinstance(exp, list):
|
|
exps: list[Experiment] = exp
|
|
|
|
# keep the order
|
|
exp_to_index: dict[Experiment, int] = {_exp: i for i, (_exp, _) in enumerate(self.hist)}
|
|
return [exp_to_index[_exp] for _exp in exps]
|
|
for i, (_exp, _) in enumerate(self.hist):
|
|
if _exp == exp:
|
|
return i
|
|
return None
|
|
|
|
def idx2exp(self, idx: int | list[int]) -> Experiment | list[Experiment]:
|
|
if isinstance(idx, list):
|
|
idxs: list[int] = idx
|
|
return [self.hist[_idx][0] for _idx in idxs]
|
|
return self.hist[idx][0]
|
|
|
|
def is_parent(self, parent_idx: int, child_idx: int) -> bool:
|
|
ancestors = self.get_parents(child_idx)
|
|
return parent_idx in ancestors
|
|
|
|
def get_parents(self, child_idx: int) -> list[int]:
|
|
if self.is_selection_new_tree((child_idx,)):
|
|
return []
|
|
|
|
ancestors: list[int] = []
|
|
curr = child_idx
|
|
while True:
|
|
ancestors.insert(0, curr)
|
|
parent_tuple = self.dag_parent[curr]
|
|
if not parent_tuple or parent_tuple[0] == curr:
|
|
break
|
|
curr = parent_tuple[0]
|
|
|
|
return ancestors
|
|
|
|
def sync_dag_parent_and_hist(
|
|
self,
|
|
exp_and_fb: NodeType,
|
|
cur_loop_id: int,
|
|
) -> None:
|
|
"""
|
|
Adding corresponding parent index to the dag_parent when the hist is going to be changed.
|
|
Should be called when the hist is changed.
|
|
"""
|
|
# Prioritize local_selection from the experiment if available
|
|
exp = exp_and_fb[0]
|
|
selection = getattr(exp, "local_selection", None)
|
|
if selection is None:
|
|
selection = self.get_current_selection()
|
|
|
|
if len(self.hist) == 0 or len(selection) == 0:
|
|
# the node we are going to add is the first node of hist / root node of a new sub-trace
|
|
self.dag_parent.append(self.NEW_ROOT)
|
|
|
|
else:
|
|
current_node_idx = selection[0]
|
|
|
|
if current_node_idx == -1:
|
|
# the current selection is the latest one
|
|
current_node_idx = len(self.hist) - 1
|
|
|
|
self.dag_parent.append((current_node_idx,))
|
|
self.hist.append(exp_and_fb)
|
|
self.idx2loop_id[len(self.hist) - 1] = cur_loop_id
|
|
|
|
def get_children(self, parent_idx: int | None = None) -> list[NodeType]:
|
|
"""
|
|
Get all children nodes for a given parent index.
|
|
If parent_idx is None, returns the root nodes (experiments starting from scratch).
|
|
"""
|
|
target_parents = (parent_idx,) if parent_idx is not None else self.NEW_ROOT
|
|
children = []
|
|
for i, parents in enumerate(self.dag_parent):
|
|
if parents == target_parents and i < len(self.hist):
|
|
children.append(self.hist[i])
|
|
return children
|
|
|
|
def get_sota_experiment(self, node_id: int | None = None) -> Experiment | None:
|
|
"""
|
|
Get the SOTA experiment from the trace by traversing ancestors backwards from node_id.
|
|
"""
|
|
# NOTE: it is first used in the finetune scenario.
|
|
if node_id is None:
|
|
selection = self.get_current_selection()
|
|
if self.is_selection_new_tree(selection):
|
|
return None
|
|
node_id = selection[0]
|
|
|
|
if node_id == -1:
|
|
if not self.hist:
|
|
return None
|
|
node_id = len(self.hist) - 1
|
|
|
|
ancestors = self.get_parents(node_id)
|
|
for i in reversed(ancestors):
|
|
if self.hist[i][1].decision:
|
|
return self.hist[i][0]
|
|
return None
|
|
|
|
|
|
class CheckpointSelector:
|
|
"""
|
|
In the trace, we may start from any check point (we'll represent it as a variable `from_checkpoint_idx`)
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_selection(self, trace: Trace) -> tuple[int, ...] | None:
|
|
"""
|
|
checkpoint_idx represents the place where we want to create a new node.
|
|
the return value should be the idx of target node (the parent of the new generating node).
|
|
- `(-1, )` represents starting from the latest trial in the trace - default value
|
|
|
|
- NOTE: we don't encourage to use this option; It is confusing when we have multiple traces.
|
|
|
|
- `(idx, )` represents starting from the `idx`-th trial in the trace.
|
|
- `None` represents starting from scratch (start a new trace)
|
|
|
|
|
|
- More advanced selection strategies in `select.py`
|
|
"""
|
|
|
|
|
|
class SOTAexpSelector:
|
|
"""
|
|
Select the SOTA experiment from the trace to submit
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_sota_exp_to_submit(self, trace: Trace) -> Experiment | None:
|
|
"""
|
|
Select the SOTA experiment from the trace to submit
|
|
"""
|
|
|
|
|
|
class ExpPlanner(ABC, Generic[ASpecificPlan]):
|
|
"""
|
|
An abstract class for planning the experiment.
|
|
The planner should generate a plan for the experiment based on the trace.
|
|
"""
|
|
|
|
def __init__(self, scen: Scenario) -> None:
|
|
self.scen = scen
|
|
|
|
@abstractmethod
|
|
def plan(self, trace: Trace) -> ASpecificPlan:
|
|
"""
|
|
Generate a plan for the experiment based on the trace.
|
|
The plan should be a dictionary that contains the plan to each stage.
|
|
"""
|
|
|
|
|
|
class ExpGen(ABC):
|
|
|
|
def __init__(self, scen: Scenario) -> None:
|
|
self.scen = scen
|
|
|
|
@abstractmethod
|
|
def gen(self, trace: Trace) -> Experiment:
|
|
"""
|
|
Generate the experiment based on the trace.
|
|
Planning is part of gen, but since we may support multi-stage planning,
|
|
we need to pass plan as optional argument.
|
|
|
|
`ExpGen().gen()` play a role like
|
|
|
|
.. code-block:: python
|
|
|
|
# ExpGen().gen() ==
|
|
Hypothesis2Experiment().convert(
|
|
HypothesisGen().gen(trace)
|
|
)
|
|
"""
|
|
|
|
async def async_gen(self, trace: Trace, loop: LoopBase) -> Experiment:
|
|
"""
|
|
generate the experiment and decide whether to stop yield generation and give up control to other routines.
|
|
"""
|
|
# we give a default implementation here.
|
|
# The proposal is set to try best to generate the experiment in max-parallel level.
|
|
while True:
|
|
if loop.get_unfinished_loop_cnt(loop.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
|
|
return self.gen(trace)
|
|
await asyncio.sleep(1)
|
|
|
|
def reset(self) -> None:
|
|
"""
|
|
Reset the proposal to the initial state.
|
|
Sometimes the main loop may want to reset the whole process to the initial state.
|
|
Default implementation does nothing; override in subclasses if needed.
|
|
"""
|
|
return
|
|
|
|
|
|
class HypothesisGen(ABC):
|
|
|
|
def __init__(self, scen: Scenario) -> None:
|
|
self.scen = scen
|
|
|
|
@abstractmethod
|
|
def gen(
|
|
self,
|
|
trace: Trace,
|
|
plan: ExperimentPlan | None = None,
|
|
) -> Hypothesis:
|
|
# def gen(self, scenario_desc: str, ) -> Hypothesis:
|
|
"""
|
|
Motivation of the variable `scenario_desc`:
|
|
- Mocking a data-scientist is observing the scenario.
|
|
|
|
scenario_desc may include:
|
|
- data observation:
|
|
- Original or derivative
|
|
- Task information:
|
|
"""
|
|
|
|
|
|
class Hypothesis2Experiment(ABC, Generic[ASpecificExp]):
|
|
"""
|
|
[Abstract description => concrete description] => Code implementation Card
|
|
"""
|
|
|
|
@abstractmethod
|
|
def convert(self, hypothesis: Hypothesis, trace: Trace) -> ASpecificExp:
|
|
"""Connect the idea proposal to implementation"""
|
|
...
|
|
|
|
|
|
# Boolean, Reason, Confidence, etc.
|
|
|
|
|
|
class Experiment2Feedback(ABC):
|
|
""" "Generated feedbacks on the hypothesis from **Executed** Implementations of different tasks
|
|
& their comparisons with previous performances"""
|
|
|
|
def __init__(self, scen: Scenario) -> None:
|
|
self.scen = scen
|
|
|
|
@abstractmethod
|
|
def generate_feedback(
|
|
self,
|
|
exp: Experiment,
|
|
trace: Trace,
|
|
exception: Exception | None = None,
|
|
) -> ExperimentFeedback:
|
|
"""
|
|
The `exp` should be executed and the results should be included, as well as the comparison
|
|
between previous results (done by LLM).
|
|
For example: `mlflow` of Qlib will be included.
|
|
"""
|
|
error_message = "generate_feedback method is not implemented."
|
|
raise NotImplementedError(error_message)
|