Files
NexQuant/rdagent/scenarios/kaggle/developer/coder.py
T
Xu Yang f7c1c4fd74 feat: implement isolated model feature selection loop (#370)
* rename meta_tpl

* use a isolated coder to deal with model feature selection and refine the structure

* fix CI

* fix: fix some errors in scenario.py, proposal.py and runner.py and several complex competition scenarios(#365)

* fix several bugs in proposal and runner

* fix a bug in feedback-prize-english-language-learning

* fix some bugs and templates

* fix the bug in optiver and nlp problem

* delete unnecessary codes

* remove unnecessary codes

* complete forest and s4e8

* push

* feedback & s4e8 &  forest

* optiver finished

* s3e11 & s3e26

* s4e9 finished

* sf-crime finished

* the last one finished

---------

Co-authored-by: WinstonLiyt <104308117+WinstonLiyt@users.noreply.github.com>
Co-authored-by: WinstonLiyte <1957922024@qq.com>
2024-09-28 00:40:25 +08:00

73 lines
2.8 KiB
Python

import json
from pathlib import Path
from jinja2 import Environment, StrictUndefined
from rdagent.components.coder.factor_coder.CoSTEER import FactorCoSTEER
from rdagent.components.coder.model_coder.CoSTEER import ModelCoSTEER
from rdagent.core.developer import Developer
from rdagent.core.prompts import Prompts
from rdagent.oai.llm_utils import APIBackend
from rdagent.scenarios.kaggle.experiment.kaggle_experiment import (
KG_SELECT_MAPPING,
KGModelExperiment,
)
KGModelCoSTEER = ModelCoSTEER
KGFactorCoSTEER = FactorCoSTEER
prompt_dict = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml")
DEFAULT_SELECTION_CODE = """
import pandas as pd
def select(X: pd.DataFrame) -> pd.DataFrame:
\"""
Select relevant features. To be used in fit & predict function.
\"""
if X.columns.nlevels == 1:
return X
{% if feature_index_list is not none %}
X = X.loc[:, X.columns.levels[0][{{feature_index_list}}].tolist()]
{% endif %}
X.columns = ["_".join(str(col)).strip() for col in X.columns.values]
return X
"""
class KGModelFeatureSelectionCoder(Developer[KGModelExperiment]):
def develop(self, exp: KGModelExperiment) -> KGModelExperiment:
target_model_type = exp.sub_tasks[0].model_type
assert target_model_type in KG_SELECT_MAPPING
if len(exp.experiment_workspace.data_description) == 1:
code = (
Environment(undefined=StrictUndefined)
.from_string(DEFAULT_SELECTION_CODE)
.render(feature_index_list=None)
)
else:
system_prompt = (
Environment(undefined=StrictUndefined)
.from_string(prompt_dict["model_feature_selection"]["system"])
.render(scenario=self.scen.get_scenario_all_desc(), model_type=exp.sub_tasks[0].model_type)
)
user_prompt = (
Environment(undefined=StrictUndefined)
.from_string(prompt_dict["model_feature_selection"]["user"])
.render(feature_groups=[desc[0] for desc in exp.experiment_workspace.data_description])
)
chosen_index = json.loads(
APIBackend().build_messages_and_create_chat_completion(
user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True
)
).get("Selected Group Index", [i + 1 for i in range(len(exp.experiment_workspace.data_description))])
chosen_index_to_list_index = [i - 1 for i in chosen_index]
code = (
Environment(undefined=StrictUndefined)
.from_string(DEFAULT_SELECTION_CODE)
.render(feature_index_list=chosen_index_to_list_index)
)
exp.experiment_workspace.inject_code(**{KG_SELECT_MAPPING[target_model_type]: code})
return exp