mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-28 07:57:44 +00:00
5baed909e7
* refactor: Update type annotations and remove unused class in evolving modules * refactor: Simplify evolving agent and feedback handling in CoSTEER module * lint & CI * mypy * ruff for core * mypy * refactor: remove unnecessary comments and update feedback handling logic * refactor: Add prev_task_feedback parameter to evolving strategies * feat: Clear folder before extracting zip file in DockerEnv * fix: Correct retrieval of last experiment from history
111 lines
4.4 KiB
Python
111 lines
4.4 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
from jinja2 import Environment, StrictUndefined
|
|
|
|
from rdagent.components.coder.CoSTEER.config import CoSTEER_SETTINGS
|
|
from rdagent.components.coder.CoSTEER.evaluators import CoSTEERSingleFeedback
|
|
from rdagent.components.coder.CoSTEER.evolving_strategy import (
|
|
MultiProcessEvolvingStrategy,
|
|
)
|
|
from rdagent.components.coder.CoSTEER.knowledge_management import (
|
|
CoSTEERQueriedKnowledge,
|
|
CoSTEERQueriedKnowledgeV2,
|
|
)
|
|
from rdagent.components.coder.model_coder.model import (
|
|
ModelExperiment,
|
|
ModelFBWorkspace,
|
|
ModelTask,
|
|
)
|
|
from rdagent.core.experiment import FBWorkspace
|
|
from rdagent.core.prompts import Prompts
|
|
from rdagent.oai.llm_conf import LLM_SETTINGS
|
|
from rdagent.oai.llm_utils import APIBackend
|
|
|
|
coder_prompts = Prompts(file_path=Path(__file__).parent / "prompts.yaml")
|
|
|
|
|
|
class ModelMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
|
|
def implement_one_task(
|
|
self,
|
|
target_task: ModelTask,
|
|
queried_knowledge: CoSTEERQueriedKnowledge = None,
|
|
workspace: FBWorkspace | None = None,
|
|
prev_task_feedback: CoSTEERSingleFeedback | None = None,
|
|
) -> str:
|
|
model_information_str = target_task.get_task_information()
|
|
|
|
queried_similar_successful_knowledge = (
|
|
queried_knowledge.task_to_similar_task_successful_knowledge[model_information_str]
|
|
if queried_knowledge is not None
|
|
else []
|
|
)
|
|
queried_former_failed_knowledge = (
|
|
queried_knowledge.task_to_former_failed_traces[model_information_str]
|
|
if queried_knowledge is not None
|
|
else []
|
|
)
|
|
|
|
queried_former_failed_knowledge_to_render = (
|
|
queried_former_failed_knowledge[0]
|
|
if isinstance(queried_knowledge, CoSTEERQueriedKnowledgeV2)
|
|
else queried_former_failed_knowledge
|
|
)
|
|
|
|
system_prompt = (
|
|
Environment(undefined=StrictUndefined)
|
|
.from_string(
|
|
coder_prompts["evolving_strategy_model_coder"]["system"],
|
|
)
|
|
.render(
|
|
scenario=self.scen.get_scenario_all_desc(filtered_tag=target_task.model_type),
|
|
queried_former_failed_knowledge=queried_former_failed_knowledge_to_render,
|
|
current_code=target_task.base_code,
|
|
)
|
|
)
|
|
|
|
queried_similar_successful_knowledge_to_render = queried_similar_successful_knowledge
|
|
for _ in range(10): # max attempt to reduce the length of user_prompt
|
|
user_prompt = (
|
|
Environment(undefined=StrictUndefined)
|
|
.from_string(
|
|
coder_prompts["evolving_strategy_model_coder"]["user"],
|
|
)
|
|
.render(
|
|
model_information_str=model_information_str,
|
|
queried_similar_successful_knowledge=queried_similar_successful_knowledge_to_render,
|
|
queried_former_failed_knowledge=queried_former_failed_knowledge_to_render,
|
|
)
|
|
.strip("\n")
|
|
)
|
|
if (
|
|
APIBackend().build_messages_and_calculate_token(
|
|
user_prompt=user_prompt,
|
|
system_prompt=system_prompt,
|
|
)
|
|
< LLM_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(
|
|
APIBackend(use_chat_cache=CoSTEER_SETTINGS.coder_use_cache).build_messages_and_create_chat_completion(
|
|
user_prompt=user_prompt,
|
|
system_prompt=system_prompt,
|
|
json_mode=True,
|
|
),
|
|
)["code"]
|
|
return code
|
|
|
|
def assign_code_list_to_evo(self, code_list, evo):
|
|
for index in range(len(evo.sub_tasks)):
|
|
if code_list[index] is None:
|
|
continue
|
|
if evo.sub_workspace_list[index] is None:
|
|
evo.sub_workspace_list[index] = ModelFBWorkspace(target_task=evo.sub_tasks[index])
|
|
evo.sub_workspace_list[index].inject_files(**{"model.py": code_list[index]})
|
|
return evo
|