Files
NexQuant/rdagent/components/coder/factor_coder/CoSTEER/scheduler.py
T
Xu Yang 768229427d feat: use unified pickle cacher & move llm config into a isolated config (#424)
* simplify RDAgent conf

* add unified cacher(untested)

* fix small bugs

* fix a bug

* fix a small bug in runner

* use hash_key = None to skip cache

* fix CI

* in factor execution, ignore cache when raise exception

* add file locker to avoid mp calling

* fix CI

* use function __module__ name as folder in cache
2024-10-14 17:34:09 +08:00

89 lines
2.7 KiB
Python

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.core.prompts import Prompts
from rdagent.core.scenario import Scenario
from rdagent.log import rdagent_logger as logger
from rdagent.oai.llm_conf import LLM_SETTINGS
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,
)
logger.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_task_information()
if target_factor_task_information in former_trace:
tasks.append((i, evo.sub_tasks[i], former_trace[target_factor_task_information][0]))
system_prompt = (
Environment(undefined=StrictUndefined)
.from_string(
scheduler_prompts["select_implementable_factor_system"],
)
.render(
scenario=scen.get_scenario_all_desc(),
)
)
for _ in range(10): # max attempt to reduce the length of user_prompt
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,
)
< LLM_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