Skip to content

[MRG] Update the advisor to support various modes #239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
huangyz0918 merged 1 commit into from
Oct 12, 2024
Merged

[MRG] Update the advisor to support various modes #239

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 43 additions & 12 deletions mle/agents/advisor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ def process_report(requirement: str, suggestions: dict):

class AdviseAgent:

def __init__(self, model, console=None):
def __init__(self, model, console=None, mode='normal'):
"""
AdviseAgent: the agent to suggest which machine learning task/model/dataset to use based on the user's
requirements. The return of the agent is an instruction to the user to modify the code based on the logs and
Expand All @@ -51,24 +51,24 @@ def __init__(self, model, console=None):

Your capabilities include:

1. Read and understand the dataset information and user's requirements, the requirements may include the task,
- Read and understand the dataset information and user's requirements, the requirements may include the task,
the model (or method), and the evaluation metrics, etc. You should always follow the user's requirements.
2. You should briefly analyze the user's dataset, and give a summary of the dataset, the dataset input can be
- You should briefly analyze the user's dataset, and give a summary of the dataset, the dataset input can be
a public dataset name or a path to a local CSV file. You can use the function `preview_csv_data` to preview
the CSV files or not if the dataset is a public dataset.
3. And then you should always use the function `search_arxiv` or `search_papers_with_code` to search the
- And then you should always use the function `search_arxiv` or `search_papers_with_code` to search the
state-of-the-art machine learning tasks/models/algorithms that can be used to solve the user's requirements,
and stay up-to-date with the latest.
4. If the user does not provide the details (task/model/algorithm/dataset/metric), you should always suggest.
5. You should provide the paper reference links of the task/model/algorithm/metric you suggest. You use the
- If the user does not provide the details (task/model/algorithm/dataset/metric), you should always suggest.
- You should provide the paper reference links of the task/model/algorithm/metric you suggest. You use the
search results from the function `search_arxiv` or `search_papers_with_code` by generated search keywords.
6. The suggestion should be as detailed as possible, include the SOTA methods for data processing, feature
- The suggestion should be as detailed as possible, include the SOTA methods for data processing, feature
extraction, model selection, training/sering methods and evaluation metrics. And the reasons why you suggest.
7. You should help user to decide which framework/tools to use for the project, such as PyTorch, TensorFlow,
- You should help user to decide which framework/tools to use for the project, such as PyTorch, TensorFlow,
MLFlow, W&B, etc.
"""
self.search_prompt = """
7. You should also use function `web_search` to search for articles, papers, or tutorials related to the
- You should also use function `web_search` to search for articles, papers, or tutorials related to the
task/model/algorithm/metric to help you decide which one to use.
"""
self.json_mode_prompt = """
Expand All @@ -89,6 +89,37 @@ def __init__(self, model, console=None):
}

"""

# precise means the user has provided specific requirements with metrics, models, etc.
if mode == 'precise':
self.sys_prompt = f"""
You are an Machine learning expert tasked with advising on the best ML task/model/algorithm to use.

Your capabilities include:

- Based on the requirements, give suggestions on how to maximize the performance of the task metrics.
- Use the function `search_arxiv` or `search_papers_with_code` to search the state-of-the-art machine learning
tasks/models/algorithms that can be used to solve the user's requirements, and stay up-to-date with the latest.
- You should help user to decide which framework/tools to use to implement the project, such as PyTorch, TensorFlow, etc.
- You should apply some tricks to improve the performance of the model, detail the implementation steps of each trick.

"""

self.json_mode_prompt = """

JSON Output Format:

{
"task":"xxxxx",
"model_or_algorithm":"xxxx",
"frameworks": ["xxxx", "xxxx"],
"training_method": "xxxx",
"suggestion": "Based on the user requirement, we suggest you to... and why...",
"tricks": ["xxx", ...]
}

"""

self.functions = [
schema_web_search,
schema_search_arxiv,
Expand All @@ -102,12 +133,12 @@ def __init__(self, model, console=None):
self.sys_prompt += self.json_mode_prompt
self.chat_history.append({"role": 'system', "content": self.sys_prompt})

def suggest(self, requirement, process=True):
def suggest(self, requirement, return_raw=False):
"""
Handle the query from the model query response.
Args:
requirement: the user requirement.
process: whether to process the report.
return_raw: whether to process the report.
"""
with self.console.status("MLE Advisor is thinking the best strategy to help you..."):
self.chat_history.append({"role": "user", "content": requirement})
Expand All @@ -124,7 +155,7 @@ def suggest(self, requirement, process=True):
except json.JSONDecodeError as e:
suggestions = clean_json_string(text)

if not process:
if return_raw:
return suggestions

return process_report(requirement, suggestions)
Expand Down
7 changes: 3 additions & 4 deletions mle/agents/summarizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,8 +137,8 @@ def summarize(self):
def kaggle_request_summarize(
self,
kaggle_overview: str,
exp_submission: str = None,
submission_file: str = "./submission.csv"):
exp_submission: str = None
):
"""
Summarize the kaggle requests.
:params: kaggle_overview: the overview json of kaggle competition
Expand All @@ -159,8 +159,7 @@ def kaggle_request_summarize(

if exp_submission:
kaggle_overview += f"""
\n\n EXAMPLE SUBMISSION FILE: {exp_submission}\n
SUBMISSION FILE LOCATION: {submission_file}
\n\n EXAMPLE SUBMISSION FILE: {exp_submission}
"""

chat_history = [
Expand Down
23 changes: 14 additions & 9 deletions mle/workflow/kaggle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ def auto_kaggle(
model = load_model(work_dir, model)

# initialize the agents
advisor = AdviseAgent(model, console)
advisor = AdviseAgent(model, console, mode="precise")
summarizer = SummaryAgent(model, console=console)
coder = CodeAgent(model, work_dir, console=console, single_file=True)
debugger = DebugAgent(model, console, analyze_only=True)
Expand All @@ -45,20 +45,25 @@ def auto_kaggle(
description = read_markdown(description)

with console.status("MLE Agent is processing the kaggle competition overview..."):
requirements = summarizer.kaggle_request_summarize(description, sub_examples, submission)
requirements = summarizer.kaggle_request_summarize(description, sub_examples)
requirements += f"\n\nLOCAL DATASET PATH:\n"
for dataset in datasets:
requirements += f" - {dataset}\n"

suggestions = advisor.suggest(requirements, process=False)
requirements += f"\nSUBMISSION FILE PATH: {submission}\n"

suggestions = advisor.suggest(requirements, return_raw=True)
requirements += f"""
\nIMPLEMENTATION PLAN:
\nIMPLEMENTATION SUGGESTIONS:

- Suggestion Summary: {suggestions.get('suggestion')}\n
- ML Task: {suggestions.get('task')}\n
- Model or Algorithm: {suggestions.get('model_or_algorithm')}
- Training Strategy: {suggestions.get('training_method')}
"""
- Suggestion summary: {suggestions.get('suggestion')}
- ML task: {suggestions.get('task')}
- Model or algorithm: {suggestions.get('model_or_algorithm')}
- Training strategy: {suggestions.get('training_method')}
- Tricks to increase performance:
"""
for trick in suggestions.get('tricks'):
requirements += f"\n - {trick}"

coder.read_requirement(requirements)
if competition_id is None:
Expand Down