aadnk commited on
Commit
67b8308
1 Parent(s): 4a46425

[Faster whisper] Add support for large-v1

Browse files
src/whisper/fasterWhisperContainer.py CHANGED
@@ -39,16 +39,21 @@ class FasterWhisperContainer(AbstractWhisperContainer):
39
  def _create_model(self):
40
  print("Loading faster whisper model " + self.model_name + " for device " + str(self.device))
41
  model_config = self._get_model_config()
42
-
43
- if model_config.type == "whisper" and model_config.url not in ["tiny", "base", "small", "medium", "large", "large-v2"]:
44
- raise Exception("FasterWhisperContainer does not yet support Whisper models. Use ct2-transformers-converter to convert the model to a faster-whisper model.")
 
 
 
 
 
45
 
46
  device = self.device
47
 
48
  if (device is None):
49
  device = "auto"
50
 
51
- model = WhisperModel(model_config.url, device=device, compute_type=self.compute_type)
52
  return model
53
 
54
  def create_callback(self, language: str = None, task: str = None, initial_prompt: str = None,
 
39
  def _create_model(self):
40
  print("Loading faster whisper model " + self.model_name + " for device " + str(self.device))
41
  model_config = self._get_model_config()
42
+ model_url = model_config.url
43
+
44
+ if model_config.type == "whisper":
45
+ if model_url not in ["tiny", "base", "small", "medium", "large", "large-v1", "large-v2"]:
46
+ raise Exception("FasterWhisperContainer does not yet support Whisper models. Use ct2-transformers-converter to convert the model to a faster-whisper model.")
47
+ if model_url == "large":
48
+ # large is an alias for large-v1
49
+ model_url = "large-v1"
50
 
51
  device = self.device
52
 
53
  if (device is None):
54
  device = "auto"
55
 
56
+ model = WhisperModel(model_url, device=device, compute_type=self.compute_type)
57
  return model
58
 
59
  def create_callback(self, language: str = None, task: str = None, initial_prompt: str = None,