Granite NAR current-API fix (cleaned-up #45)

CONTEXT.md

@@ -42,7 +42,8 @@ service manager, and supported runtime kinds. Linux x86_64 defaults to Granite
 NAR with systemd; Linux CPU defaults to Granite AR. macOS Apple Silicon defaults
 to MLX Whisper via `mlx-audio` with launchd. Windows defaults to Granite AR
 (`ibm-granite/granite-speech-4.1-2b`) with CUDA when available and Task
-Scheduler for the background service. Granite NAR is not supported on Windows.
+Scheduler for the background service. Granite NAR is also selectable on Windows
+CUDA.
 
 **ASR runtime**: The local speech-to-text execution path for a WAV file. It
 includes audio preparation, backend selection, local model loading, transcript

README.md

@@ -69,7 +69,8 @@ transclip tray
 
 Install a CUDA-enabled PyTorch wheel before prefetching Granite AR models, then
 run `transclip models prefetch --model ibm-granite/granite-speech-4.1-2b`.
-Granite NAR is not supported on Windows.
+Granite NAR (`asr_backend = "granite_nar"`) is also selectable on Windows CUDA;
+Granite AR is the default.
 
 Check readiness and logs:
 
@@ -252,9 +253,11 @@ Supported ASR on Windows:
 
 - `ibm-granite/granite-speech-4.1-2b` (default Granite autoregressive)
 - `ibm-granite/granite-speech-4.1-2b-plus` (speaker/timestamp features)
+- `ibm-granite/granite-speech-4.1-2b-nar` (Granite NAR, lower latency; selectable on CUDA)
 
-Granite Speech 4.1 NAR and ROCm are not supported on Windows. Eval thresholds
-for Windows Granite AR are in `eval/windows/manifest.json` (relaxed vs Linux NAR).
+ROCm is not supported on Windows; Granite Speech 4.1 NAR is selectable on Windows
+CUDA (Granite AR remains the default). Eval thresholds for Windows Granite AR are
+in `eval/windows/manifest.json` (relaxed vs Linux NAR).
 
 ### Intel acceleration (integrated GPU / NPU) via OpenVINO
 
@@ -378,7 +381,7 @@ after a restart does not pay a multi-second ROCm shape compile.
 |----------|---------------|-------------|-------|
 | Linux GPU (CUDA/ROCm) | Granite NAR | Opt-in | No extra dependencies |
 | Linux/Windows CPU | Granite CPU/AR | No | Requires the granite_nar GPU backend |
-| Windows CUDA | Granite AR | No | granite_nar is not supported on Windows |
+| Windows CUDA | Granite AR (NAR selectable) | No | NAR selectable on CUDA; incremental pending validation |
 | macOS MLX | MLX Whisper | Pending benchmark | Run `scripts/bench_nar_mlx.py` on an M-series (gate: warm 8 s pass <= 900 ms) |
 
 While recording, `GET /record/partial` and the tray's **Copy partial

tests/test_asr.py

@@ -14,6 +14,8 @@
     OpenVINOWhisperBackend,
     _configure_rocm_nar_attention_env,
     _granite_nar_dtype,
+    _nar_decode,
+    _nar_infer,
     _pad_nar_waveform_to_bucket,
     _whisper_language_token,
     build_asr_backend,
@@ -28,6 +30,7 @@ class ASRTests(unittest.TestCase):
     @staticmethod
     def _linux_runtime() -> FakeRuntime:
         return FakeRuntime(system="Linux", home=Path("/home/user"))
+
     def test_granite_prompt_requests_punctuation(self):
         self.assertEqual(
             granite_user_prompt(),
@@ -421,6 +424,68 @@ def test_granite_nar_bucket_padding_leaves_exact_bucket_unchanged(self):
 
         self.assertIs(padded, waveform)
 
+    def test_nar_infer_prefers_transcribe_over_generate(self):
+        # Current NAR revisions expose transcribe(); generate() is gone.
+        calls = []
+
+        class Model:
+            def transcribe(self, **kwargs):
+                calls.append("transcribe")
+                return "via-transcribe"
+
+            def generate(self, **kwargs):
+                calls.append("generate")
+                return "via-generate"
+
+        self.assertEqual(_nar_infer(Model(), {"input_features": 1}), "via-transcribe")
+        self.assertEqual(calls, ["transcribe"])
+
+    def test_nar_infer_falls_back_to_generate_for_older_models(self):
+        # Older revisions (e.g. the friend's pinned Linux model) only had generate().
+        class OldModel:
+            def generate(self, **kwargs):
+                return "via-generate"
+
+        self.assertEqual(_nar_infer(OldModel(), {}), "via-generate")
+
+    def test_nar_decode_uses_text_preds_when_present(self):
+        # Older revisions returned already-decoded strings in text_preds.
+        output = SimpleNamespace(text_preds=["hello world"])
+        self.assertEqual(_nar_decode(output, lambda: None), "hello world")
+
+    def test_nar_decode_detokenizes_token_id_preds(self):
+        # Current revisions return token-id tensors in preds needing a tokenizer.
+        class FakeTokenizer:
+            def decode(self, ids, skip_special_tokens=False):
+                assert skip_special_tokens
+                return "decoded:" + ",".join(str(i) for i in ids)
+
+        output = SimpleNamespace(preds=[[1, 2, 3]])
+        self.assertEqual(_nar_decode(output, lambda: FakeTokenizer()), "decoded:1,2,3")
+
+    def test_nar_decode_raises_on_token_ids_without_tokenizer(self):
+        output = SimpleNamespace(preds=[[1, 2, 3]])
+        with self.assertRaisesRegex(RuntimeError, "no tokenizer"):
+            _nar_decode(output, lambda: None)
+
+    def test_nar_decode_prefers_text_preds_when_both_present(self):
+        # Already-decoded text_preds wins over raw token-id preds.
+        class BoomTokenizer:
+            def decode(self, *_a, **_k):
+                raise AssertionError("must not detokenize when text_preds is present")
+
+        output = SimpleNamespace(text_preds=["winner"], preds=[[9, 9]])
+        self.assertEqual(_nar_decode(output, lambda: BoomTokenizer()), "winner")
+
+    def test_nar_decode_falls_through_empty_text_preds_to_token_ids(self):
+        # An empty text_preds must not IndexError; fall through to preds.
+        class FakeTokenizer:
+            def decode(self, ids, skip_special_tokens=False):
+                return "decoded:" + ",".join(str(i) for i in ids)
+
+        output = SimpleNamespace(text_preds=[], preds=[[7, 8]])
+        self.assertEqual(_nar_decode(output, lambda: FakeTokenizer()), "decoded:7,8")
+
     def test_audio_preparer_folds_channels_and_resamples_without_model_runtime(self):
         samples = np.array([[1.0, 3.0], [5.0, 7.0]], dtype=np.float32)
         resample_calls = []

tests/test_models.py

@@ -170,7 +170,8 @@ def test_disk_space_failure_uses_catalog_estimate(self):
             settings = Settings(model_cache_dir=tmp)
             ensure_disk_space(settings, SUPPORTED_MODELS[0])
 
-    def test_windows_catalog_excludes_nar_and_defaults_to_granite_ar(self):
+    def test_windows_catalog_includes_nar_but_defaults_to_granite_ar(self):
+        # NAR is selectable on Windows (CUDA), but AR stays the proven default.
         runtime = FakeRuntime(system="Windows", home=Path("C:/Users/test"))
         with patch("transclip.device.torch_cuda_usable", return_value=True):
             defaults = default_settings(runtime)
@@ -182,21 +183,21 @@ def test_windows_catalog_excludes_nar_and_defaults_to_granite_ar(self):
         self.assertEqual(defaults.asr_model, "ibm-granite/granite-speech-4.1-2b")
         self.assertIn("ibm-granite/granite-speech-4.1-2b", model_ids)
         self.assertIn("ibm-granite/granite-speech-4.1-2b-plus", model_ids)
-        self.assertNotIn("ibm-granite/granite-speech-4.1-2b-nar", model_ids)
+        self.assertIn("ibm-granite/granite-speech-4.1-2b-nar", model_ids)
         default_row = next(row for row in rows if "default" in row.marker)
         self.assertEqual(default_row.backend, "granite")
         self.assertEqual(default_row.model_id, "ibm-granite/granite-speech-4.1-2b")
 
-    def test_windows_rejects_granite_nar_backend(self):
+    def test_windows_accepts_granite_nar_backend(self):
         runtime = FakeRuntime(system="Windows", home=Path("C:/Users/test"))
-        with (
-            patch("transclip.device.torch_cuda_usable", return_value=True),
-            self.assertRaisesRegex(ValueError, "not supported on Windows"),
-        ):
-            validate_asr_model_backend(
+        with patch("transclip.device.torch_cuda_usable", return_value=True):
+            self.assertEqual(
+                validate_asr_model_backend(
+                    "granite_nar",
+                    "ibm-granite/granite-speech-4.1-2b-nar",
+                    runtime,
+                ),
                 "granite_nar",
-                "ibm-granite/granite-speech-4.1-2b-nar",
-                runtime,
             )
 
     def test_cli_models_list_uses_local_catalog(self):

tests/test_platform_runtime.py

@@ -101,9 +101,7 @@ def test_granite_nar_accepted_on_darwin_arm(self):
     def test_granite_nar_rejected_on_darwin_intel(self):
         runtime = FakeRuntime(system="Darwin", home=Path("/Users/test"), check_output_text="x86_64")
         with self.assertRaisesRegex(ValueError, "requires aarch64, arm64"):
-            validate_asr_model_backend(
-                "granite_nar", "ibm-granite/granite-speech-4.1-2b-nar", runtime
-            )
+            validate_asr_model_backend("granite_nar", "ibm-granite/granite-speech-4.1-2b-nar", runtime)
 
     def test_supported_catalog_entries_filter_by_platform(self):
         linux_runtime = FakeRuntime(system="Linux", home=Path("/home/user"))
@@ -168,23 +166,25 @@ def test_windows_cpu_profile_defaults(self):
         self.assertEqual(profile.profile_id, "windows_cpu")
         self.assertEqual(settings.asr_device, "cpu")
 
-    def test_granite_nar_rejected_on_windows(self):
+    def test_granite_nar_accepted_on_windows(self):
         runtime = FakeRuntime(system="Windows", home=Path("C:/Users/tester"))
         with patch("transclip.device.torch_cuda_usable", return_value=True):
             settings = replace(
                 default_settings(runtime),
                 asr_backend="granite_nar",
                 asr_model="ibm-granite/granite-speech-4.1-2b-nar",
             )
-            with self.assertRaisesRegex(ValueError, "not supported on Windows"):
-                validate_asr_model_backend(settings.asr_backend, settings.asr_model, runtime)
+            self.assertEqual(
+                validate_asr_model_backend(settings.asr_backend, settings.asr_model, runtime),
+                "granite_nar",
+            )
 
-    def test_supported_catalog_entries_exclude_nar_on_windows(self):
+    def test_supported_catalog_entries_include_nar_on_windows(self):
         runtime = FakeRuntime(system="Windows", home=Path("C:/Users/tester"))
         model_ids = {entry.model_id for entry in supported_catalog_entries(runtime)}
 
         self.assertIn("ibm-granite/granite-speech-4.1-2b", model_ids)
-        self.assertNotIn("ibm-granite/granite-speech-4.1-2b-nar", model_ids)
+        self.assertIn("ibm-granite/granite-speech-4.1-2b-nar", model_ids)
 
     def test_windows_openvino_profile_selected_when_intel_accelerator_present(self):
         runtime = FakeRuntime(system="Windows", home=Path("C:/Users/tester"))

transclip/asr.py

@@ -5,6 +5,7 @@
 import platform as py_platform
 import tempfile
 import threading
+from collections.abc import Callable
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Protocol
@@ -251,6 +252,7 @@ def __init__(
         self.local_files_only = local_files_only
         self.cache_dir = cache_dir
         self._loaded = None
+        self._tokenizer = None
         self.audio_preparer = TorchAudioPreparer()
 
     def _device(self):
@@ -287,6 +289,22 @@ def _load(self, device: TorchDevice):
         self._loaded = (feature_extractor, model)
         return self._loaded
 
+    def _get_tokenizer(self):
+        # Lazy + cached: only current NAR revisions (which return token-id `preds`)
+        # need a tokenizer. The default revision returns decoded `text_preds`, so
+        # this never runs on the primary Linux gfx1151 path -- keeping its startup
+        # identical to before the version-robust decode was added.
+        if self._tokenizer is None:
+            from transformers import AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                self.model,
+                trust_remote_code=True,
+                local_files_only=self.local_files_only,
+                cache_dir=self.cache_dir or None,
+            )
+        return self._tokenizer
+
     def transcribe(self, wav_path: Path, keywords: list[str] | None = None) -> TranscriptionResult:
         del keywords
         audio = self.audio_preparer.prepare(wav_path)
@@ -305,15 +323,14 @@ def transcribe_waveform(self, waveform: Any, sample_rate: int = 16000) -> Transc
             if sample_rate != GRANITE_NAR_SAMPLE_RATE:
                 import torchaudio
 
-                waveform = torchaudio.functional.resample(
-                    waveform, sample_rate, GRANITE_NAR_SAMPLE_RATE
-                )
+                waveform = torchaudio.functional.resample(waveform, sample_rate, GRANITE_NAR_SAMPLE_RATE)
                 sample_rate = GRANITE_NAR_SAMPLE_RATE
             waveform = _pad_nar_waveform_to_bucket(waveform, sample_rate=sample_rate)
             inputs = feature_extractor([waveform], device=device)
             with torch.inference_mode():
-                output = model.generate(**inputs)
-        return TranscriptionResult(output.text_preds[0].strip(), timings, self.name, self.model)
+                output = _nar_infer(model, inputs)
+            text = _nar_decode(output, self._get_tokenizer)
+        return TranscriptionResult(text.strip(), timings, self.name, self.model)
 
 
 class MlxAudioASRBackend:
@@ -638,6 +655,40 @@ def _configure_rocm_nar_attention_env(os_module, torch, device: TorchDevice) ->
         os_module.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
 
 
+def _nar_infer(model: Any, inputs: Any) -> Any:
+    """Run Granite NAR inference across model revisions.
+
+    The published NAR model code renamed its inference entrypoint from
+    ``generate`` to ``transcribe`` (``can_generate()`` now returns False), so
+    prefer ``transcribe`` and fall back to ``generate`` for older revisions.
+    """
+    runner = getattr(model, "transcribe", None) or getattr(model, "generate", None)
+    if runner is None:
+        raise RuntimeError("Granite NAR model exposes neither transcribe() nor generate()")
+    return runner(**inputs)
+
+
+def _nar_decode(output: Any, load_tokenizer: Callable[[], Any]) -> str:
+    """Decode a Granite NAR output across model revisions.
+
+    Older revisions return decoded strings in ``text_preds`` (preferred when
+    present); newer ones return token-id tensors in ``preds``. ``load_tokenizer``
+    is called only on the ``preds`` path, so the default revision never builds a
+    tokenizer it won't use. Truthiness (not ``is not None``) lets an empty
+    ``text_preds`` fall through instead of raising IndexError.
+    """
+    text_preds = getattr(output, "text_preds", None)
+    if text_preds:
+        return text_preds[0]
+    preds = getattr(output, "preds", None)
+    if preds:
+        tokenizer = load_tokenizer()
+        if tokenizer is None:
+            raise RuntimeError("Granite NAR returned token ids but no tokenizer was loaded to decode them")
+        return tokenizer.decode(preds[0], skip_special_tokens=True)
+    raise RuntimeError(f"Unexpected Granite NAR output without text_preds or preds: {type(output).__name__}")
+
+
 def _linear_resample(samples: Any, source_rate: int, target_rate: int) -> Any:
     if source_rate == target_rate:
         return samples

transclip/models/catalog.py

@@ -19,7 +19,7 @@
         display_name="Fast local ASR - Granite 4.1 NAR",
         runtime_kind="torch",
         estimated_bytes=8 * GIB,
-        supported_platforms=frozenset({"Darwin", "Linux"}),
+        supported_platforms=frozenset({"Darwin", "Linux", "Windows"}),
         supported_architectures=frozenset({"arm64", "aarch64"}),
         dependency_extra="models",
         prefetch_strategy="transformers",
@@ -209,9 +209,7 @@ def catalog_entry_for_backend(backend: str, model_id: str) -> ModelCatalogEntry:
         raise ValueError("file backends do not use the model catalog")
     entry = catalog_entry_for_model(model_id)
     if entry.backend != normalized:
-        raise ValueError(
-            f"Model {model_id} requires asr_backend={entry.backend!r}, not {backend!r}"
-        )
+        raise ValueError(f"Model {model_id} requires asr_backend={entry.backend!r}, not {backend!r}")
     return entry
 
 
@@ -286,11 +284,7 @@ def supported_catalog_entries(runtime: PlatformRuntime | None = None) -> list[Mo
     for entry in MODEL_CATALOG:
         if entry.supported_platforms and profile.system not in entry.supported_platforms:
             continue
-        if (
-            entry.supported_architectures is not None
-            and profile.system == "Darwin"
-            and not is_apple_silicon(runtime)
-        ):
+        if entry.supported_architectures is not None and profile.system == "Darwin" and not is_apple_silicon(runtime):
             continue
         entries.append(entry)
     return entries
@@ -365,5 +359,3 @@ def model_rows(settings: Settings, runtime: PlatformRuntime | None = None) -> li
 
 def asr_model_rows(settings: Settings, runtime: PlatformRuntime | None = None) -> list[ModelRow]:
     return [row for row in model_rows(settings, runtime) if row.backend != "text_generation"]
-
-

transclip/models/prefetch.py

@@ -25,9 +25,7 @@ def prefetch_model(model_id: str, settings: Settings, runtime: PlatformRuntime |
         try:
             handler = _PREFETCH_BY_BACKEND[entry.backend]
         except KeyError as exc:
-            raise RuntimeError(
-                f"No prefetch handler for backend {entry.backend!r} (model {entry.model_id!r})"
-            ) from exc
+            raise RuntimeError(f"No prefetch handler for backend {entry.backend!r} (model {entry.model_id!r})") from exc
         handler(entry, cache_dir)
         return model_cache_path(entry.model_id, settings, runtime)
 
@@ -71,7 +69,7 @@ def _prefetch_text_generation(entry: ModelCatalogEntry, cache_dir: str) -> None:
 
 def _prefetch_granite_nar(entry: ModelCatalogEntry, cache_dir: str) -> None:
     try:
-        from transformers import AutoModel, AutoProcessor
+        from transformers import AutoModel, AutoProcessor, AutoTokenizer
     except ImportError as exc:
         raise RuntimeError("transformers is required. Install transclip[models].") from exc
     AutoModel.from_pretrained(
@@ -86,6 +84,14 @@ def _prefetch_granite_nar(entry: ModelCatalogEntry, cache_dir: str) -> None:
         local_files_only=False,
         cache_dir=cache_dir,
     )
+    # The NAR backend loads a standalone tokenizer with local_files_only=True to
+    # detokenize output.preds; cache it explicitly so a fresh prefetch suffices.
+    AutoTokenizer.from_pretrained(
+        entry.model_id,
+        trust_remote_code=True,
+        local_files_only=False,
+        cache_dir=cache_dir,
+    )
 
 
 def _prefetch_transformers_speech(entry: ModelCatalogEntry, cache_dir: str) -> None: