chore: serialization methods for LocalWhisperTranscriber

haystack/preview/components/audio/whisper_local.py

@@ -6,7 +6,7 @@
 import torch
 import whisper
 
-from haystack.preview import component, Document
+from haystack.preview import component, Document, default_to_dict, default_from_dict
 
 
 logger = logging.getLogger(__name__)
@@ -59,14 +59,16 @@ def to_dict(self) -> Dict[str, Any]:
         """
         Serialize this component to a dictionary.
         """
-        # return default_to_dict(self, model_name_or_path=self.model_name, device=self.device, whisper_params=self.whisper_params)
+        return default_to_dict(
+            self, model_name_or_path=self.model_name, device=str(self.device), whisper_params=self.whisper_params
+        )
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "LocalWhisperTranscriber":
         """
         Deserialize this component from a dictionary.
         """
-        # return default_from_dict(cls, data)
+        return default_from_dict(cls, data)
 
     @component.output_types(documents=List[Document])
     def run(self, audio_files: List[Path], whisper_params: Optional[Dict[str, Any]] = None):

test/preview/components/audio/test_whisper_local.py

@@ -26,6 +26,47 @@ def test_init_wrong_model(self):
         with pytest.raises(ValueError, match="Model name 'whisper-1' not recognized"):
             LocalWhisperTranscriber(model_name_or_path="whisper-1")
 
+    @pytest.mark.unit
+    def test_to_dict(self):
+        transcriber = LocalWhisperTranscriber()
+        data = transcriber.to_dict()
+        assert data == {
+            "type": "LocalWhisperTranscriber",
+            "init_parameters": {"model_name_or_path": "large", "device": "cpu", "whisper_params": {}},
+        }
+
+    @pytest.mark.unit
+    def test_to_dict_with_custom_init_parameters(self):
+        transcriber = LocalWhisperTranscriber(
+            model_name_or_path="tiny",
+            device="cuda",
+            whisper_params={"return_segments": True, "temperature": [0.1, 0.6, 0.8]},
+        )
+        data = transcriber.to_dict()
+        assert data == {
+            "type": "LocalWhisperTranscriber",
+            "init_parameters": {
+                "model_name_or_path": "tiny",
+                "device": "cuda",
+                "whisper_params": {"return_segments": True, "temperature": [0.1, 0.6, 0.8]},
+            },
+        }
+
+    @pytest.mark.unit
+    def test_from_dict(self):
+        data = {
+            "type": "LocalWhisperTranscriber",
+            "init_parameters": {
+                "model_name_or_path": "tiny",
+                "device": "cuda",
+                "whisper_params": {"return_segments": True, "temperature": [0.1, 0.6, 0.8]},
+            },
+        }
+        transcriber = LocalWhisperTranscriber.from_dict(data)
+        assert transcriber.model_name == "tiny"
+        assert transcriber.device == torch.device("cuda")
+        assert transcriber.whisper_params == {"return_segments": True, "temperature": [0.1, 0.6, 0.8]}
+
     @pytest.mark.unit
     def test_warmup(self):
         with patch("haystack.preview.components.audio.whisper_local.whisper") as mocked_whisper: