weco 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

weco/__init__.py

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

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()
@@ -98,54 +104,59 @@ def main() -> None:
             api_keys=api_keys,
         )
 
-        # Define the runs directory (.runs/<session-id>)
-        session_id = session_response["session_id"]
-        runs_dir = pathlib.Path(".runs") / session_id
-        runs_dir.mkdir(parents=True, exist_ok=True)
-
-        # Save the original code (.runs/<session-id>/original.py)
-        runs_copy_source_fp = runs_dir / "original.py"
-        write_to_path(fp=runs_copy_source_fp, content=source_code)
-
-        # Write the code string to the source file path
-        # Do this after the original code is saved
-        write_to_path(fp=source_fp, content=session_response["code"])
-
-        # Update the panels with the initial solution
-        # Add session id now that we have it
-        summary_panel.session_id = session_id
-        # Set the step of the progress bar
-        summary_panel.set_step(step=0)
-        # Update the token counts
-        summary_panel.update_token_counts(usage=session_response["usage"])
-        # Update the plan
-        plan_panel.update(plan=session_response["plan"])
-        # Build the metric tree
-        tree_panel.build_metric_tree(
-            nodes=[
-                {
-                    "solution_id": session_response["solution_id"],
-                    "parent_id": None,
-                    "code": session_response["code"],
-                    "step": 0,
-                    "metric_value": None,
-                    "is_buggy": False,
-                }
-            ]
-        )
-        # Set the current solution as unevaluated since we haven't run the evaluation function and fed it back to the model yet
-        tree_panel.set_unevaluated_node(node_id=session_response["solution_id"])
-        # Update the solution panels with the initial solution and get the panel displays
-        solution_panels.update(
-            current_node=Node(
-                id=session_response["solution_id"], parent_id=None, code=session_response["code"], metric=None, is_buggy=False
-            ),
-            best_node=None,
-        )
-        current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=0)
         # Define the refresh rate
         refresh_rate = 4
         with Live(layout, refresh_per_second=refresh_rate, screen=True) as live:
+            # Define the runs directory (.runs/<session-id>)
+            session_id = session_response["session_id"]
+            runs_dir = pathlib.Path(".runs") / session_id
+            runs_dir.mkdir(parents=True, exist_ok=True)
+
+            # Save the original code (.runs/<session-id>/original.py)
+            runs_copy_source_fp = runs_dir / "original.py"
+            write_to_path(fp=runs_copy_source_fp, content=source_code)
+
+            # Write the code string to the source file path
+            # Do this after the original code is saved
+            write_to_path(fp=source_fp, content=session_response["code"])
+
+            # Update the panels with the initial solution
+            # Add session id now that we have it
+            summary_panel.session_id = session_id
+            # Set the step of the progress bar
+            summary_panel.set_step(step=0)
+            # Update the token counts
+            summary_panel.update_token_counts(usage=session_response["usage"])
+            # Update the plan
+            plan_panel.update(plan=session_response["plan"])
+            # Build the metric tree
+            tree_panel.build_metric_tree(
+                nodes=[
+                    {
+                        "solution_id": session_response["solution_id"],
+                        "parent_id": None,
+                        "code": session_response["code"],
+                        "step": 0,
+                        "metric_value": None,
+                        "is_buggy": False,
+                    }
+                ]
+            )
+            # Set the current solution as unevaluated since we haven't run the evaluation function and fed it back to the model yet
+            tree_panel.set_unevaluated_node(node_id=session_response["solution_id"])
+            # Update the solution panels with the initial solution and get the panel displays
+            solution_panels.update(
+                current_node=Node(
+                    id=session_response["solution_id"],
+                    parent_id=None,
+                    code=session_response["code"],
+                    metric=None,
+                    is_buggy=False,
+                ),
+                best_node=None,
+            )
+            current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=0)
+
             # Update the entire layout
             smooth_update(
                 live=live,

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,12 +43,15 @@ 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}[/]")
+        summary_table.add_row(f"[bold cyan]Logs:[/] [blue underline]{runs_dir}[/]")
         summary_table.add_row("")
         # Model used
