Release 2.0.4-beta.9: PyPI-ready with tiktoken workaround

Changes: - Bundle tiktoken assets locally (mlxk2/assets/whisper/*.tiktoken) - Monkey-patch mlx-audio to use bundled assets at import time - Update pyproject.toml: mlx-audio>=0.3.1 from PyPI (no Git install) - Simplify README: 3 clear installation methods - Add NOTICE file for tiktoken asset attribution Fixes: mlx-audio Issue #479 (missing tiktoken assets in PyPI wheels) Workaround: Temporary until upstream fix available
2026-07-18 12:46:54 -04:00 · 2026-02-04 14:41:52 +01:00
parent bf7480d042
commit 69a3c19c16
4 changed files with 117 additions and 48 deletions
@@ -85,69 +85,38 @@ This license applies **only** to the `mlx-knife` code and **does not extend** to

 ## Installation

-### Via PyPI (Stable - v2.0.3, Text only)
+### 1. PyPI Stable (2.0.3 - Text models only)

 ```bash
 pip install mlx-knife
-
-# Verify
 mlxk --version  # → mlxk 2.0.3
 ```

-**Requirements:**
- macOS with Apple Silicon (M1/M2/M3/M4)
- Python 3.10-3.12
+**Requirements:** macOS Apple Silicon, Python 3.9-3.12

-> **Note:** PyPI stable (2.0.3) supports **Text models only**. For Vision + Audio, use the beta version below.
-
-### Via GitHub (Beta - v2.0.4-beta.9, Text + Vision + Audio)
+### 2. PyPI Beta (2.0.4-beta.9 - Text + Vision + Audio)

 ```bash
-# Step 1: Base + Vision
-pip install "git+https://github.com/mzau/mlx-knife.git@v2.0.4-beta.9#egg=mlx-knife[vision]"
-
-# Step 2: Audio (optional - requires Git install due to PyPI regression)
-pip install -e "git+https://github.com/Blaizzy/mlx-audio.git@9349644#egg=mlx-audio"
-pip install tiktoken
-
-# Verify
+pip install mlx-knife[all]==2.0.4b9
 mlxk --version  # → mlxk 2.0.4b9
 ```

-**Beta.9 features:**
- **Vision**: mlx-vlm 0.3.10 - Image analysis with EXIF metadata
- **Audio**: mlx-audio (Git) - Whisper STT (WAV, MP3, M4A)
- **Recommended models**: `whisper-large-v3-turbo-4bit`, `pixtral-12b-4bit`
+**Requirements:** macOS Apple Silicon, Python 3.10-3.12
+**Features:** Audio STT (Whisper), Vision with EXIF metadata, tiktoken workaround bundled

-> **⚠️ Audio Installation Note:** mlx-audio 0.3.1 (PyPI) has a tiktoken regression. For audio support, install manually:
-> ```bash
-> pip install -e "git+https://github.com/Blaizzy/mlx-audio.git@9349644#egg=mlx-audio"
-> pip install tiktoken
-> ```
-
-### Development Installation
+### 3. Developer Installation

 ```bash
-# Clone and install from source
 git clone https://github.com/mzau/mlx-knife.git
 cd mlx-knife
+pip install -e ".[all,dev,test]"

-# Step 1: Base + Vision + Dev tools
-pip install -e ".[vision,dev,test]"
-
-# Step 2: Audio from Git (PyPI 0.3.1 broken)
-pip install -e "git+https://github.com/Blaizzy/mlx-audio.git@9349644#egg=mlx-audio"
-pip install tiktoken
-
-# Verify
 mlxk --version  # → mlxk 2.0.4b9
-python -c "import mlx_vlm; print('vision ok')"
-python -c "import mlx_audio; print('audio ok')"
-
-# Run tests
 pytest -v
 ```

+**Requirements:** macOS Apple Silicon, Python 3.10-3.12
+
 ### Migrating from 1.x

 If you're upgrading from MLX Knife 1.x, see [MIGRATION.md](MIGRATION.md) for important information about the license change (MIT → Apache 2.0) and behavior changes.
@@ -22,3 +22,9 @@ This product includes software developed by:
  The embedded library is dynamically loaded via CFFI (permitted under LGPL §6).
  MP3 codec support is provided by the embedded libsndfile without additional
  system dependencies (no ffmpeg or Homebrew required).
+
+  mlx-knife bundles tiktoken vocabulary files (gpt2.tiktoken, multilingual.tiktoken)
+  in mlxk2/assets/whisper/ from mlx-audio commit 9349644. These files are used as
+  a workaround for mlx-audio Issue #479 (assets removed in f7328a4 but code still
+  requires them). Original files are from the OpenAI Whisper model, distributed by
+  mlx-audio under the MIT License.
@@ -17,6 +17,94 @@ from typing import Dict, List, Optional, Sequence, Tuple
 from ..operations.workspace import is_workspace_path


+# ============================================================================
+# CRITICAL: Monkey-patch mlx-audio tokenizer BEFORE any imports
+# ============================================================================
+# Workaround for mlx-audio Issue #479: tiktoken assets were removed in commit
+# f7328a4 (Jan 29, 2026), but code still tries to load them.
+#
+# We bundle the assets from commit 9349644 in mlxk2/assets/whisper/ and patch
+# get_encoding() to use them instead.
+#
+# This MUST happen at module import time, before any mlx-audio code runs!
+# ============================================================================
+
+def _apply_tiktoken_patch():
+    """Apply tiktoken asset patch globally at module import time."""
+    try:
+        import base64
+        import tiktoken
+        from functools import lru_cache
+
+        # Import tokenizer module first (but don't trigger get_encoding yet)
+        import mlx_audio.stt.models.whisper.tokenizer as whisper_tokenizer
+
+        # Get path to our bundled tiktoken assets
+        assets_dir = Path(__file__).parent.parent / "assets" / "whisper"
+        if not assets_dir.exists():
+            # Assets not found - skip patching (will fall back to HF WhisperProcessor)
+            return
+
+        @lru_cache(maxsize=None)
+        def patched_get_encoding(name: str = "gpt2", num_languages: int = 99):
+            """Patched get_encoding using mlxk2's bundled tiktoken files."""
+            vocab_path = assets_dir / f"{name}.tiktoken"
+
+            if not vocab_path.exists():
+                raise FileNotFoundError(
+                    f"Tiktoken vocabulary file not found: {vocab_path}\n"
+                    f"mlx-audio Issue #479: assets were removed in f7328a4"
+                )
+
+            with open(vocab_path) as fid:
+                ranks = {
+                    base64.b64decode(token): int(rank)
+                    for token, rank in (line.split() for line in fid if line)
+                }
+
+            n_vocab = len(ranks)
+            special_tokens = {}
+
+            # Build special tokens (from mlx-audio tokenizer.py:343-358)
+            from mlx_audio.stt.models.whisper.tokenizer import LANGUAGES
+
+            specials = [
+                "<|endoftext|>",
+                "<|startoftranscript|>",
+                *[f"<|{lang}|>" for lang in list(LANGUAGES.keys())[:num_languages]],
+                "<|translate|>",
+                "<|transcribe|>",
+                "<|startoflm|>",
+                "<|startofprev|>",
+                "<|nospeech|>",
+                "<|notimestamps|>",
+                *[f"<|{i * 0.02:.2f}|>" for i in range(1501)],
+            ]
+
+            for token in specials:
+                special_tokens[token] = n_vocab
+                n_vocab += 1
+
+            return tiktoken.Encoding(
+                name=vocab_path.name,
+                explicit_n_vocab=n_vocab,
+                pat_str=r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""",
+                mergeable_ranks=ranks,
+                special_tokens=special_tokens,
+            )
+
+        # Patch the module globally
+        whisper_tokenizer.get_encoding = patched_get_encoding
+
+    except ImportError:
+        # mlx-audio not installed - skip patching
+        pass
+
+
+# Apply patch immediately at module import
+_apply_tiktoken_patch()
+
+
 class AudioRunner:
    """Wrapper around mlx-audio STT API for dedicated transcription models.

@@ -79,6 +167,7 @@ class AudioRunner:
        """Internal model loading - called with progress bars suppressed."""
        try:
            # Import mlx-audio STT module (0.3.0 API)
+            # Note: tiktoken patch was already applied at module import time
            from mlx_audio.stt import load_model
            from mlx_audio.stt.generate import generate_transcription
        except ImportError as e:
@@ -69,16 +69,16 @@ vision = [
    "mlx-vlm>=0.3.10",  # Vision support (ADR-012)
 ]
 audio = [
-    # mlx-audio 0.3.1 has tiktoken fallback regression - use post-0.3.1 commit
-    # See: https://github.com/Blaizzy/mlx-audio/issues/445
-    "mlx-audio @ git+https://github.com/Blaizzy/mlx-audio.git@9349644",
-    "tiktoken>=0.7.0",  # Required by mlx-audio Whisper (not declared as transitive dep)
+    # mlx-audio 0.3.1+ (tiktoken assets workaround bundled in mlxk2/assets/whisper/)
+    # See Issue #479: https://github.com/Blaizzy/mlx-audio/issues/479
+    "mlx-audio>=0.3.1",
+    "tiktoken>=0.7.0",  # Required by bundled tiktoken assets
 ]
 all = [
    "mlx-vlm>=0.3.10",
-    # mlx-audio 0.3.1 has tiktoken fallback regression - use post-0.3.1 commit
-    "mlx-audio @ git+https://github.com/Blaizzy/mlx-audio.git@9349644",
-    "tiktoken>=0.7.0",  # Required by mlx-audio Whisper (not declared as transitive dep)
+    # mlx-audio 0.3.1+ (tiktoken assets workaround bundled in mlxk2/assets/whisper/)
+    "mlx-audio>=0.3.1",
+    "tiktoken>=0.7.0",  # Required by bundled tiktoken assets
 ]

 [tool.setuptools]
@@ -86,3 +86,8 @@ license-files = [
  "LICENSE",
  "mlxk2/NOTICE",
 ]
+
+[tool.setuptools.package-data]
+mlxk2 = [
+  "assets/whisper/*.tiktoken",
+]