Feat: Add WebUI (please merge with squash) (#27)

* 💄 Add webUI for aideml * add gradio webui * update requirements * update gitignore * 💄 Update webUI and use logs to display results * 💄 Add graio theme and strealit app * 💄 Update streamlit webui * 💄 fix issues for serialization on config display * 💄 Fix results at bottom issue * 💄 Update the streamlit webui theme * 💄 Update the streamlit app theme * 🔥 Remove gradio and use streamlit * 🚨 add linter workflow and fix issues * Add github action for linter * Fix current linter issues * 👷 Add github templates * Add issue templates * Add pr template * 🚨 Add black formatter (#25) * 🚨 Add black formatter * Add black to linter * Fix the formatting * 🚨 Update black linter to see the suggestion change * 💄 Add webUI for aideml * add gradio webui * update requirements * update gitignore * 💄 Update streamlit app * 🚨 fix linter issue * 👷 Fix styling issues for tree viz * 👷 Finalize the theme * ♻️ Refactor webui code * 📝 Update README * 🚨 fix linter * 🐛 fix conflict * ➕ Add dependecy
WecoAI · Nov 26, 2024 · 5c7fa16 · 5c7fa16
1 parent 5ddb3df
commit 5c7fa16
Show file tree

Hide file tree

Showing 7 changed files with 816 additions and 44 deletions.
diff --git a/.gitignore b/.gitignore
@@ -167,4 +167,4 @@ logs
 .trunk
 
 .gradio/
-.ruff_cache/
+.ruff_cache/
diff --git a/.streamlit/config.toml b/.streamlit/config.toml
@@ -0,0 +1,17 @@
+[theme]
+# Primary colors
+primaryColor="#0D0F18"  # --wecopink: 343 98% 63%
+backgroundColor="#F0EFE9" # --background: 49 10% 94%
+secondaryBackgroundColor="#FFFFFF" # --card: 60 33.3% 98%
+textColor="#0A0A0A" # --primary: 0 0% 17%
+
+# Font
+font="sans serif"
+
+[ui]
+hideTopBar = true
+
+[client]
+toolbarMode = "minimal"
+showErrorDetails = true
+showSidebarNavigation = false
diff --git a/README.md b/README.md
@@ -11,14 +11,15 @@ AIDE is an LLM agent that generates solutions for machine learning tasks just fr
 AIDE is the state-of-the-art agent on OpenAI's [MLE-bench](https://arxiv.org/pdf/2410.07095), a benchmark composed of 75 Kaggle machine learning tasks, where we achieved four times more medals compared to the runner-up agent architecture.
 
 In our own benchmark composed of over 60 Kaggle data science competitions, AIDE demonstrated impressive performance, surpassing 50% of Kaggle participants on average (see our [technical report](https://www.weco.ai/blog/technical-report) for details).
+
 More specifically, AIDE has the following features:
 
 1. **Instruct with Natural Language**: Describe your problem or additional requirements and expert insights, all in natural language.
 2. **Deliver Solution in Source Code**: AIDE will generate Python scripts for the **tested** machine learning pipeline. Enjoy full transparency, reproducibility, and the freedom to further improve the source code!
 3. **Iterative Optimization**: AIDE iteratively runs, debugs, evaluates, and improves the ML code, all by itself.
 4. **Visualization**: We also provide tools to visualize the solution tree produced by AIDE for a better understanding of its experimentation process. This gives you insights not only about what works but also what doesn't.
 
-# How to use AIDE?
+# How to Use AIDE?
 
 ## Setup
 
@@ -38,7 +39,7 @@ export OPENAI_API_KEY=<your API key>
 export ANTHROPIC_API_KEY=<your API key>
 ```
 
-## Running AIDE via the command line
+## Running AIDE via the Command Line
 
 To run AIDE:
 
@@ -54,9 +55,9 @@ aide data_dir="example_tasks/house_prices" goal="Predict the sales price for eac
 
 Options:
 
-- `data_dir` (required): a directory containing all the data relevant for your task (`.csv` files, images, etc.).
-- `goal`: describe what you want the models to predict in your task, for example, "Build a timeseries forcasting model for bitcoin close price" or "Predict sales price for houses".
-- `eval`: the evaluation metric used to evaluate the ML models for the task (e.g., accuracy, F1, Root-Mean-Squared-Error, etc.)
+- `data_dir` (required): A directory containing all the data relevant for your task (`.csv` files, images, etc.).
+- `goal`: Describe what you want the models to predict in your task, for example, "Build a time series forecasting model for bitcoin close price" or "Predict sales price for houses".
+- `eval`: The evaluation metric used to evaluate the ML models for the task (e.g., accuracy, F1, Root-Mean-Squared-Error, etc.).
 
 Alternatively, you can provide the entire task description as a `desc_str` string, or write it in a plaintext file and pass its path as `desc_file` ([example file](aide/example_tasks/house_prices.md)).
 
@@ -66,19 +67,19 @@ aide data_dir="my_data_dir" desc_file="my_task_description.txt"
 
 The result of the run will be stored in the `logs` directory.
 
-- `logs/<experiment-id>/best_solution.py`: Python code of _best solution_ according to the validation metric
-- `logs/<experiment-id>/journal.json`: a JSON file containing the metadata of the experiment runs, including all the code generated in intermediate steps, plan, evaluation results, etc.
-- `logs/<experiment-id>/tree_plot.html`: you can open it in your browser. It contains visualization of solution tree, which details the experimentation process of finding and optimizing ML code. You can explore and interact with the tree visualization to view what plan and code AIDE comes up with in each step.
+- `logs/<experiment-id>/best_solution.py`: Python code of the _best solution_ according to the validation metric.
+- `logs/<experiment-id>/journal.json`: A JSON file containing the metadata of the experiment runs, including all the code generated in intermediate steps, plan, evaluation results, etc.
+- `logs/<experiment-id>/tree_plot.html`: You can open it in your browser. It contains a visualization of the solution tree, which details the experimentation process of finding and optimizing ML code. You can explore and interact with the tree visualization to view what plan and code AIDE comes up with in each step.
 
 The `workspaces` directory will contain all the files and data that the agent generated.
 
 ### Advanced Usage
 
-To further customize the behaviour of AIDE, some useful options might be:
+To further customize the behavior of AIDE, some useful options might be:
 
-- `agent.code.model=...` to configure which model the agent should use for coding (default is `gpt-4-turbo`)
-- `agent.steps=...` to configure how many improvement iterations the agent should run (default is 20)
-- `agent.search.num_drafts=...` to configure the number of initial drafts the agent should generate (default is 5)
+- `agent.code.model=...` to configure which model the agent should use for coding (default is `gpt-4-turbo`).
+- `agent.steps=...` to configure how many improvement iterations the agent should run (default is 20).
+- `agent.search.num_drafts=...` to configure the number of initial drafts the agent should generate (default is 5).
 
 You can check the [`config.yaml`](aide/utils/config.yaml) file for more options.
 
@@ -88,23 +89,73 @@ AIDE supports using local LLMs through OpenAI-compatible APIs. Here's how to set
 
 1. Set up a local LLM server with an OpenAI-compatible API endpoint. You can use:
    - [Ollama](https://github.com/ollama/ollama)
-   - or similar solutions
+   - or similar solutions.
 
 2. Configure your environment to use the local endpoint:
+
+   ```bash
+   export OPENAI_BASE_URL="http://localhost:11434/v1"  # For Ollama
+   export OPENAI_API_KEY="local-llm"  # Can be any string if your local server doesn't require authentication
+   ```
+
+3. Update the model configuration in your AIDE command or config. For example, with Ollama:
+
+   ```bash
+   # Example with house prices dataset
+   aide agent.code.model="qwen2.5" agent.feedback.model="qwen2.5" report.model="qwen2.5" \
+       data_dir="example_tasks/house_prices" \
+       goal="Predict the sales price for each house" \
+       eval="Use the RMSE metric between the logarithm of the predicted and observed values."
+   ```
+
+## Running AIDE via the Web UI
+
+We have developed a user-friendly Web UI using Streamlit to make it even easier to interact with AIDE.
+
+### Prerequisites
+
+Ensure you have installed the development version of AIDE and its dependencies as described in the [Development](#development) section.
+
+### Running the Web UI
+
+Navigate to the `aide/webui` directory and run the Streamlit application:
+
 ```bash
-export OPENAI_BASE_URL="http://localhost:11434/v1"  # For Ollama
-export OPENAI_API_KEY="local-llm"  # Can be any string if your local server doesn't require authentication
+cd aide/webui
+streamlit run app.py
 ```
 
-3. Update the model configuration in your AIDE command or config. For example, with Ollama:
+Alternatively, you can run it from the root directory:
+
 ```bash
-# Example with house prices dataset
-aide agent.code.model="qwen2.5" agent.feedback.model="qwen2.5" report.model="qwen2.5" \
-    data_dir="example_tasks/house_prices" \
-    goal="Predict the sales price for each house" \
-    eval="Use the RMSE metric between the logarithm of the predicted and observed values."
+streamlit run aide/webui/app.py
 ```
 
+### Using the Web UI
+
+1. **API Key Configuration**: In the sidebar, input your OpenAI API key or Anthropic API key and click "Save API Keys".
+
+2. **Input Data**:
+   - You can either **upload your dataset files** (`.csv`, `.txt`, `.json`, `.md`) using the "Upload Data Files" feature.
+   - Or click on "Load Example Experiment" to use the example house prices dataset.
+
+3. **Define Goal and Evaluation Criteria**:
+   - In the "Goal" text area, describe what you want the model to achieve (e.g., "Predict the sales price for each house").
+   - In the "Evaluation Criteria" text area, specify the evaluation metric (e.g., "Use the RMSE metric between the logarithm of the predicted and observed values.").
+
+4. **Configure Steps**:
+   - Use the slider to set the number of steps (iterations) for the experiment.
+
+5. **Run the Experiment**:
+   - Click on "Run AIDE" to start the experiment.
+   - Progress and status updates will be displayed in the "Results" section.
+
+6. **View Results**:
+   - **Tree Visualization**: Explore the solution tree to understand how AIDE experimented and optimized the models.
+   - **Best Solution**: View the Python code of the best solution found.
+   - **Config**: Review the configuration used for the experiment.
+   - **Journal**: Examine the detailed journal entries for each step.
+
 ## Using AIDE in Python
 
 Using AIDE within your Python script/project is easy. Follow the setup steps above, and then create an AIDE experiment like below and start running:
@@ -113,7 +164,7 @@ Using AIDE within your Python script/project is easy. Follow the setup steps abo
 import aide
 exp = aide.Experiment(
     data_dir="example_tasks/bitcoin_price",  # replace this with your own directory
-    goal="Build a timeseries forcasting model for bitcoin close price.",  # replace with your own goal description
+    goal="Build a time series forecasting model for bitcoin close price.",  # replace with your own goal description
     eval="RMSLE"  # replace with your own evaluation metric
 )
 
@@ -125,41 +176,53 @@ print(f"Best solution code: {best_solution.code}")
 
 ## Development
 
-To install AIDE for development, clone this repository and install it locally.
+To install AIDE for development, clone this repository and install it locally:
 
 ```bash
 git clone https://github.com/WecoAI/aideml.git
 cd aideml
 pip install -e .
 ```
 
-## Using AIDE with Docker
+### Running the Web UI in Development Mode
 
-You can also run AIDE using Docker:
+Ensure that you have all the required development dependencies installed. Then, you can run the Web UI as follows:
 
-1. Build the Docker image:
 ```bash
-docker build -t aide .
+cd aide/webui
+streamlit run app.py
 ```
 
-2. Run AIDE with Docker (example with house prices task):
-```bash
-# Set custom workspace and logs location (optional)
-export WORKSPACE_BASE=$(pwd)/workspaces
-export LOGS_DIR=$(pwd)/logs
-
-docker run -it --rm \
-          -v "${LOGS_DIR:-$(pwd)/logs}:/app/logs" \
-          -v "${WORKSPACE_BASE:-$(pwd)/workspaces}:/app/workspaces" \
-          -v "$(pwd)/aide/example_tasks:/app/data" \
-          -e OPENAI_API_KEY="your-actual-api-key" \
-          aide \
-          data_dir=/app/data/house_prices \
-          goal="Predict the sales price for each house" \
-          eval="Use the RMSE metric between the logarithm of the predicted and observed values."
-```
+## Using AIDE with Docker
+
+You can also run AIDE using Docker:
+
+1. **Build the Docker Image**:
+
+   ```bash
+   docker build -t aide .
+   ```
+
+2. **Run AIDE with Docker** (example with house prices task):
+
+   ```bash
+   # Set custom workspace and logs location (optional)
+   export WORKSPACE_BASE=$(pwd)/workspaces
+   export LOGS_DIR=$(pwd)/logs
+
+   docker run -it --rm \
+             -v "${LOGS_DIR:-$(pwd)/logs}:/app/logs" \
+             -v "${WORKSPACE_BASE:-$(pwd)/workspaces}:/app/workspaces" \
+             -v "$(pwd)/aide/example_tasks:/app/data" \
+             -e OPENAI_API_KEY="your-actual-api-key" \
+             aide \
+             data_dir=/app/data/house_prices \
+             goal="Predict the sales price for each house" \
+             eval="Use the RMSE metric between the logarithm of the predicted and observed values."
+   ```
 
 You can customize the location of workspaces and logs by setting environment variables before running the container:
+
 - `WORKSPACE_BASE`: Sets the base directory for AIDE workspaces (default: `$(pwd)/workspaces`)
 - `LOGS_DIR`: Sets the directory for AIDE logs (default: `$(pwd)/logs`)
 

diff --git a/aide/webui/__init__.py b/aide/webui/__init__.py