weco 0.1.10__tar.gz → 0.2.0__tar.gz

{weco-0.1.10 → weco-0.2.0}/.github/workflows/lint.yml

@@ -6,12 +6,8 @@ on:
       - main
       - dev
 
-  pull_request:
-    branches:
-      - main
-
 jobs:
-  build:
+  lint:
     runs-on: ubuntu-latest
 
     steps:
@@ -23,26 +19,20 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v3
       with:
-        python-version: "3.10"
+        python-version: "3.12.0"
 
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install flake8 flake8-pyproject black isort
+        pip install ruff
 
-    - name: Lint with flake8
+    - name: Run linter
       run: |
-        # uses the flake8 configuration in pyproject.toml
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 .
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --exit-zero
-
-    - name: Run black
-      run: black .
-
-    - name: Run isort
-      run: isort .
+        ruff check . --fix
+      
+    - name: Run formatter
+      run: |
+        ruff format .
 
     - name: Commit changes
       run: |
@@ -52,6 +42,6 @@ jobs:
         if git diff --exit-code --staged; then
           echo "No changes to commit"
         else
-          git commit -m "[PROJ] Format code with Black"
+          git commit -m "[GitHub Action] Lint and format code with Ruff"
           git push https://${{ secrets.GITHUB_TOKEN }}@github.com/${{ github.repository }}
-        fi
+        fi

{weco-0.1.10 → weco-0.2.0}/.github/workflows/release.yml

@@ -1,9 +1,10 @@
 name: Publish Python Package
 
 on:
-  push:
-    branches:
-      - 'main'
+  workflow_run:
+    workflows: ["Lint and Format Code"]
+    types:
+      - completed
 
   release:
     types: [published]
@@ -14,18 +15,25 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v4
+    - name: Checkout code
+      uses: actions/checkout@v4
+      with:
+        ref: ${{ github.head_ref }}
+
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: "3.10"
+        python-version: "3.12"
+
     - name: Install build dependencies
       run: >-
         python3 -m pip install build --user
+
     - name: Build a source distribution and a wheel
       run: python3 -m build
+
     - name: Store the distribution packages
-      uses: actions/upload-artifact@v3
+      uses: actions/upload-artifact@v4
       with:
         name: python-package-distributions
         path: dist/
@@ -44,10 +52,11 @@ jobs:
 
     steps:
     - name: Download the dists
-      uses: actions/download-artifact@v3
+      uses: actions/download-artifact@v4
       with:
         name: python-package-distributions
         path: dist/
+
     - name: Publish to PyPI
       uses: pypa/gh-action-pypi-publish@release/v1
 
@@ -64,24 +73,27 @@ jobs:
 
     steps:
     - name: Download dists
-      uses: actions/download-artifact@v3
+      uses: actions/download-artifact@v4
       with:
         name: python-package-distributions
         path: dist/
