weco 0.2.22__tar.gz → 0.2.23__tar.gz

{weco-0.2.22 → weco-0.2.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.22
+Version: 0.2.23
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -30,10 +30,11 @@ Dynamic: license-file
 </div>
 
 [![Python](https://img.shields.io/badge/Python-3.8.0+-blue)](https://www.python.org)
+[![PyPI version](https://img.shields.io/pypi/v/weco?label=PyPI%20version&color=f05138&labelColor=555555)](https://badge.fury.io/py/weco)
 [![docs](https://img.shields.io/website?url=https://docs.weco.ai/&label=docs)](https://docs.weco.ai/)
-[![PyPI version](https://badge.fury.io/py/weco.svg)](https://badge.fury.io/py/weco)
-[![AIDE](https://img.shields.io/badge/AI--Driven_Exploration-arXiv-orange?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2502.13138)
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
+[![PyPI Downloads](https://static.pepy.tech/badge/weco?color=4c1)](https://pepy.tech/projects/weco)
+[![arXiv on AIDE](https://img.shields.io/badge/arXiv-AIDE-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2502.13138)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg?labelColor=ffffff&color=F17E01)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
 
 `pip install weco`
 
@@ -73,9 +74,9 @@ The `weco` CLI leverages a tree search approach guided by LLMs to iteratively ex
 
     `weco` requires API keys for the LLMs it uses internally. You **must** provide these keys via environment variables:
 
-    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your API key [here](https://platform.openai.com/api-keys))
-    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your API key [here](https://console.anthropic.com/settings/keys))
-    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
+    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your OpenAI API key [here](https://platform.openai.com/api-keys))
+    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your Anthropic API key [here](https://console.anthropic.com/settings/keys))
+    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your Gemini API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
 
 ---
 
@@ -252,20 +253,17 @@ Weco will parse this output to extract the numerical value (1.5 in this case) as
 We welcome your contributions! To get started:
 
 1.  **Fork & Clone the Repository:**
-
     ```bash
     git clone https://github.com/WecoAI/weco-cli.git
     cd weco-cli
     ```
 
 2.  **Install Dependencies:**
-
     ```bash
     pip install -e ".[dev]"
     ```
 
 3.  **Create a Feature Branch:**
-
     ```bash
     git checkout -b feature/your-feature-name
     ```

{weco-0.2.22 → weco-0.2.23}/README.md

@@ -6,10 +6,11 @@
 </div>
 
 [![Python](https://img.shields.io/badge/Python-3.8.0+-blue)](https://www.python.org)
+[![PyPI version](https://img.shields.io/pypi/v/weco?label=PyPI%20version&color=f05138&labelColor=555555)](https://badge.fury.io/py/weco)
 [![docs](https://img.shields.io/website?url=https://docs.weco.ai/&label=docs)](https://docs.weco.ai/)
-[![PyPI version](https://badge.fury.io/py/weco.svg)](https://badge.fury.io/py/weco)
-[![AIDE](https://img.shields.io/badge/AI--Driven_Exploration-arXiv-orange?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2502.13138)
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
+[![PyPI Downloads](https://static.pepy.tech/badge/weco?color=4c1)](https://pepy.tech/projects/weco)
+[![arXiv on AIDE](https://img.shields.io/badge/arXiv-AIDE-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2502.13138)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg?labelColor=ffffff&color=F17E01)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
 
 `pip install weco`
 
@@ -49,9 +50,9 @@ The `weco` CLI leverages a tree search approach guided by LLMs to iteratively ex
 
     `weco` requires API keys for the LLMs it uses internally. You **must** provide these keys via environment variables:
 
-    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your API key [here](https://platform.openai.com/api-keys))
-    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your API key [here](https://console.anthropic.com/settings/keys))
-    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
+    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your OpenAI API key [here](https://platform.openai.com/api-keys))
+    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your Anthropic API key [here](https://console.anthropic.com/settings/keys))
+    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your Gemini API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
 
 ---
 
@@ -228,20 +229,17 @@ Weco will parse this output to extract the numerical value (1.5 in this case) as
 We welcome your contributions! To get started:
 
 1.  **Fork & Clone the Repository:**
-
     ```bash
     git clone https://github.com/WecoAI/weco-cli.git
     cd weco-cli
     ```
 
 2.  **Install Dependencies:**
-
     ```bash
     pip install -e ".[dev]"
     ```
 
 3.  **Create a Feature Branch:**
-
     ```bash
     git checkout -b feature/your-feature-name
     ```

{weco-0.2.22 → weco-0.2.23}/examples/cuda/README.md

@@ -11,7 +11,7 @@ Install the CLI using `pip`:
 pip install weco
 ```
 
-Set up your API key:
+Create your OpenAI API key [here](https://platform.openai.com/api-keys), then run:
 ```bash
 export OPENAI_API_KEY="your_key_here"
 ```
@@ -46,3 +46,9 @@ weco run --source optimize.py \
 *   `--additional-instructions guide.md`: Provides guidance to the LLM on the optimization approach.
 
 Weco will iteratively modify `optimize.py`, generating and integrating CUDA C++ code, guided by the evaluation results and the instructions in `guide.md`.
+
+## Next Steps
+
+Now that you've optimized your code with CUDA kernels, try [Triton Optimization](/examples/triton/README.md) for a higher-level GPU programming approach. If you're more interested in [Model Development](/examples/spaceship-titanic/README.md) or [Prompt Engineering](/examples/prompt/README.md), we've got you covered! 
+
+You can check out our [CLI Reference](https://docs.weco.ai/cli/cli-reference) to learn more about what you can do with the tool.

{weco-0.2.22 → weco-0.2.23}/examples/cuda/evaluate.py

@@ -110,6 +110,7 @@ if __name__ == "__main__":
 
     # benchmarking parameters
     n_correctness_trials = 10
+    correctness_tolerance = 1e-5
     n_warmup = 1000
     n_rep = 5000
 
@@ -152,6 +153,8 @@ if __name__ == "__main__":
             max_diff_avg += torch.max(torch.abs(optimized_output - baseline_output))
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
+    if max_diff_avg > correctness_tolerance:
+        print("invalid solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size=batch_size, seq_len=seq_len, n_embd=n_embd, device="cuda")

{weco-0.2.22 → weco-0.2.23}/examples/hello-kernel-world/evaluate.py

@@ -101,6 +101,7 @@ if __name__ == "__main__":
 
     # benchmark parameters
     n_correctness_trials = 10
+    correctness_tolerance = 1e-5
     n_warmup = 1000
     n_rep = 5000
 
@@ -130,6 +131,8 @@ if __name__ == "__main__":
         max_diff_avg += torch.max(torch.abs(optimized_output - baseline_output))
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
+    if max_diff_avg > correctness_tolerance:
+        print("invalid solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size, input_size, args.device)

{weco-0.2.22 → weco-0.2.23}/examples/prompt/README.md

@@ -16,7 +16,7 @@ The experiment runs locally, requires only two short Python files and a prompt g
    pip install weco openai datasets
    ```
 
-3. This example uses `o4-mini` via the OpenAI API by default. Set up your API key:
+3. This example uses `o4-mini` via the OpenAI API by default. Create your OpenAI API key [here](https://platform.openai.com/api-keys), then run:
    ```bash
    export OPENAI_API_KEY="your_key_here"
    ```
@@ -62,3 +62,10 @@ Weco then mutates the prompt instructions in `optimize.py`, tries again, and gra
 *   The script sends model calls in parallel via `ThreadPoolExecutor`, so network latency is hidden.
 *   Every five completed items, the script logs progress and elapsed time.
 *   The final line `accuracy: value` is the only part Weco needs for guidance.
+
+## Next Steps
+
+Now that you've automated prompt engineering for yourself, check out our guide on [Model Development](/examples/spaceship-titanic/README.md) or [CUDA Kernel Engineering](/examples/cuda/README.md).
+
+You can check out our [CLI Reference](https://docs.weco.ai/cli/cli-reference) to learn more about what you can do with the tool.
+

{weco-0.2.22 → weco-0.2.23}/examples/spaceship-titanic/README.md

@@ -10,7 +10,7 @@ The goal is to improve the model's `accuracy` metric by optimizing the `train.py
     ```bash
     pip install weco pandas numpy scikit-learn torch xgboost lightgbm catboost
     ```
-3. Set up your API key:
+3. Create your OpenAI API key [here](https://platform.openai.com/api-keys), then run:
     ```bash
     export OPENAI_API_KEY="your_key_here"
     ```
@@ -44,3 +44,7 @@ weco run --source train.py \
 *   `--log-dir .runs/spaceship-titanic`: Specifies the directory where Weco should save logs and results for this run.
 
 Weco will iteratively update the feature engineering or modeling code within `train.py` guided by the evaluation method defined in `evaluate.py`
+
+## Next Steps
+
+With model development covered, you might be curious to see how you can make your AI code run faster, saving you time and more importantly GPU credits. Check out our example on automating kernel engineering in [CUDA](/examples/cuda/README.md) and [Triton](/examples/triton/README.md), or dive into the [CLI Reference](https://docs.weco.ai/cli/cli-reference).

{weco-0.2.22 → weco-0.2.23}/examples/triton/README.md

@@ -12,7 +12,7 @@ Install the CLI using `pip`:
 pip install weco
 ```
 
-Set up your API key:
+Create your OpenAI API key [here](https://platform.openai.com/api-keys), then run:
 ```bash
 export OPENAI_API_KEY="your_key_here"
 ```
@@ -47,3 +47,7 @@ weco run --source optimize.py \
 *   `--additional-instructions "..."`: Provides specific guidance to the LLM. In this case, it directs the model to use Triton for optimization, ensure the numerical difference ("max float diff") between the original and optimized code remains small, and keep the overall code structure consistent.
 
 Weco will iteratively modify `optimize.py`, incorporating Triton kernels, guided by the performance feedback (`speedup`) from the evaluation script and the instructions provided.
+
+## Next Steps
+
+After mastering Triton kernels, explore [CUDA Optimization](/examples/cuda/README.md) for even lower-level GPU programming, or check the [CLI Reference](https://docs.weco.ai/cli/cli-reference) to improve the results you get with Weco.

{weco-0.2.22 → weco-0.2.23}/examples/triton/evaluate.py

@@ -105,6 +105,7 @@ if __name__ == "__main__":
 
     # benchmarking parameters
     n_correctness_trials = 10
+    correctness_tolerance = 1e-5
     n_warmup = 1000
     n_rep = 5000
 
@@ -147,6 +148,8 @@ if __name__ == "__main__":
             max_diff_avg += torch.max(torch.abs(optimized_output - baseline_output))
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
+    if max_diff_avg > correctness_tolerance:
+        print("invalid solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size=batch_size, seq_len=seq_len, n_embd=n_embd, device="cuda")

{weco-0.2.22 → weco-0.2.23}/pyproject.toml

@@ -8,7 +8,7 @@ name = "weco"
 authors = [{ name = "Weco AI Team", email = "contact@weco.ai" }]
 description = "Documentation for `weco`, a CLI for using Weco AI's code optimizer."
 readme = "README.md"
-version = "0.2.22"
+version = "0.2.23"
 license = { text = "MIT" }
 requires-python = ">=3.8"
 dependencies = [

{weco-0.2.22 → weco-0.2.23}/weco/api.py

@@ -1,29 +1,11 @@
 import sys
 from typing import Dict, Any, Optional, Union, Tuple, List
-
 import requests
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
 from rich.console import Console
 
 from weco import __pkg_version__, __base_url__
 
 
-# --- Session Configuration ---
-def _get_weco_session() -> requests.Session:
-    session = requests.Session()
-    retry_strategy = Retry(
-        total=3,
-        status_forcelist=[429, 500, 502, 503, 504],  # Retry on these server errors and rate limiting
-        allowed_methods=["HEAD", "GET", "PUT", "POST", "DELETE", "OPTIONS"],  # Case-insensitive
-        backoff_factor=1,  # e.g., sleep for 0s, 2s, 4s between retries (factor * (2 ** ({number of total retries} - 1)))
-    )
-    adapter = HTTPAdapter(max_retries=retry_strategy)
-    session.mount("http://", adapter)
-    session.mount("https://", adapter)
-    return session
-
-
 def handle_api_error(e: requests.exceptions.HTTPError, console: Console) -> None:
     """Extract and display error messages from API responses in a structured format."""
     try:
@@ -53,8 +35,7 @@ def start_optimization_run(
     """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
         try:
-            session = _get_weco_session()
-            response = session.post(
+            response = requests.post(
                 f"{__base_url__}/runs",
                 json={
                     "source_code": source_code,
@@ -76,8 +57,8 @@ def start_optimization_run(
         except requests.exceptions.HTTPError as e:
             handle_api_error(e, console)
             sys.exit(1)
-        except requests.exceptions.RequestException as e:
-            console.print(f"[bold red]Network Error starting run: {e}[/]")
+        except Exception as e:
+            console.print(f"[bold red]Error starting run: {e}[/]")
             sys.exit(1)
 
 
@@ -91,8 +72,7 @@ def evaluate_feedback_then_suggest_next_solution(
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/runs/{run_id}/suggest",
             json={
                 "execution_output": execution_output,
@@ -108,8 +88,8 @@ def evaluate_feedback_then_suggest_next_solution(
         # Allow caller to handle suggest errors, maybe retry or terminate
         handle_api_error(e, Console())  # Use default console if none passed
         raise  # Re-raise the exception
-    except requests.exceptions.RequestException as e:
-        print(f"Network Error during suggest: {e}")  # Use print as console might not be available
+    except Exception as e:
+        print(f"Error: {e}")  # Use print as console might not be available
         raise  # Re-raise the exception
 
 
@@ -118,8 +98,7 @@ def get_optimization_run_status(
 ) -> Dict[str, Any]:
     """Get the current status of the optimization run."""
     try:
-        session = _get_weco_session()
-        response = session.get(
+        response = requests.get(
             f"{__base_url__}/runs/{run_id}", params={"include_history": include_history}, headers=auth_headers, timeout=timeout
         )
         response.raise_for_status()
@@ -127,16 +106,15 @@ def get_optimization_run_status(
     except requests.exceptions.HTTPError as e:
         handle_api_error(e, Console())  # Use default console
         raise  # Re-raise
-    except requests.exceptions.RequestException as e:
-        print(f"Network Error getting status: {e}")
+    except Exception as e:
+        print(f"Error getting run status: {e}")
         raise  # Re-raise
 
 
 def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 10) -> bool:
     """Send a heartbeat signal to the backend."""
     try:
-        session = _get_weco_session()
-        response = session.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
+        response = requests.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
         response.raise_for_status()
         return True
     except requests.exceptions.HTTPError as e:
@@ -145,8 +123,8 @@ def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tup
         else:
             print(f"Heartbeat failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
         return False
-    except requests.exceptions.RequestException as e:
-        print(f"Heartbeat network error for run {run_id}: {e}", file=sys.stderr)
+    except Exception as e:
+        print(f"Error sending heartbeat for run {run_id}: {e}", file=sys.stderr)
         return False
 
 
@@ -160,8 +138,7 @@ def report_termination(
 ) -> bool:
     """Report the termination reason to the backend."""
     try:
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/runs/{run_id}/terminate",
             json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
             headers=auth_headers,
@@ -169,7 +146,7 @@ def report_termination(
         )
         response.raise_for_status()
         return True
-    except requests.exceptions.RequestException as e:
+    except Exception as e:
         print(f"Warning: Failed to report termination to backend for run {run_id}: {e}", file=sys.stderr)
         return False
 
@@ -211,8 +188,7 @@ def get_optimization_suggestions_from_codebase(
     """Analyze codebase and get optimization suggestions using the model-agnostic backend API."""
     try:
         model, api_key_dict = _determine_model_and_api_key()
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/onboard/analyze-codebase",
             json={
                 "gitingest_summary": gitingest_summary,
@@ -231,11 +207,8 @@ def get_optimization_suggestions_from_codebase(
     except requests.exceptions.HTTPError as e:
         handle_api_error(e, console)
         return None
-    except requests.exceptions.RequestException as e:
-        console.print(f"[bold red]Network Error getting optimization suggestions: {e}[/]")
-        return None
     except Exception as e:
-        console.print(f"[bold red]Error calling backend API: {e}[/]")
+        console.print(f"[bold red]Error: {e}[/]")
         return None
 
 
@@ -250,8 +223,7 @@ def generate_evaluation_script_and_metrics(
     """Generate evaluation script and determine metrics using the model-agnostic backend API."""
     try:
         model, api_key_dict = _determine_model_and_api_key()
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/onboard/generate-script",
             json={
                 "target_file": target_file,
@@ -266,15 +238,11 @@ def generate_evaluation_script_and_metrics(
         response.raise_for_status()
         result = response.json()
         return result.get("script_content"), result.get("metric_name"), result.get("goal"), result.get("reasoning")
-
     except requests.exceptions.HTTPError as e:
         handle_api_error(e, console)
         return None, None, None, None
-    except requests.exceptions.RequestException as e:
-        console.print(f"[bold red]Network Error generating evaluation script: {e}[/]")
-        return None, None, None, None
     except Exception as e:
-        console.print(f"[bold red]Error calling backend API: {e}[/]")
+        console.print(f"[bold red]Error: {e}[/]")
         return None, None, None, None
 
 
@@ -291,8 +259,7 @@ def analyze_evaluation_environment(
     """Analyze existing evaluation scripts and environment using the model-agnostic backend API."""
     try:
         model, api_key_dict = _determine_model_and_api_key()
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/onboard/analyze-environment",
             json={
                 "target_file": target_file,
@@ -312,11 +279,8 @@ def analyze_evaluation_environment(
     except requests.exceptions.HTTPError as e:
         handle_api_error(e, console)
         return None
-    except requests.exceptions.RequestException as e:
-        console.print(f"[bold red]Network Error analyzing evaluation environment: {e}[/]")
-        return None
     except Exception as e:
-        console.print(f"[bold red]Error calling backend API: {e}[/]")
+        console.print(f"[bold red]Error: {e}[/]")
         return None
 
 
@@ -331,8 +295,7 @@ def analyze_script_execution_requirements(
     """Analyze script to determine proper execution command using the model-agnostic backend API."""
     try:
         model, api_key_dict = _determine_model_and_api_key()
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/onboard/analyze-script",
             json={
                 "script_content": script_content,
@@ -351,9 +314,6 @@ def analyze_script_execution_requirements(
     except requests.exceptions.HTTPError as e:
         handle_api_error(e, console)
         return f"python {script_path}"
-    except requests.exceptions.RequestException as e:
-        console.print(f"[bold red]Network Error analyzing script execution: {e}[/]")
-        return f"python {script_path}"
     except Exception as e:
-        console.print(f"[bold red]Error calling backend API: {e}[/]")
+        console.print(f"[bold red]Error: {e}[/]")
         return f"python {script_path}"

{weco-0.2.22 → weco-0.2.23}/weco.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.22
+Version: 0.2.23
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -30,10 +30,11 @@ Dynamic: license-file
 </div>
 
 [![Python](https://img.shields.io/badge/Python-3.8.0+-blue)](https://www.python.org)
+[![PyPI version](https://img.shields.io/pypi/v/weco?label=PyPI%20version&color=f05138&labelColor=555555)](https://badge.fury.io/py/weco)
 [![docs](https://img.shields.io/website?url=https://docs.weco.ai/&label=docs)](https://docs.weco.ai/)
-[![PyPI version](https://badge.fury.io/py/weco.svg)](https://badge.fury.io/py/weco)
-[![AIDE](https://img.shields.io/badge/AI--Driven_Exploration-arXiv-orange?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2502.13138)
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
+[![PyPI Downloads](https://static.pepy.tech/badge/weco?color=4c1)](https://pepy.tech/projects/weco)
+[![arXiv on AIDE](https://img.shields.io/badge/arXiv-AIDE-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2502.13138)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg?labelColor=ffffff&color=F17E01)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
 
 `pip install weco`
 
@@ -73,9 +74,9 @@ The `weco` CLI leverages a tree search approach guided by LLMs to iteratively ex
 
     `weco` requires API keys for the LLMs it uses internally. You **must** provide these keys via environment variables:
 
-    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your API key [here](https://platform.openai.com/api-keys))
-    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your API key [here](https://console.anthropic.com/settings/keys))
-    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
+    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your OpenAI API key [here](https://platform.openai.com/api-keys))
+    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your Anthropic API key [here](https://console.anthropic.com/settings/keys))
+    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your Gemini API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
 
 ---
 
@@ -252,20 +253,17 @@ Weco will parse this output to extract the numerical value (1.5 in this case) as
 We welcome your contributions! To get started:
 
 1.  **Fork & Clone the Repository:**
-
     ```bash
     git clone https://github.com/WecoAI/weco-cli.git
     cd weco-cli
     ```
 
 2.  **Install Dependencies:**
-
     ```bash
     pip install -e ".[dev]"
     ```
 
 3.  **Create a Feature Branch:**
-
     ```bash
     git checkout -b feature/your-feature-name
     ```