Pin MedASR to Dec 22 revision + transformers==4.47.1

backend/models/medasr.py

@@ -7,12 +7,9 @@ from pathlib import Path
 
 import librosa
 try:
-    import torch
-    from transformers import AutoProcessor, AutoModelForCTC
-except ModuleNotFoundError:
-    torch = None
-    AutoProcessor = None
-    AutoModelForCTC = None
+    from transformers import pipeline
+except ModuleNotFoundError:  # pragma: no cover - mock mode support
+    pipeline = None
 
 from backend.config import get_settings
 from backend.errors import ModelExecutionError
@@ -21,54 +18,90 @@ from backend.schemas import Transcript
 
 
 class MedASRModel:
-    """Load and run MedASR speech recognition or a deterministic mock implementation."""
+    """Load and run MedASR speech recognition or a deterministic mock implementation.
+
+    Args:
+        model_manager (ModelManager | None): Optional shared model registry manager.
+
+    Returns:
+        None: Initialised model wrapper with lazy-loaded pipeline.
+    """
 
     def __init__(self, model_manager: ModelManager | None = None) -> None:
+        """Initialise MedASR wrapper.
+
+        Args:
+            model_manager (ModelManager | None): Optional model manager instance.
+
+        Returns:
+            None: Sets internal settings and model state.
+        """
         self.settings = get_settings()
         self.model_manager = model_manager or ModelManager()
-        self._model = None
-        self._processor = None
-        self._device = "cpu"
+        self._pipeline = None
 
     @property
     def is_mock_mode(self) -> bool:
+        """Return whether MedASR should operate in deterministic mock mode.
+
+        Args:
+            None: Reads current settings.
+
+        Returns:
+            bool: True when configured model id is "mock".
+        """
         return self.settings.MEDASR_MODEL_ID.lower() == "mock"
 
     def load_model(self) -> None:
+        """Load the MedASR transformer pipeline unless running in mock mode.
+
+        Args:
+            None: Uses settings for model id and device selection.
+
+        Returns:
+            None: Caches loaded pipeline instance.
+        """
         if self.is_mock_mode:
-            self._model = "mock"
-            self.model_manager.register_model("medasr", self._model)
+            self._pipeline = "mock"
+            self.model_manager.register_model("medasr", self._pipeline)
             return
 
-        if self._model is not None:
+        if self._pipeline is not None:
             return
 
-        if AutoModelForCTC is None:
+        if pipeline is None:
             raise ModelExecutionError("transformers is required for non-mock MedASR mode")
 
         device = "cuda:0"
         if self.model_manager.check_gpu()["vram_total_bytes"] == 0:
             device = "cpu"
 
-        model_id = self.settings.MEDASR_MODEL_ID
-
         try:
-            self._processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-            self._model = AutoModelForCTC.from_pretrained(model_id, trust_remote_code=True)
-            self._model = self._model.to(device)
-            self._model.eval()
-            self._device = device
+            self._pipeline = pipeline(
+                "automatic-speech-recognition",
+                model=self.settings.MEDASR_MODEL_ID,
+                revision="2625be4f1377ac544b451c6938eaf955c19a9c38",
+                device=device,
+            )
         except Exception as exc:
             raise ModelExecutionError(f"Failed to load MedASR model: {exc}") from exc
 
-        self.model_manager.register_model("medasr", self._model)
+        self.model_manager.register_model("medasr", self._pipeline)
 
     def transcribe(self, audio_path: str) -> Transcript:
+        """Transcribe audio input into a Transcript schema object.
+
+        Args:
+            audio_path (str): Path to 16kHz mono audio WAV.
+
+        Returns:
+            Transcript: Structured transcript result.
+        """
         source = Path(audio_path)
         if not source.exists():
             raise ModelExecutionError(f"Audio path not found: {source}")
 
-        if self._model is None:
+        if self._pipeline is None:
             self.load_model()
 
         if self.is_mock_mode:
@@ -80,30 +113,30 @@ class MedASRModel:
         duration_s = float(librosa.get_duration(y=waveform, sr=16000))
 
         try:
-            inputs = self._processor(
+            result = self._pipeline(
                 waveform,
-                sampling_rate=16000,
-                return_tensors="pt",
-                padding=True,
+                chunk_length_s=20,
+                stride_length_s=(4, 2),
+                return_timestamps=True,
+                generate_kwargs={"language": "en", "task": "transcribe"},
             )
-            inputs = inputs.to(self._device)
-
-            with torch.no_grad():
-                outputs = self._model.generate(**inputs)
-                transcript_text = self._processor.batch_decode(outputs, skip_special_tokens=True)[0]
-
-            # Clean up special tokens that may remain
-            import re
-            transcript_text = transcript_text.replace("<epsilon>", "")
-            transcript_text = transcript_text.replace("</s>", "").replace("<s>", "")
-            transcript_text = re.sub(r'\s+', ' ', transcript_text).strip()
-
         except Exception as exc:
             raise ModelExecutionError(f"MedASR inference failed: {exc}") from exc
 
+        transcript_text = str(result.get("text", "")).strip()
         return self._make_transcript(source, transcript_text, duration_s)
 
     def _make_transcript(self, audio_path: Path, text: str, duration_s: float) -> Transcript:
+        """Build a Transcript object from model output values.
+
+        Args:
+            audio_path (Path): Source audio path.
+            text (str): Transcript text.
+            duration_s (float): Audio duration seconds.
+
+        Returns:
+            Transcript: Pydantic transcript model.
+        """
         now = datetime.now(tz=timezone.utc).isoformat()
         consultation_id = audio_path.stem
         return Transcript(
@@ -116,20 +149,33 @@ class MedASRModel:
 
     @staticmethod
     def _duration(audio_path: Path) -> float:
+        """Compute audio duration in seconds using librosa.
+
+        Args:
+            audio_path (Path): Audio file path.
+
+        Returns:
+            float: Duration in seconds.
+        """
         waveform, sample_rate = librosa.load(audio_path, sr=16000, mono=True)
         return float(librosa.get_duration(y=waveform, sr=sample_rate))
 
     @staticmethod
     def _get_mock_text(audio_path: Path) -> str:
+        """Return ground-truth transcript for known demo files in mock mode.
+
+        Args:
+            audio_path (Path): Audio file path used for lookup.
+
+        Returns:
+            str: Transcript text from fixture file or fallback placeholder.
+        """
         transcript_map = {
             "mrs_thompson": Path("data/demo/mrs_thompson_transcript.txt"),
             "mr_okafor": Path("data/demo/mr_okafor_transcript.txt"),
             "ms_patel": Path("data/demo/ms_patel_transcript.txt"),
-            "mr_williams": Path("data/demo/mr_williams_transcript.txt"),
-            "mrs_khan": Path("data/demo/mrs_khan_transcript.txt"),
         }
         for key, transcript_path in transcript_map.items():
             if key in audio_path.stem:
-                if transcript_path.exists():
-                    return transcript_path.read_text(encoding="utf-8").strip()
+                return transcript_path.read_text(encoding="utf-8").strip()
         return "Mock transcript placeholder for non-demo audio input."

requirements.txt

@@ -1,5 +1,5 @@
 torch==2.4.1
-transformers @ git+https://github.com/huggingface/transformers.git@65dc261512cbdb1ee72b88ae5b222f2605aad8e5
+transformers==4.47.1
 bitsandbytes>=0.46.1
 accelerate>=1.2.1
 gradio>=5.10.0