-        summary_table.add_row(f"[bold cyan]Model:[/] {self.model}")
+        summary_table.add_row(f"[bold cyan]Model:[/] [yellow]{self.model}[/]")
         summary_table.add_row("")
         # Token counts
         summary_table.add_row(
@@ -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"
@@ -220,7 +224,7 @@ class MetricTreePanel:
         """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)
+            self._build_rich_tree(), title="[bold]🔎 Exploring Solutions...", border_style="green", expand=True, padding=(0, 1)
         )
 
 
@@ -240,17 +244,19 @@ class EvaluationOutputPanel:
 
     def get_display(self) -> Panel:
         """Create a panel displaying the evaluation output with truncation if needed."""
-        return Panel(self.output, title="[bold]📋 Evaluation Output", border_style="red", expand=True, padding=(0, 1))
+        return Panel(self.output, title="[bold]📋 Evaluation Output", border_style="blue", expand=True, padding=(0, 1))
 
 
 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 +282,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/utils.py

@@ -28,10 +28,14 @@ def read_additional_instructions(additional_instructions: str | None) -> str | N
 
     # Try interpreting as a path first
     potential_path = pathlib.Path(additional_instructions)
-    if potential_path.exists() and potential_path.is_file():
-        return read_from_path(potential_path, is_json=False)  # type: ignore # read_from_path returns str when is_json=False
-    else:
-        # If it's not a valid file path, return the string itself
+    try:
+        if potential_path.exists() and potential_path.is_file():
+            return read_from_path(potential_path, is_json=False)  # type: ignore # read_from_path returns str when is_json=False
+        else:
+            # If it's not a valid file path, return the string itself
+            return additional_instructions
+    except OSError:
+        # If the path can't be read, return the string itself
         return additional_instructions
 
 

{weco-0.2.2.dist-info → weco-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.2
+Version: 0.2.4
 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,22 +85,22 @@ 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" \
+     --model o3-mini \
      --additional-instructions "Fuse operations in the forward method while ensuring the max float deviation remains small. Maintain the same format of the code."
 ```
 
-Optimizing these same using mlx and metal:
+Sometimes we have a bit more context we'd like to provide. Its not easy to fit all of this in a string like shown above with `additional-instructions`. Thats why you can also provide a path to any file you'd like to me read as in context. In this example, we optimize the same operations using mlx and metal with additional instructions:
 ```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" \
-     --additional-instructions "examples/simple-mlx/metal-examples.rst"
+     --model o3-mini \
+     --additional-instructions examples/simple-mlx/metal-examples.rst
 ```
 ---
 

weco-0.2.4.dist-info/RECORD

@@ -0,0 +1,11 @@
+weco/__init__.py,sha256=5TUzMXAo2-PEbB3T_z4hNMJpe1yYKi36KZIeo5s1TCI,124
+weco/api.py,sha256=JHAm60sO51Y2tzLnyAAFv9OLm75o_RNWeoK-gYF4VHE,3342
+weco/cli.py,sha256=daKXAnwjnvUptaVA8nShzJiO9piU91c24s294rzu2d4,16492
+weco/panels.py,sha256=5HrDrYzx_vDjV41mlZmOOZYxdJIXD5RzBBNKXxUOLEY,12022
+weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
+weco-0.2.4.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
+weco-0.2.4.dist-info/METADATA,sha256=T2f_Fa3CdnfX1nwPFCbpY6a3lGYWx5BudSw4lcKxdCg,5518
+weco-0.2.4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+weco-0.2.4.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.2.4.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.2.4.dist-info/RECORD,,

weco-0.2.2.dist-info/RECORD

@@ -1,11 +0,0 @@
-weco/__init__.py,sha256=piNl2k8gLYD32s0kA8QqBDe7vANYnzN_QN9bWtfRsks,124
-weco/api.py,sha256=JHAm60sO51Y2tzLnyAAFv9OLm75o_RNWeoK-gYF4VHE,3342
-weco/cli.py,sha256=OneQ8a7ySgtzrNGo0l53Ij_I_p0RTVQ8GV2n66wIJv0,16050
-weco/panels.py,sha256=iUZWtZfohaWeI6-x6pf2kKN1c8IL7U_AhpBBeFhtFhk,11642
-weco/utils.py,sha256=g1Lchd3_L-eGRNkqr2IyRZ5FJfZTrFf9aTmmpMtdhko,3761
-weco-0.2.2.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
-weco-0.2.2.dist-info/METADATA,sha256=jy-GYv-J4qga01Ic2HtzkG1P3MDlgNjq_aas4EukZ0o,5270
-weco-0.2.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-weco-0.2.2.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.2.2.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.2.2.dist-info/RECORD,,