weco 0.2.1__tar.gz → 0.2.3__tar.gz

{weco-0.2.1 → weco-0.2.3}/.github/workflows/release.yml

@@ -90,7 +90,7 @@ jobs:
         GITHUB_TOKEN: ${{ github.token }}
       run: >-
         gh release create
-        'v0.2.1'
+        'v0.2.3'
         --repo '${{ github.repository }}'
         --notes ""
 
@@ -102,5 +102,5 @@ jobs:
       # sigstore-produced signatures and certificates.
       run: >-
         gh release upload
-        'v0.2.1' dist/**
+        'v0.2.3' dist/**
         --repo '${{ github.repository }}'

{weco-0.2.1 → weco-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.1
+Version: 0.2.3
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <dhruv@weco.ai>
 License: MIT
@@ -20,18 +20,28 @@ Requires-Dist: build; extra == "dev"
 Requires-Dist: setuptools_scm; extra == "dev"
 Dynamic: license-file
 
-# Weco CLI – Optimize Your Code Effortlessly
+# Weco CLI – Code Optimizer for Machine Learning Engineers
 
 [![Python](https://img.shields.io/badge/Python-3.12.0-blue)](https://www.python.org)
 [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)
 
-`weco` is a powerful command-line interface for interacting with Weco AI's code optimizer. Whether you are looking to improve performance or refine code quality, our CLI streamlines your workflow for a better development experience.
+`weco` is a command-line interface for interacting with Weco AI's code optimizer, powerred by [AI-Driven Exploration](https://arxiv.org/abs/2502.13138).
+
+
+
+https://github.com/user-attachments/assets/cb724ef1-bff6-4757-b457-d3b2201ede81
+
+
 
 ---
 
 ## Overview
 
-The `weco` CLI leverages advanced optimization techniques and language model strategies to iteratively improve your source code. It supports multiple language models and offers a flexible configuration to suit different optimization tasks.
+The weco CLI leverages a tree search approach with LLMs to iteratively improve your code.
+
+![image](https://github.com/user-attachments/assets/a6ed63fa-9c40-498e-aa98-a873e5786509)
+
+
 
 ---
 
@@ -59,13 +69,13 @@ The `weco` CLI leverages advanced optimization techniques and language model str
 
 | Argument                    | Description                                                                                                                                   | Required |
 |-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------|----------|
-| `--source`                  | Path to the Python source code that will be optimized (e.g. optimize.py).                                                                    | Yes      |
+| `--source`                  | Path to the Python source code that will be optimized (e.g. optimize.py).                                                                     | Yes      |
 | `--eval-command`            | Command to run for evaluation (e.g. 'python eval.py --arg1=val1').                                                                            | Yes      |
 | `--metric`                  | Metric to optimize.                                                                                                                           | Yes      |
-| `--maximize`                | Boolean flag indicating whether to maximize the metric.                                                                                       | Yes      |
+| `--maximize`                | Whether to maximize ('true') or minimize ('false') the metric.                                                                                | Yes      |
 | `--steps`                   | Number of optimization steps to run.                                                                                                          | Yes      |
 | `--model`                   | Model to use for optimization.                                                                                                                | Yes      |
-| `--additional-instructions` | (Optional) Description of additional instructions or path to a file containing additional instructions.                                       | No       |
+| `--additional-instructions` | (Optional) Description of additional instructions OR path to a file containing additional instructions.                                       | No       |
 
 ---
 
@@ -75,7 +85,7 @@ Optimizing common operations in pytorch:
 ```bash
 weco --source examples/simple-torch/optimize.py \
      --eval-command "python examples/simple-torch/evaluate.py --solution-path examples/simple-torch/optimize.py --device mps" \
-     --metric "speedup" \
+     --metric speedup \
      --maximize true \
      --steps 15 \
      --model "o3-mini" \
@@ -86,7 +96,7 @@ Optimizing these same using mlx and metal:
 ```bash
 weco --source examples/simple-mlx/optimize.py \
      --eval-command "python examples/simple-mlx/evaluate.py --solution-path examples/simple-mlx/optimize.py" \
-     --metric "speedup" \
+     --metric speedup \
      --maximize true \
      --steps 30 \
      --model "o3-mini" \

{weco-0.2.1 → weco-0.2.3}/README.md

@@ -1,15 +1,25 @@
-# Weco CLI – Optimize Your Code Effortlessly
+# Weco CLI – Code Optimizer for Machine Learning Engineers
 
 [![Python](https://img.shields.io/badge/Python-3.12.0-blue)](https://www.python.org)
 [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)
 
-`weco` is a powerful command-line interface for interacting with Weco AI's code optimizer. Whether you are looking to improve performance or refine code quality, our CLI streamlines your workflow for a better development experience.
+`weco` is a command-line interface for interacting with Weco AI's code optimizer, powerred by [AI-Driven Exploration](https://arxiv.org/abs/2502.13138).
+
+
+
+https://github.com/user-attachments/assets/cb724ef1-bff6-4757-b457-d3b2201ede81
+
+
 
 ---
 
 ## Overview
 
-The `weco` CLI leverages advanced optimization techniques and language model strategies to iteratively improve your source code. It supports multiple language models and offers a flexible configuration to suit different optimization tasks.
+The weco CLI leverages a tree search approach with LLMs to iteratively improve your code.
+
+![image](https://github.com/user-attachments/assets/a6ed63fa-9c40-498e-aa98-a873e5786509)
+
+
 
 ---
 
@@ -37,13 +47,13 @@ The `weco` CLI leverages advanced optimization techniques and language model str
 
 | Argument                    | Description                                                                                                                                   | Required |
 |-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------|----------|
-| `--source`                  | Path to the Python source code that will be optimized (e.g. optimize.py).                                                                    | Yes      |
+| `--source`                  | Path to the Python source code that will be optimized (e.g. optimize.py).                                                                     | Yes      |
 | `--eval-command`            | Command to run for evaluation (e.g. 'python eval.py --arg1=val1').                                                                            | Yes      |
 | `--metric`                  | Metric to optimize.                                                                                                                           | Yes      |
-| `--maximize`                | Boolean flag indicating whether to maximize the metric.                                                                                       | Yes      |
+| `--maximize`                | Whether to maximize ('true') or minimize ('false') the metric.                                                                                | Yes      |
 | `--steps`                   | Number of optimization steps to run.                                                                                                          | Yes      |
 | `--model`                   | Model to use for optimization.                                                                                                                | Yes      |
-| `--additional-instructions` | (Optional) Description of additional instructions or path to a file containing additional instructions.                                       | No       |
+| `--additional-instructions` | (Optional) Description of additional instructions OR path to a file containing additional instructions.                                       | No       |
 
 ---
 
@@ -53,7 +63,7 @@ Optimizing common operations in pytorch:
 ```bash
 weco --source examples/simple-torch/optimize.py \
      --eval-command "python examples/simple-torch/evaluate.py --solution-path examples/simple-torch/optimize.py --device mps" \
-     --metric "speedup" \
+     --metric speedup \
      --maximize true \
      --steps 15 \
      --model "o3-mini" \
@@ -64,7 +74,7 @@ Optimizing these same using mlx and metal:
 ```bash
 weco --source examples/simple-mlx/optimize.py \
      --eval-command "python examples/simple-mlx/evaluate.py --solution-path examples/simple-mlx/optimize.py" \
-     --metric "speedup" \
+     --metric speedup \
      --maximize true \
      --steps 30 \
      --model "o3-mini" \

{weco-0.2.1 → weco-0.2.3}/examples/simple-mlx/optimize.py

@@ -19,8 +19,8 @@ class Model(nn.Module):
         Returns:
             mx.array: Output tensor of shape (batch_size, hidden_size).
         """
-        x = mx.matmul(x, mx.transpose(self.weight))  # Gemm
-        x = x / 2  # Divide
-        x = mx.sum(x, axis=1, keepdims=True)  # Sum
-        x = x * self.scaling_factor  # Scaling
+        x = mx.matmul(x, mx.transpose(self.weight))
+        x = x / 2
+        x = mx.sum(x, axis=1, keepdims=True)
+        x = x * self.scaling_factor
         return x

{weco-0.2.1 → weco-0.2.3}/examples/simple-torch/optimize.py

@@ -19,8 +19,8 @@ class Model(nn.Module):
         Returns:
             torch.Tensor: Output tensor of shape (batch_size, hidden_size).
         """
-        x = torch.matmul(x, self.weight.T)  # Gemm
-        x = x / 2  # Divide
-        x = torch.sum(x, dim=1, keepdim=True)  # Sum
-        x = x * self.scaling_factor  # Scaling
+        x = torch.matmul(x, self.weight.T)
+        x = x / 2
+        x = torch.sum(x, dim=1, keepdim=True)
+        x = x * self.scaling_factor
         return x

{weco-0.2.1 → weco-0.2.3}/pyproject.toml

@@ -10,7 +10,7 @@ authors = [
 ]
 description = "Documentation for `weco`, a CLI for using Weco AI's code optimizer."
 readme = "README.md"
-version = "0.2.1"
+version = "0.2.3"
 license = {text = "MIT"}
 requires-python = ">=3.12"
 dependencies = ["requests", "rich"]

{weco-0.2.1 → weco-0.2.3}/weco/__init__.py

@@ -1,4 +1,4 @@
 # DO NOT EDIT
-__pkg_version__ = "0.2.1"
+__pkg_version__ = "0.2.3"
 __api_version__ = "v1"
 __base_url__ = f"https://api.aide.weco.ai/{__api_version__}"

{weco-0.2.1 → weco-0.2.3}/weco/cli.py

@@ -41,7 +41,13 @@ def main() -> None:
         "--eval-command", type=str, required=True, help="Command to run for evaluation (e.g. 'python eval.py --arg1=val1')"
     )
     parser.add_argument("--metric", type=str, required=True, help="Metric to optimize")
-    parser.add_argument("--maximize", type=bool, required=True, help="Maximize the metric")
+    parser.add_argument(
+        "--maximize",
+        type=str,
+        choices=["true", "false"],
+        required=True,
+        help="Specify 'true' to maximize the metric or 'false' to minimize.",
+    )
     parser.add_argument("--steps", type=int, required=True, help="Number of steps to run")
     parser.add_argument("--model", type=str, required=True, help="Model to use for optimization")
     parser.add_argument(
@@ -57,7 +63,7 @@ def main() -> None:
             # Define optimization session config
             evaluation_command = args.eval_command
             metric_name = args.metric
-            maximize = args.maximize
+            maximize = args.maximize == "true"
             steps = args.steps
             code_generator_config = {"model": args.model}
             evaluator_config = {"model": args.model}
@@ -75,9 +81,9 @@ def main() -> None:
             api_keys = read_api_keys_from_env()
 
         # Initialize panels
-        summary_panel = SummaryPanel(total_steps=steps, model=args.model)
+        summary_panel = SummaryPanel(maximize=maximize, metric_name=metric_name, total_steps=steps, model=args.model)
         plan_panel = PlanPanel()
-        solution_panels = SolutionPanels()
+        solution_panels = SolutionPanels(metric_name=metric_name)
         eval_output_panel = EvaluationOutputPanel()
         tree_panel = MetricTreePanel(maximize=maximize)
         layout = create_optimization_layout()
@@ -185,6 +191,9 @@ def main() -> None:
                 # Save next solution (.runs/<session-id>/step_<step>.py)
                 write_to_path(fp=runs_dir / f"step_{step}.py", content=eval_and_next_solution_response["code"])
 
+                # Write the next solution to the source file
+                write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
+
                 # Get the optimization session status for
                 # the best solution, its score, and the history to plot the tree
                 status_response = get_optimization_session_status(console=console, session_id=session_id, include_history=True)

{weco-0.2.1 → weco-0.2.3}/weco/panels.py

@@ -11,12 +11,13 @@ from .utils import format_number
 class SummaryPanel:
     """Holds a summary of the optimization session."""
 
-    def __init__(self, total_steps: int, model: str, session_id: str = None):
+    def __init__(self, maximize: bool, metric_name: str, total_steps: int, model: str, session_id: str = None):
+        self.goal = ("Maximizing" if maximize else "Minimizing") + f" {metric_name}"
         self.total_input_tokens = 0
         self.total_output_tokens = 0
         self.total_steps = total_steps
         self.model = model
-        self.session_id = session_id or "unknown"
+        self.session_id = session_id or "N/A"
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=20),
@@ -42,6 +43,9 @@ class SummaryPanel:
         """Create a summary panel with the relevant information."""
         layout = Layout(name="summary")
         summary_table = Table(show_header=False, box=None, padding=(0, 1))
+        # Goal
+        summary_table.add_row(f"[bold cyan]Goal:[/] {self.goal}")
+        summary_table.add_row("")
         # Log directory
         runs_dir = f".runs/{self.session_id}"
         summary_table.add_row(f"[bold cyan]Logs:[/] [blue]{runs_dir}[/]")
@@ -172,7 +176,7 @@ class MetricTreePanel:
     def _build_rich_tree(self) -> Tree:
         """Get a Rich Tree representation of the solution tree using a DFS like traversal."""
         if len(self.metric_tree.nodes) == 0:
-            return Tree("[bold green]Building 🌳")
+            return Tree("[bold green]Building first solution...")
 
         best_node = self.metric_tree.get_best_node()
 
@@ -193,7 +197,7 @@ class MetricTreePanel:
                     # best node
                     color = "green"
                     style = "bold"
-                    text = f"{node.metric:.3f} (best)"
+                    text = f"🏆 {node.metric:.3f}"
                 elif node.metric is None:
                     # metric not extracted from evaluated solution
                     color = "yellow"
@@ -210,7 +214,7 @@ class MetricTreePanel:
             for child in node.children:
                 append_rec(child, subtree)
 
-        tree = Tree("🌳")
+        tree = Tree("solutions")
         for n in self.metric_tree.get_draft_nodes():
             append_rec(n, tree)
 
@@ -219,9 +223,7 @@ class MetricTreePanel:
     def get_display(self) -> Panel:
         """Get a panel displaying the solution tree."""
         # Make sure the metric tree is built before calling build_rich_tree
-        return Panel(
-            self._build_rich_tree(), title="[bold]🌳 Solution Tree", border_style="green", expand=True, padding=(0, 1)
-        )
+        return Panel(self._build_rich_tree(), title="[bold]🌳 Exploring...", border_style="green", expand=True, padding=(0, 1))
 
 
 class EvaluationOutputPanel:
@@ -246,11 +248,13 @@ class EvaluationOutputPanel:
 class SolutionPanels:
     """Displays the current and best solutions side by side."""
 
-    def __init__(self):
+    def __init__(self, metric_name: str):
         # Current solution
         self.current_node = None
         # Best solution
         self.best_node = None
+        # Metric name
+        self.metric_name = metric_name.capitalize()
 
     def update(self, current_node: Union[Node, None], best_node: Union[Node, None]):
         """Update the current and best solutions."""
@@ -276,7 +280,7 @@ class SolutionPanels:
         )
 
         # Best solution
-        best_title = f"[bold]🏆 Best Solution ([green]Score: {f'{best_score:.4f}' if best_score is not None else 'N/A'}[/])"
+        best_title = f"[bold]🏆 Best Solution ([green]{self.metric_name}: {f'{best_score:.4f}' if best_score is not None else 'N/A'}[/])"
         best_panel = Panel(
             Syntax(str(best_code), "python", theme="monokai", line_numbers=True, word_wrap=False),
             title=best_title,

{weco-0.2.1 → weco-0.2.3}/weco.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.1
+Version: 0.2.3
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <dhruv@weco.ai>
 License: MIT
@@ -20,18 +20,28 @@ Requires-Dist: build; extra == "dev"
 Requires-Dist: setuptools_scm; extra == "dev"
 Dynamic: license-file
 
-# Weco CLI – Optimize Your Code Effortlessly
+# Weco CLI – Code Optimizer for Machine Learning Engineers
 
 [![Python](https://img.shields.io/badge/Python-3.12.0-blue)](https://www.python.org)
 [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)
 
-`weco` is a powerful command-line interface for interacting with Weco AI's code optimizer. Whether you are looking to improve performance or refine code quality, our CLI streamlines your workflow for a better development experience.
+`weco` is a command-line interface for interacting with Weco AI's code optimizer, powerred by [AI-Driven Exploration](https://arxiv.org/abs/2502.13138).
+
+
+
+https://github.com/user-attachments/assets/cb724ef1-bff6-4757-b457-d3b2201ede81
+
+
 
 ---
 
 ## Overview
 
-The `weco` CLI leverages advanced optimization techniques and language model strategies to iteratively improve your source code. It supports multiple language models and offers a flexible configuration to suit different optimization tasks.
+The weco CLI leverages a tree search approach with LLMs to iteratively improve your code.
+
+![image](https://github.com/user-attachments/assets/a6ed63fa-9c40-498e-aa98-a873e5786509)
+
+
 
 ---
 
@@ -59,13 +69,13 @@ The `weco` CLI leverages advanced optimization techniques and language model str
 
 | Argument                    | Description                                                                                                                                   | Required |
 |-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------|----------|
-| `--source`                  | Path to the Python source code that will be optimized (e.g. optimize.py).                                                                    | Yes      |
+| `--source`                  | Path to the Python source code that will be optimized (e.g. optimize.py).                                                                     | Yes      |
 | `--eval-command`            | Command to run for evaluation (e.g. 'python eval.py --arg1=val1').                                                                            | Yes      |
 | `--metric`                  | Metric to optimize.                                                                                                                           | Yes      |
-| `--maximize`                | Boolean flag indicating whether to maximize the metric.                                                                                       | Yes      |
+| `--maximize`                | Whether to maximize ('true') or minimize ('false') the metric.                                                                                | Yes      |
 | `--steps`                   | Number of optimization steps to run.                                                                                                          | Yes      |
 | `--model`                   | Model to use for optimization.                                                                                                                | Yes      |
-| `--additional-instructions` | (Optional) Description of additional instructions or path to a file containing additional instructions.                                       | No       |
+| `--additional-instructions` | (Optional) Description of additional instructions OR path to a file containing additional instructions.                                       | No       |
 
 ---
 
@@ -75,7 +85,7 @@ Optimizing common operations in pytorch:
 ```bash
 weco --source examples/simple-torch/optimize.py \
      --eval-command "python examples/simple-torch/evaluate.py --solution-path examples/simple-torch/optimize.py --device mps" \
-     --metric "speedup" \
+     --metric speedup \
      --maximize true \
      --steps 15 \
      --model "o3-mini" \
@@ -86,7 +96,7 @@ Optimizing these same using mlx and metal:
 ```bash
 weco --source examples/simple-mlx/optimize.py \
      --eval-command "python examples/simple-mlx/evaluate.py --solution-path examples/simple-mlx/optimize.py" \
-     --metric "speedup" \
+     --metric speedup \
      --maximize true \
      --steps 30 \
      --model "o3-mini" \