weco 0.2.3__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.3"
+__pkg_version__ = "0.2.4"
 __api_version__ = "v1"
 __base_url__ = f"https://api.aide.weco.ai/{__api_version__}"

weco/cli.py

@@ -104,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

@@ -12,7 +12,7 @@ class SummaryPanel:
     """Holds a summary of the optimization session."""
 
     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.goal = ("Maximizing" if maximize else "Minimizing") + f" {metric_name}..."
         self.total_input_tokens = 0
         self.total_output_tokens = 0
         self.total_steps = total_steps
@@ -48,10 +48,10 @@ class SummaryPanel:
         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(
@@ -197,7 +197,7 @@ class MetricTreePanel:
                     # best node
                     color = "green"
                     style = "bold"
-                    text = f"🏆 {node.metric:.3f}"
+                    text = f"{node.metric:.3f} 🏆"
                 elif node.metric is None:
                     # metric not extracted from evaluated solution
                     color = "yellow"
@@ -214,7 +214,7 @@ class MetricTreePanel:
             for child in node.children:
                 append_rec(child, subtree)
 
-        tree = Tree("solutions")
+        tree = Tree("🌳")
         for n in self.metric_tree.get_draft_nodes():
             append_rec(n, tree)
 
@@ -223,7 +223,9 @@ 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]🌳 Exploring...", border_style="green", expand=True, padding=(0, 1))
+        return Panel(
+            self._build_rich_tree(), title="[bold]🔎 Exploring Solutions...", border_style="green", expand=True, padding=(0, 1)
+        )
 
 
 class EvaluationOutputPanel:
@@ -242,7 +244,7 @@ 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:

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.3.dist-info → weco-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.3
+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
@@ -88,19 +88,19 @@ weco --source examples/simple-torch/optimize.py \
      --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 \
      --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.3.dist-info/RECORD

@@ -1,11 +0,0 @@
-weco/__init__.py,sha256=bR_VMLdj9CKuS7r20P1w045efBJ-Jndx8NFtQVc3_78,124
-weco/api.py,sha256=JHAm60sO51Y2tzLnyAAFv9OLm75o_RNWeoK-gYF4VHE,3342
-weco/cli.py,sha256=Raq6X5ZSalPe0BV18UA_arC_jrNkQkJ4gDqWNTMlfsI,16242
-weco/panels.py,sha256=Y0K6Nx4uCBREdVL1c3ee7a3ZNjyhXJFThZBCToG3xO0,11970
-weco/utils.py,sha256=g1Lchd3_L-eGRNkqr2IyRZ5FJfZTrFf9aTmmpMtdhko,3761
-weco-0.2.3.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
-weco-0.2.3.dist-info/METADATA,sha256=GXFnAZjqNTXZGq8NC_zWvuq243ZZSsW0zYGrVEP8K0Y,5229
-weco-0.2.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-weco-0.2.3.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.2.3.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.2.3.dist-info/RECORD,,