Spaces:
Running
Running
File size: 9,149 Bytes
95261ed 295de00 7fd072f 33a2c1e 44d964a 295de00 33ee1bb 295de00 33a2c1e 95261ed 33a2c1e 8031785 295de00 95261ed c0e541b 295de00 31f7bdb 295de00 33ee1bb 95261ed 7c5d37e 7fd072f 295de00 7fd072f 7c5d37e 7fd072f 44d964a 7c5d37e 44d964a 7c5d37e 295de00 44d964a c0e541b 295de00 44d964a 295de00 44d964a c0e541b 8031785 95261ed 8031785 95261ed 8031785 95261ed 295de00 95261ed 295de00 95261ed 295de00 95261ed 295de00 8031785 95261ed 8031785 95261ed 33a2c1e 95261ed 295de00 95261ed 33a2c1e 33ee1bb 8031785 95261ed 33a2c1e 8031785 33ee1bb 295de00 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
# External programs
import abc
import os
import sys
from typing import List
from urllib.parse import urlparse
import torch
import urllib3
from src.hooks.progressListener import ProgressListener
import whisper
from whisper import Whisper
from src.config import ModelConfig, VadInitialPromptMode
from src.hooks.whisperProgressHook import create_progress_listener_handle
from src.modelCache import GLOBAL_MODEL_CACHE, ModelCache
from src.utils import download_file
from src.whisper.abstractWhisperContainer import AbstractWhisperCallback, AbstractWhisperContainer
class WhisperContainer(AbstractWhisperContainer):
def __init__(self, model_name: str, device: str = None, compute_type: str = "float16",
download_root: str = None,
cache: ModelCache = None, models: List[ModelConfig] = []):
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
super().__init__(model_name, device, compute_type, download_root, cache, models)
def ensure_downloaded(self):
"""
Ensure that the model is downloaded. This is useful if you want to ensure that the model is downloaded before
passing the container to a subprocess.
"""
# Warning: Using private API here
try:
root_dir = self.download_root
model_config = self._get_model_config()
if root_dir is None:
root_dir = os.path.join(os.path.expanduser("~"), ".cache", "whisper")
if self.model_name in whisper._MODELS:
whisper._download(whisper._MODELS[self.model_name], root_dir, False)
else:
# If the model is not in the official list, see if it needs to be downloaded
model_config.download_url(root_dir)
return True
except Exception as e:
# Given that the API is private, it could change at any time. We don't want to crash the program
print("Error pre-downloading model: " + str(e))
return False
def _get_model_config(self) -> ModelConfig:
"""
Get the model configuration for the model.
"""
for model in self.models:
if model.name == self.model_name:
return model
return None
def _create_model(self):
print("Loading whisper model " + self.model_name)
model_config = self._get_model_config()
# Note that the model will not be downloaded in the case of an official Whisper model
model_path = self._get_model_path(model_config, self.download_root)
return whisper.load_model(model_path, device=self.device, download_root=self.download_root)
def create_callback(self, language: str = None, task: str = None, initial_prompt: str = None,
initial_prompt_mode: VadInitialPromptMode = VadInitialPromptMode.PREPREND_FIRST_SEGMENT,
**decodeOptions: dict) -> AbstractWhisperCallback:
"""
Create a WhisperCallback object that can be used to transcript audio files.
Parameters
----------
language: str
The target language of the transcription. If not specified, the language will be inferred from the audio content.
task: str
The task - either translate or transcribe.
initial_prompt: str
The initial prompt to use for the transcription.
initial_prompt_mode: VadInitialPromptMode
The mode to use for the initial prompt. If set to PREPEND_FIRST_SEGMENT, the initial prompt will be prepended to the first segment of audio.
If set to PREPEND_ALL_SEGMENTS, the initial prompt will be prepended to all segments of audio.
decodeOptions: dict
Additional options to pass to the decoder. Must be pickleable.
Returns
-------
A WhisperCallback object.
"""
return WhisperCallback(self, language=language, task=task, initial_prompt=initial_prompt, initial_prompt_mode=initial_prompt_mode, **decodeOptions)
def _get_model_path(self, model_config: ModelConfig, root_dir: str = None):
from src.conversion.hf_converter import convert_hf_whisper
"""
Download the model.
Parameters
----------
model_config: ModelConfig
The model configuration.
"""
# See if path is already set
if model_config.path is not None:
return model_config.path
if root_dir is None:
root_dir = os.path.join(os.path.expanduser("~"), ".cache", "whisper")
model_type = model_config.type.lower() if model_config.type is not None else "whisper"
if model_type in ["huggingface", "hf"]:
model_config.path = model_config.url
destination_target = os.path.join(root_dir, model_config.name + ".pt")
# Convert from HuggingFace format to Whisper format
if os.path.exists(destination_target):
print(f"File {destination_target} already exists, skipping conversion")
else:
print("Saving HuggingFace model in Whisper format to " + destination_target)
convert_hf_whisper(model_config.url, destination_target)
model_config.path = destination_target
elif model_type in ["whisper", "w"]:
model_config.path = model_config.url
# See if URL is just a file
if model_config.url in whisper._MODELS:
# No need to download anything - Whisper will handle it
model_config.path = model_config.url
elif model_config.url.startswith("file://"):
# Get file path
model_config.path = urlparse(model_config.url).path
# See if it is an URL
elif model_config.url.startswith("http://") or model_config.url.startswith("https://"):
# Extension (or file name)
extension = os.path.splitext(model_config.url)[-1]
download_target = os.path.join(root_dir, model_config.name + extension)
if os.path.exists(download_target) and not os.path.isfile(download_target):
raise RuntimeError(f"{download_target} exists and is not a regular file")
if not os.path.isfile(download_target):
download_file(model_config.url, download_target)
else:
print(f"File {download_target} already exists, skipping download")
model_config.path = download_target
# Must be a local file
else:
model_config.path = model_config.url
else:
raise ValueError(f"Unknown model type {model_type}")
return model_config.path
class WhisperCallback(AbstractWhisperCallback):
def __init__(self, model_container: WhisperContainer, language: str = None, task: str = None, initial_prompt: str = None,
initial_prompt_mode: VadInitialPromptMode=VadInitialPromptMode.PREPREND_FIRST_SEGMENT, **decodeOptions: dict):
self.model_container = model_container
self.language = language
self.task = task
self.initial_prompt = initial_prompt
self.initial_prompt_mode = initial_prompt_mode
self.decodeOptions = decodeOptions
def invoke(self, audio, segment_index: int, prompt: str, detected_language: str, progress_listener: ProgressListener = None):
"""
Peform the transcription of the given audio file or data.
Parameters
----------
audio: Union[str, np.ndarray, torch.Tensor]
The audio file to transcribe, or the audio data as a numpy array or torch tensor.
segment_index: int
The target language of the transcription. If not specified, the language will be inferred from the audio content.
task: str
The task - either translate or transcribe.
progress_listener: ProgressListener
A callback to receive progress updates.
"""
model = self.model_container.get_model()
if progress_listener is not None:
with create_progress_listener_handle(progress_listener):
return self._transcribe(model, audio, segment_index, prompt, detected_language)
else:
return self._transcribe(model, audio, segment_index, prompt, detected_language)
def _transcribe(self, model: Whisper, audio, segment_index: int, prompt: str, detected_language: str):
decodeOptions = self.decodeOptions.copy()
# Add fp16
if self.model_container.compute_type in ["fp16", "float16"]:
decodeOptions["fp16"] = True
initial_prompt = self._get_initial_prompt(self.initial_prompt, self.initial_prompt_mode, prompt, segment_index)
return model.transcribe(audio, \
language=self.language if self.language else detected_language, task=self.task, \
initial_prompt=initial_prompt, \
**decodeOptions
) |