+
     - name: Sign dists with Sigstore
       uses: sigstore/gh-action-sigstore-python@v2.1.1
       with:
         inputs: >-
           ./dist/*.tar.gz
           ./dist/*.whl
+
     - name: Create GitHub Release
       env:
         GITHUB_TOKEN: ${{ github.token }}
       run: >-
         gh release create
-        'v0.1.10'
+        'v0.2.0'
         --repo '${{ github.repository }}'
         --notes ""
+
     - name: Upload artifact signatures to GitHub Release
       env:
         GITHUB_TOKEN: ${{ github.token }}
@@ -90,5 +102,5 @@ jobs:
       # sigstore-produced signatures and certificates.
       run: >-
         gh release upload
-        'v0.1.9' dist/**
-        --repo '${{ github.repository }}'
+        'v0.2.0' dist/**
+        --repo '${{ github.repository }}'

weco-0.2.0/.gitignore

@@ -0,0 +1,71 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+.pytest_cache/
+.coverage
+htmlcov/
+.env
+.venv
+venv/
+ENV/
+
+# VSCode Extension
+node_modules/
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+.npm
+*.tsbuildinfo
+.eslintcache
+.next
+out/
+build/
+dist/
+
+# IDEs and editors
+.idea/
+.vscode/
+*.swp
+*.swo
+.project
+.classpath
+.settings/
+*.sublime-workspace
+
+# OS
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+
+# Linting
+.ruff_cache/
+
+# Miscellaneous
+etc/
+
+# AI generated files
+digest.txt
+.runs/

{weco-0.1.10 → weco-0.2.0}/LICENSE

@@ -1,6 +1,7 @@
+
 MIT License
 
-Copyright (c) 2024 Weco AI
+Copyright (c) 2025 Weco AI
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

weco-0.2.0/PKG-INFO

@@ -0,0 +1,129 @@
+Metadata-Version: 2.4
+Name: weco
+Version: 0.2.0
+Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
+Author-email: Weco AI Team <dhruv@weco.ai>
+License: MIT
+Project-URL: Homepage, https://github.com/WecoAI/weco-cli
+Keywords: AI,Code Optimization,Code Generation
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: License :: OSI Approved :: MIT License
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: requests
+Requires-Dist: rich
+Provides-Extra: dev
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: setuptools_scm; extra == "dev"
+Dynamic: license-file
+
+# Weco CLI – Optimize Your Code Effortlessly
+
+[![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.
+
+---
+
+## 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.
+
+---
+
+## Setup
+
+1. **Install the Package:**
+
+   ```bash
+   pip install weco
+   ```
+
+2. **Configure API Keys:**
+
+   Set the appropriate environment variables for your language model provider:
+   
+   - **OpenAI:** `export OPENAI_API_KEY="your_key_here"`
+   - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"`
+
+---
+
+## Usage
+
+### Command Line Arguments
+
+| Argument                    | Description                                                                                                                                   | Required |
+|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------|----------|
+| `--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      |
+| `--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       |
+
+---
+
+### Example
+
+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" \
+     --maximize true \
+     --steps 15 \
+     --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:
+```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"
+```
+---
+
+## Supported Providers
+
+The CLI supports the following model providers:
+
+- **OpenAI:** Set your API key using `OPENAI_API_KEY`.
+- **Anthropic:** Set your API key using `ANTHROPIC_API_KEY`.
+
+---
+
+## Contributing
+
+We welcome contributions! To get started:
+
+1. **Fork and Clone the Repository:**
+   ```bash
+   git clone https://github.com/WecoAI/weco-cli.git
+   cd weco-cli
+   ```
+
+2. **Install Development Dependencies:**
+   ```bash
+   pip install -e ".[dev]"
+   ```
+
+3. **Create a Feature Branch:**
+   ```bash
+   git checkout -b feature/your-feature-name
+   ```
+
+4. **Make Your Changes:** Ensure your code adheres to our style guidelines and includes relevant tests.
+
+5. **Commit and Push** your changes, then open a pull request with a clear description of your enhancements.
+
+---

weco-0.2.0/README.md

@@ -0,0 +1,107 @@
+# Weco CLI – Optimize Your Code Effortlessly
+
+[![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.
+
+---
+
+## 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.
+
+---
+
+## Setup
+
+1. **Install the Package:**
+
+   ```bash
+   pip install weco
+   ```
+
+2. **Configure API Keys:**
+
+   Set the appropriate environment variables for your language model provider:
+   
+   - **OpenAI:** `export OPENAI_API_KEY="your_key_here"`
+   - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"`
+
+---
+
+## Usage
+
+### Command Line Arguments
+
+| Argument                    | Description                                                                                                                                   | Required |
+|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------|----------|
+| `--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      |
+| `--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       |
+
+---
+
+### Example
+
+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" \
+     --maximize true \
+     --steps 15 \
+     --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:
+```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"
+```
+---
+
+## Supported Providers
+
+The CLI supports the following model providers:
+
+- **OpenAI:** Set your API key using `OPENAI_API_KEY`.
+- **Anthropic:** Set your API key using `ANTHROPIC_API_KEY`.
+
+---
+
+## Contributing
+
+We welcome contributions! To get started:
+
+1. **Fork and Clone the Repository:**
+   ```bash
+   git clone https://github.com/WecoAI/weco-cli.git
+   cd weco-cli
+   ```
+
+2. **Install Development Dependencies:**
+   ```bash
+   pip install -e ".[dev]"
+   ```
+
+3. **Create a Feature Branch:**
+   ```bash
+   git checkout -b feature/your-feature-name
+   ```
+
+4. **Make Your Changes:** Ensure your code adheres to our style guidelines and includes relevant tests.
+
+5. **Commit and Push** your changes, then open a pull request with a clear description of your enhancements.
+
+---

weco-0.2.0/examples/simple-mlx/evaluate.py

@@ -0,0 +1,133 @@
+import time
+import sys
+import pathlib
+import importlib
+import traceback
+import mlx.core as mx
+import mlx.nn as nn
+
+
+########################################################
+# Baseline
+########################################################
+class Model(nn.Module):
+    """
+    Model that performs a matrix multiplication, division, summation, and scaling.
+    """
+
+    def __init__(self, input_size, hidden_size, scaling_factor):
+        super(Model, self).__init__()
+        self.weight = mx.random.normal(shape=(hidden_size, input_size))
+        self.scaling_factor = scaling_factor
+
+    def __call__(self, x):
+        """
+        Args:
+            x (mx.array): Input tensor of shape (batch_size, input_size).
+        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
+        return x
+
+
+########################################################
+# Weco Solution
+########################################################
+def load_module_from_path(module_path: str, add_to_sys_modules: bool = False):
+    # Clean out all old compiled extensions to prevent namespace collisions during build
+    module_path = pathlib.Path(module_path)
+    name = module_path.stem
+    spec = importlib.util.spec_from_file_location(name, module_path)
+    mod = importlib.util.module_from_spec(spec)  # type: ignore
+    if add_to_sys_modules:
+        sys.modules[name] = mod
+    spec.loader.exec_module(mod)  # type: ignore
+    return mod
+
+
+########################################################
+# Benchmark
+########################################################
+def get_inputs(B, N):
+    # MLX doesn't use device parameter like PyTorch, as it automatically uses Metal
+    return mx.random.normal(shape=(B, N), dtype=mx.float32)
+
+
+def bench(f, inputs, n_warmup, n_rep):
+    # Warm up
+    for _ in range(n_warmup):
+        result = f(inputs)
+        mx.eval(result)  # Force computation due to lazy evaluation
+
+    t_avg = 0.0
+    for _ in range(n_rep):
+        # Clear cache before timing
+        mx.metal.clear_cache()
+
+        start_time = time.time()
+        result = f(inputs)
+        mx.eval(result)  # Force computation
+        mx.synchronize()  # Wait for all computations to complete
+        t_avg += time.time() - start_time
+
+    t_avg /= n_rep * 1e-3
+    return t_avg
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--solution-path", type=str, required=True)
+    args = parser.parse_args()
+
+    # init and input parameters
+    B, N, H, S = 128, 10, 20, 1.5
+
+    # Set the default device to 0
+    mx.set_default_device(mx.gpu)
+
+    # load solution module
+    try:
+        mx.random.seed(0)
+        solution_module = load_module_from_path(args.solution_path, add_to_sys_modules=False)
+        solution_model = solution_module.Model(N, H, S)
+        assert hasattr(solution_model, "__call__")
+    except Exception:
+        print(f"Candidate module initialization failed: {traceback.format_exc()}")
+        exit(1)
+
+    mx.random.seed(0)
+    baseline_model = Model(N, H, S)
+
+    # measure correctness
+    n_correctness_trials = 10
+    max_diff_avg = 0
+    for _ in range(n_correctness_trials):
+        inputs = get_inputs(B, N)
+        baseline_output = baseline_model(inputs)
+        optimized_output = solution_model(inputs)
+        max_diff = mx.max(mx.abs(optimized_output - baseline_output))
+        mx.eval(max_diff)  # Force computation
+        max_diff_avg += max_diff.item()  # Convert to Python scalar
+    max_diff_avg /= n_correctness_trials
+    print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
+
+    # measure performance
+    inputs = get_inputs(B, N)
+    n_warmup = 100
+    n_rep = 500
+
+    # baseline
+    t_avg_baseline = bench(baseline_model, inputs, n_warmup, n_rep)
+    print(f"baseline time: {t_avg_baseline:.2f}ms")
+
+    # optimized
+    t_avg_optimized = bench(solution_model, inputs, n_warmup, n_rep)
+    print(f"optimized time: {t_avg_optimized:.2f}ms")
+
+    print(f"speedup: {t_avg_baseline / t_avg_optimized:.2f}x")