diff --git a/.gitignore b/.gitignore index 6a1d2b4f..bfc4f067 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,8 @@ cdk.out/ .langchain.db output_*.csv runs/ +*.index +*.metadata # MCP specific .flask.pid diff --git a/README.md b/README.md index fa057824..582519d8 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,18 @@ AutoGluon Assistant provides multiple interfaces: ![Demo](https://github.com/autogluon/autogluon-assistant/blob/main/docs/assets/mcp_demo.gif) +### Integration with Upgini + +To enable Upgini integration for enriching your dataset with additional relevant features and selecting the most relevant features from the input dataset, set the `UPGINI_API_KEY` environment variable before invoking MLZero. You can obtain your API key in your [Upgini profile](https://profile.upgini.com). + +```bash +export UPGINI_API_KEY="" +# then run MLZero as usual +mlzero -i +``` + + + ## Citation If you use Autogluon Assistant (MLZero) in your research, please cite our paper: diff --git a/src/autogluon/assistant/cli/app.py b/src/autogluon/assistant/cli/app.py index 785a0d4c..1a851663 100644 --- a/src/autogluon/assistant/cli/app.py +++ b/src/autogluon/assistant/cli/app.py @@ -95,7 +95,7 @@ def main( # 3) Invoke the core run_agent function # Override config path if provider is specified and config path is default provider_config_path = config_path - if llm_provider in ["bedrock", "openai", "anthropic", "sagemaker"] and config_path == DEFAULT_CONFIG_PATH: + if llm_provider in ["bedrock", "openai", "anthropic", "sagemaker", "azure"] and config_path == DEFAULT_CONFIG_PATH: provider_config_path = Path(DEFAULT_CONFIG_PATH).parent / f"{llm_provider}.yaml" if not provider_config_path.exists(): provider_config_path = DEFAULT_CONFIG_PATH diff --git a/src/autogluon/assistant/configs/azure.yaml b/src/autogluon/assistant/configs/azure.yaml new file mode 100644 index 00000000..c25c1f28 --- /dev/null +++ b/src/autogluon/assistant/configs/azure.yaml @@ -0,0 +1,83 @@ +# SageMaker Configuration + +per_execution_timeout: 86400 + +# Data Perception +max_file_group_size_to_show: 5 +num_example_files_to_show: 1 + +max_chars_per_file: 768 +num_tutorial_retrievals: 30 +max_num_tutorials: 5 +max_user_input_length: 2048 +max_error_message_length: 2048 +max_tutorial_length: 32768 +create_venv: false +condense_tutorials: True +use_tutorial_summary: True +continuous_improvement: False +optimize_system_resources: False +cleanup_unused_env: True +enable_meta_prompting: False + +llm: &default_llm + provider: azure + model: o3-mini + max_tokens: 65535 + proxy_url: null + top_p: 0.9 + verbose: True + multi_turn: False + template: null + add_coding_format_instruction: false + apply_meta_prompting: False + +# Ensure all agent types inherit the SageMaker LLM config +python_coder: + <<: *default_llm # Merge llm_config + multi_turn: True + apply_meta_prompting: True + +bash_coder: + <<: *default_llm # Merge llm_config + multi_turn: True + +executer: + <<: *default_llm # Merge llm_config + max_stdout_length: 8192 + max_stderr_length: 2048 + +meta_prompting: + <<: *default_llm # Merge llm_config + multi_turn: False + +reader: + <<: *default_llm # Merge llm_config + details: False + +error_analyzer: + <<: *default_llm # Merge llm_config + +retriever: + <<: *default_llm # Merge llm_config + +reranker: + <<: *default_llm # Merge llm_config + temperature: 0. + top_p: 1. + +description_file_retriever: + <<: *default_llm # Merge llm_config + temperature: 0. + top_p: 1. + +task_descriptor: + <<: *default_llm # Merge llm_config + max_description_files_length_to_show: 1024 + max_description_files_length_for_summarization: 16384 + apply_meta_prompting: True + +tool_selector: + <<: *default_llm # Merge llm_config + temperature: 0. + top_p: 1. diff --git a/src/autogluon/assistant/llm/azure_openai_chat.py b/src/autogluon/assistant/llm/azure_openai_chat.py index 30ba099f..2c19241d 100644 --- a/src/autogluon/assistant/llm/azure_openai_chat.py +++ b/src/autogluon/assistant/llm/azure_openai_chat.py @@ -8,6 +8,7 @@ from .base_chat import BaseAssistantChat logger = logging.getLogger(__name__) +NO_TEMPERATURE_MODELS = {"o1", "o1-mini", "o3", "o3-mini"} class AssistantAzureChatOpenAI(AzureChatOpenAI, BaseAssistantChat): @@ -53,7 +54,7 @@ def create_azure_openai_chat(config, session_name: str) -> AssistantAzureChatOpe "max_tokens": config.max_tokens, } - if hasattr(config, "temperature"): + if hasattr(config, "temperature") and model not in NO_TEMPERATURE_MODELS: kwargs["temperature"] = config.temperature if hasattr(config, "verbose"): diff --git a/src/autogluon/assistant/prompts/bash_coder_prompt.py b/src/autogluon/assistant/prompts/bash_coder_prompt.py index 5bb07907..e7e45423 100644 --- a/src/autogluon/assistant/prompts/bash_coder_prompt.py +++ b/src/autogluon/assistant/prompts/bash_coder_prompt.py @@ -97,12 +97,13 @@ def get_env_prompt(self): common_env_file = self.manager.common_env_file selected_tool_env_file = self.manager.selected_tool_env_file + install_cmd = f"uv pip install --system -r {selected_tool_env_file} -r {common_env_file}" env_prompt = f""" Create and configure a conda environment in "{ENV_FOLDER_NAME}" folder under {iteration_folder}: - Python version: 3.11 - - Activate the environment + - Activate the environment (conda activate "{ENV_FOLDER_NAME}") - pip install uv - - Install required packages from {common_env_file} and {selected_tool_env_file} using uv pip install -r {selected_tool_env_file} -r {common_env_file}""" + - Install required packages from {common_env_file} and {selected_tool_env_file} using {install_cmd}""" if not create_venv: env_prompt += f"\n - Do not install or update any package unless there is an error due to the missing package.\n - Do NOT upgrade {selected_tool} which is already installed." diff --git a/src/autogluon/assistant/prompts/python_coder_prompt.py b/src/autogluon/assistant/prompts/python_coder_prompt.py index e385e420..4dcf573a 100644 --- a/src/autogluon/assistant/prompts/python_coder_prompt.py +++ b/src/autogluon/assistant/prompts/python_coder_prompt.py @@ -34,25 +34,65 @@ def default_template(self) -> str: ONLY save files to the working directory: {per_iteration_output_folder}. 1. Data preprocessing: - - Remove training data samples without valid labels (drop NA values from training dataset ONLY, NOT from test dataset) unless explicitly instructed otherwise. - - Remove the unneccesary index column (if applicable) - -2. Model training: + - Remove training data samples without valid labels (drop NA values from the training dataset ONLY, NOT from test dataset) unless explicitly instructed otherwise. + - Remove any unneccesary index column (if present). + - If the environment variable MLZERO_SAMPLE_SIZE is set and the training dataset contains more than MLZERO_SAMPLE_SIZE rows, draw a random sample: + ```python + import os + + if "MLZERO_SAMPLE_SIZE" in os.environ: + sample_size = int(os.getenv("MLZERO_SAMPLE_SIZE")) + if len(train_data) > sample_size: + train_data = train_data.sample(n=sample_size, random_state=42) + else: + print(f"Training dataset has {len(train_data)} rows, which is less than MLZERO_SAMPLE_SIZE {sample_size}. Skipping sampling.") + ``` + +2. Data enrichment: + - Perform this step only if UPGINI_API_KEY is set in the environment and the data is tabular. Otherwise, skip this step. + - Use Upgini to enrich the data and select the most relevant features from internal and external sources. + - Example of usage: + ```python + from upgini import FeaturesEnricher + enricher = FeaturesEnricher() + train_data_features = enricher.fit_transform(train_data.drop(columns=label), train_data[label], eval_set=(test_data.drop(columns=label, errors='ignore'), ), calculate_metrics=False) + train_data = pd.concat([train_data_features, train_data[[label]]], axis=1) + test_data_for_prediction = enricher.transform(test_data.drop(columns=label), errors='ignore') + # If the dataset includes a stable identifier column (e.g., id, row_id, key, uid, or any column with 100% unique values), after enrichment, include it in the test data for prediction + test_data_for_prediction = pd.concat([test_data[[id_column_name]], test_data_for_prediction], axis=1) + ``` + - If there are multiple target columns, enrich the data separately for each target: + ```python + train_data_features = dict() + test_data_for_prediction = dict() + for label in labels: + train_data_features[label] = enricher.fit_transform(train_data.drop(columns=labels), train_data[label], eval_set=(test_data.drop(columns=labels, errors='ignore'), ), calculate_metrics=False) + test_data_features = enricher.transform(test_data.drop(columns=labels, errors='ignore'), errors='ignore') + test_data_for_prediction[label] = pd.concat([test_data[[id_column_name]], test_data_features], axis=1) + + predictions = dict() + for label in labels: + predictor = ... + predictor.fit(train_data=train_data_features[label], ...) + predictions[label] = predictor.predict(test_data_for_prediction[label]) + ``` + +3. Model training: - Use {selected_tool} with appropriate parameters for the task - If a model is trained, save it in a folder with random timestamp within {per_iteration_output_folder} -3. Prediction: +4. Prediction: - Make predictions on the test data. Always preserve and use the ORIGINAL INDICES from the test data to maintain exact row correspondence - DO NOT generate new indices or rely on assumed ordering. - Save the predicted results to {per_iteration_output_folder}, result file name should be "results", the format and extension should be same as the test data file - Output column names must exactly match those in the training or sample submission files without adding "predicted_" prefixes or creating any new columns. - At the end, implement validation checks that assert the prediction file maintains exact test data indices, verify correct column names match requirements, and confirm proper output format. -4. Documentation: +5. Documentation: - Add a brief docstring at the beginning of the script explaining its purpose - Include additional installation steps with comments at the beginning of the script - Include comments explaining any complex operations or design decisions -5. Others: +6. Others: - To avoid DDP errors, wrap the code in: if __name__ == "__main__": - Ensure errors are propagated up and not silently caught - do not use try/except blocks unless you explicitly re-raise the exception. diff --git a/src/autogluon/assistant/tools_registry/_common/requirements.txt b/src/autogluon/assistant/tools_registry/_common/requirements.txt index 0ce744e7..86c6a312 100644 --- a/src/autogluon/assistant/tools_registry/_common/requirements.txt +++ b/src/autogluon/assistant/tools_registry/_common/requirements.txt @@ -12,3 +12,4 @@ python-dateutil datasets>=3.6 six pyarrow +upgini