From 426d0df672e16de4af4e78d1344105a3cc799e03 Mon Sep 17 00:00:00 2001
From: Roy Shilkrot <roy.shil@gmail.com>
Date: Fri, 18 Oct 2024 09:57:23 -0400
Subject: [PATCH] Add build and test workflow, Python dependencies, and build
 configuration options

---
 .github/workflows/build.yaml | 69 ++++++++++++++++++++++++++++++++++++
 README.md                    | 26 ++++++++++++++
 setup.py                     | 12 ++++++-
 3 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 .github/workflows/build.yaml

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
new file mode 100644
index 0000000..8302d55
--- /dev/null
+++ b/.github/workflows/build.yaml
@@ -0,0 +1,69 @@
+name: Build and Test
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          # Windows builds
+          - os: windows-latest
+            platform: win64
+            acceleration: cpu
+          - os: windows-latest
+            platform: win64
+            acceleration: cuda
+          - os: windows-latest
+            platform: win64
+            acceleration: hipblas
+          - os: windows-latest
+            platform: win64
+            acceleration: vulkan
+          
+          # macOS builds
+          - os: macos-latest
+            platform: x86_64
+          - os: macos-latest
+            platform: arm64
+          
+          # Linux build
+          - os: ubuntu-latest
+            platform: x86_64
+
+    runs-on: ${{ matrix.os }}
+
+    steps:
+    - uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install numpy cmake wheel setuptools
+
+    - name: Install package
+      env:
+        SIMPLER_WHISPER_ACCELERATION: ${{ matrix.acceleration }}
+        SIMPLER_WHISPER_PLATFORM: ${{ matrix.platform }}
+      run: |
+        pip install .
+
+    - name: Test import
+      run: |
+        python -c "import simpler_whisper; print(simpler_whisper.__file__)"
+
+    - name: Upload wheel
+      uses: actions/upload-artifact@v2
+      with:
+        name: wheel-${{ matrix.os }}-${{ matrix.platform }}-${{ matrix.acceleration }}
+        path: dist/*.whl
\ No newline at end of file
diff --git a/README.md b/README.md
index beb6496..1b92bd6 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,7 @@
 # Simpler Whisper
 
+![Build and Test](https://github.com/locaal-ai/simpler-whisper/workflows/Build%20and%20Test/badge.svg)
+
 A simple Python wrapper for whisper.cpp, providing an easy-to-use interface for speech recognition using the Whisper model. This package uses a CMake-based build process to create a Python extension that interfaces with the whisper.cpp library, supporting static libraries on Mac and Linux, and dynamic libraries on Windows.
 
 ## Installation
@@ -53,6 +55,30 @@ If you're building from source:
 
 This will run the CMake build process and compile the extension.
 
+## Build Configuration
+
+Simpler Whisper supports various build configurations to optimize for different hardware and acceleration methods. You can specify the build configuration using environment variables:
+
+- `SIMPLER_WHISPER_ACCELERATION`: Specifies the acceleration method. Options are:
+  - `cpu` (default)
+  - `cuda` (for NVIDIA GPUs)
+  - `hipblas` (for AMD GPUs)
+  - `vulkan` (for cross-platform GPU acceleration)
+
+- `SIMPLER_WHISPER_PLATFORM`: Specifies the target platform. This is mainly used for macOS builds to differentiate between x86_64 and arm64 architectures.
+
+### Example: Building with CUDA acceleration
+
+```bash
+SIMPLER_WHISPER_ACCELERATION=cuda pip install simpler-whisper
+```
+
+### Example: Building for macOS ARM64
+
+```bash
+SIMPLER_WHISPER_PLATFORM=arm64 pip install simpler-whisper
+```
+
 ## License
 
 This project is licensed under the MIT License - see the LICENSE file for details.
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 35614d8..20f2e32 100644
--- a/setup.py
+++ b/setup.py
@@ -25,6 +25,10 @@ def run(self):
     def build_extension(self, ext):
         extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
         
+        # Get acceleration and platform from environment variables
+        acceleration = os.environ.get('SIMPLER_WHISPER_ACCELERATION', 'cpu')
+        target_platform = os.environ.get('SIMPLER_WHISPER_PLATFORM', platform.machine())
+
         # Correctly identify the Python executable and other Python-related paths
         if platform.system() == "Windows":
             python_executable = os.path.join(sys.prefix, 'python.exe')
@@ -42,9 +46,15 @@ def build_extension(self, ext):
             f'-DPYTHON_EXECUTABLE={python_executable}',
             f'-DPYTHON_INCLUDE_DIR={python_include}',
             f'-DNUMPY_INCLUDE_DIR={numpy_include}',
-            '-DACCELERATION=cpu',
+            f'-DACCELERATION={acceleration}',
         ]
 
+        # Add platform-specific arguments
+        if platform.system() == "Darwin":  # macOS
+            cmake_args.append(f'-DCMAKE_OSX_ARCHITECTURES={target_platform}')
+            # add MACOS_ARCH env variable to specify the target platform
+            os.environ["MACOS_ARCH"] = target_platform
+
         cfg = 'Debug' if self.debug else 'Release'
         build_args = ['--config', cfg]