From 08d09eb6666607db2bd508552bfade0ae8969044 Mon Sep 17 00:00:00 2001
From: mmmkkaaayy <18418787+mmmkkaaayy@users.noreply.github.com>
Date: Sat, 18 Nov 2023 14:52:50 -0800
Subject: [PATCH] Enable whisper test in CI for more backends (#2355)

---
 .github/workflows/test.yml  | 21 +++++++++++++++++++++
 test/models/test_whisper.py |  3 ++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 27977099d8..5d0c2baee3 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -81,6 +81,13 @@ jobs:
       with:
         path: ${{ env.Python3_ROOT_DIR }}/lib/python3.8/site-packages
         key: testing-packages-${{ hashFiles('**/setup.py') }}
+    - name: Cache model weights
+      uses: actions/cache@v3
+      with:
+        path: |
+          weights/whisper-tiny.en.pt
+          weights/gpt2.tiktoken
+        key: model-weights-v1
     - name: Install Dependencies
       run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
     - name: Run Pytest
@@ -109,6 +116,13 @@ jobs:
       with:
         path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
         key: testing-packages-${{ hashFiles('**/setup.py') }}
+    - name: Cache model weights
+      uses: actions/cache@v3
+      with:
+        path: |
+          weights/whisper-tiny.en.pt
+          weights/gpt2.tiktoken
+        key: model-weights-v1
     - name: Install Dependencies
       run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
     - name: Run Pytest
@@ -285,6 +299,13 @@ jobs:
         with:
           path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
           key: ${{ matrix.backend }}-packages-${{ hashFiles('**/setup.py') }}
+      - name: Cache model weights
+        uses: actions/cache@v3
+        with:
+          path: |
+            weights/whisper-tiny.en.pt
+            weights/gpt2.tiktoken
+          key: model-weights-v1
       - name: Set env
         run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'clang' && 'CLANG=1' || matrix.backend == 'gpu' && 'GPU=1' || matrix.backend == 'cuda' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nCUDACPU=1\n' || matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nCUDACPU=1\nPTX=1' || matrix.backend == 'triton' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nCUDACPU=1\nTRITON=1\nTRITON_PTXAS_PATH=/usr/bin/ptxas'}}" >> $GITHUB_ENV
       - name: Install OpenCL
diff --git a/test/models/test_whisper.py b/test/models/test_whisper.py
index 835c238fe4..70bb8aa67a 100644
--- a/test/models/test_whisper.py
+++ b/test/models/test_whisper.py
@@ -1,9 +1,10 @@
 import unittest
 import pathlib
 from examples.whisper import init_whisper, load_file_waveform, transcribe_file, transcribe_waveform
+from tinygrad.helpers import CI
 from tinygrad.ops import Device
 
-@unittest.skipUnless(Device.DEFAULT == "METAL", "Some non-metal backends spend too long trying to allocate a 20GB array")
+@unittest.skipIf(CI and Device.DEFAULT in ["LLVM", "CLANG", "CPU"], "Not working on LLVM, slow on others")
 class TestWhisper(unittest.TestCase):
   @classmethod
   def setUpClass(cls):