Marko Henning
376 lines
13 KiB
Python
376 lines
13 KiB
Python
"""
|
|
Transcriber Module
|
|
------------------
|
|
|
|
This module provides the Transcriber class, a comprehensive tool for working with Whisper models.
|
|
The Transcriber class offers functionalities such as loading different Whisper models, transcribing audio files,
|
|
and saving transcriptions to text files. It acts as an interface between various Whisper models and the user,
|
|
simplifying the process of audio transcription.
|
|
|
|
Main Features:
|
|
- Loading different sizes and versions of Whisper models.
|
|
- Transcribing audio in various formats including str, Tensor, and nparray.
|
|
- Saving the transcriptions to the specified paths.
|
|
- Adaptable to various language specifications.
|
|
- Options to control the verbosity of the transcription process.
|
|
|
|
Constants:
|
|
WHISPER_DEFAULT_PATH: Default path for downloading and loading Whisper models.
|
|
|
|
Usage:
|
|
>>> from your_package import Transcriber
|
|
>>> transcriber = Transcriber.load_model(model="medium")
|
|
>>> transcript = transcriber.transcribe(audio="path/to/audio.wav")
|
|
>>> transcriber.save_transcript(transcript, "path/to/save.txt")
|
|
"""
|
|
|
|
from whisper import Whisper
|
|
from whisper import load_model as whisper_load_model
|
|
from whisperx.asr import WhisperModel
|
|
from whisperx import load_model as whisperx_load_model
|
|
from typing import TypeVar , Union , Optional
|
|
from torch import Tensor, device
|
|
from numpy import ndarray
|
|
from inspect import getfullargspec
|
|
from abc import ABC, abstractmethod
|
|
|
|
from .misc import WHISPER_DEFAULT_PATH
|
|
whisper = TypeVar('whisper')
|
|
|
|
|
|
class Transcriber:
|
|
"""
|
|
Transcriber Class
|
|
-----------------
|
|
|
|
The Transcriber class serves as a wrapper around Whisper models for efficient audio
|
|
transcription. By encapsulating the intricacies of loading models, processing audio,
|
|
and saving transcripts, it offers an easy-to-use interface
|
|
for users to transcribe audio files.
|
|
|
|
Attributes:
|
|
model (whisper): The Whisper model used for transcription.
|
|
|
|
Methods:
|
|
transcribe: Transcribes the given audio file.
|
|
save_transcript: Saves the transcript to a file.
|
|
load_model: Loads a specific Whisper model.
|
|
_get_whisper_kwargs: Private method to get valid keyword arguments for the whisper model.
|
|
|
|
Examples:
|
|
>>> transcriber = Transcriber.load_model(model="medium")
|
|
>>> transcript = transcriber.transcribe(audio="path/to/audio.wav")
|
|
>>> transcriber.save_transcript(transcript, "path/to/save.txt")
|
|
|
|
Note:
|
|
The class supports various sizes and versions of Whisper models. Please refer to
|
|
the load_model method for available options.
|
|
"""
|
|
def __init__(self, model: whisper, model_name: str) -> None:
|
|
"""
|
|
Initialize the Transcriber class with a Whisper model.
|
|
|
|
Args:
|
|
model (whisper): The Whisper model to use for transcription.
|
|
model_name (str): The name of the model.
|
|
"""
|
|
|
|
self.model = model
|
|
|
|
self.model_name = model_name
|
|
|
|
@abstractmethod
|
|
def transcribe(self, audio: Union[str, Tensor, ndarray] ,
|
|
*args, **kwargs) -> str:
|
|
"""
|
|
Transcribe an audio file.
|
|
|
|
Args:
|
|
audio (Union[str, Tensor, nparray]): The audio file to transcribe.
|
|
*args: Additional arguments.
|
|
**kwargs: Additional keyword arguments,
|
|
such as the language of the audio file.
|
|
|
|
Returns:
|
|
str: The transcript as a string.
|
|
"""
|
|
pass
|
|
|
|
@staticmethod
|
|
def save_transcript(transcript : str , save_path : str) -> None:
|
|
"""
|
|
Save a transcript to a file.
|
|
|
|
Args:
|
|
transcript (str): The transcript as a string.
|
|
save_path (str): The path to save the transcript.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
|
|
with open(save_path, 'w') as f:
|
|
f.write(transcript)
|
|
|
|
print(f'Transcript saved to {save_path}')
|
|
|
|
@classmethod
|
|
def load_model(cls,
|
|
model: str = "medium",
|
|
whisper_type: str = 'whisper',
|
|
download_root: str = WHISPER_DEFAULT_PATH,
|
|
device: Optional[Union[str, device]] = None,
|
|
in_memory: bool = False,
|
|
*args, **kwargs
|
|
) -> 'Transcriber':
|
|
"""
|
|
Load whisper model.
|
|
|
|
Args:
|
|
model (str): Whisper model. Available models include:
|
|
- 'tiny.en'
|
|
- 'tiny'
|
|
- 'base.en'
|
|
- 'base'
|
|
- 'small.en'
|
|
- 'small'
|
|
- 'medium.en'
|
|
- 'medium'
|
|
- 'large-v1'
|
|
- 'large-v2'
|
|
- 'large-v3'
|
|
- 'large'
|
|
|
|
download_root (str, optional): Path to download the model.
|
|
Defaults to WHISPER_DEFAULT_PATH.
|
|
|
|
device (Optional[Union[str, torch.device]], optional):
|
|
Device to load model on. Defaults to None.
|
|
in_memory (bool, optional): Whether to load model in memory.
|
|
Defaults to False.
|
|
args: Additional arguments only to avoid errors.
|
|
kwargs: Additional keyword arguments only to avoid errors.
|
|
|
|
Returns:
|
|
Transcriber: A Transcriber object initialized with the specified model.
|
|
"""
|
|
if whisper_type.lower() == 'whisper':
|
|
_model = WhisperTranscriber.load_model(
|
|
model, download_root, device, in_memory, *args, **kwargs)
|
|
return _model
|
|
elif whisper_type.lower() == 'whisperx':
|
|
_model = WhisperXTranscriber.load_model(
|
|
model, download_root, device, *args, **kwargs)
|
|
return _model
|
|
else:
|
|
raise ValueError(f'Model type not recognized, exptected "whisper" '
|
|
f'or "whisperx", got {whisper_type}.')
|
|
pass
|
|
|
|
@staticmethod
|
|
def _get_whisper_kwargs(**kwargs) -> dict:
|
|
"""
|
|
Get kwargs for whisper model. Ensure that kwargs are valid.
|
|
|
|
Returns:
|
|
dict: Keyword arguments for whisper model.
|
|
"""
|
|
pass
|
|
|
|
def __repr__(self) -> str:
|
|
return f"Transcriber(model_name={self.model_name}, model={self.model})"
|
|
|
|
|
|
class WhisperTranscriber(Transcriber):
|
|
def __init__(self, model: whisper, model_name: str) -> None:
|
|
super().__init__(model, model_name)
|
|
|
|
def transcribe(self, audio: Union[str, Tensor, ndarray],
|
|
*args, **kwargs) -> str:
|
|
"""
|
|
Transcribe an audio file.
|
|
|
|
Args:
|
|
audio (Union[str, Tensor, nparray]): The audio file to transcribe.
|
|
*args: Additional arguments.
|
|
**kwargs: Additional keyword arguments,
|
|
such as the language of the audio file.
|
|
|
|
Returns:
|
|
str: The transcript as a string.
|
|
"""
|
|
|
|
kwargs = self._get_whisper_kwargs(**kwargs)
|
|
|
|
if not kwargs.get("verbose"):
|
|
kwargs["verbose"] = None
|
|
|
|
result = self.model.transcribe(audio, *args, **kwargs)
|
|
return result["text"]
|
|
|
|
@classmethod
|
|
def load_model(cls,
|
|
model: str = "medium",
|
|
download_root: str = WHISPER_DEFAULT_PATH,
|
|
device: Optional[Union[str, device]] = None,
|
|
in_memory: bool = False,
|
|
*args, **kwargs
|
|
) -> 'Transcriber':
|
|
"""
|
|
Load whisper model.
|
|
|
|
Args:
|
|
model (str): Whisper model. Available models include:
|
|
- 'tiny.en'
|
|
- 'tiny'
|
|
- 'base.en'
|
|
- 'base'
|
|
- 'small.en'
|
|
- 'small'
|
|
- 'medium.en'
|
|
- 'medium'
|
|
- 'large-v1'
|
|
- 'large-v2'
|
|
- 'large-v3'
|
|
- 'large'
|
|
|
|
download_root (str, optional): Path to download the model.
|
|
Defaults to WHISPER_DEFAULT_PATH.
|
|
|
|
device (Optional[Union[str, torch.device]], optional):
|
|
Device to load model on. Defaults to None.
|
|
in_memory (bool, optional): Whether to load model in memory.
|
|
Defaults to False.
|
|
args: Additional arguments only to avoid errors.
|
|
kwargs: Additional keyword arguments only to avoid errors.
|
|
|
|
Returns:
|
|
Transcriber: A Transcriber object initialized with the specified model.
|
|
"""
|
|
|
|
_model = whisper_load_model(model, download_root=download_root,
|
|
device=device, in_memory=in_memory)
|
|
|
|
return cls(_model, model_name=model)
|
|
|
|
@staticmethod
|
|
def _get_whisper_kwargs(**kwargs) -> dict:
|
|
"""
|
|
Get kwargs for whisper model. Ensure that kwargs are valid.
|
|
|
|
Returns:
|
|
dict: Keyword arguments for whisper model.
|
|
"""
|
|
# _possible_kwargs = WhisperModel.transcribe.__code__.co_varnames
|
|
_args = getfullargspec(Whisper.transcribe).args
|
|
_kwargs = getfullargspec(Whisper.transcribe).kwonlyargs
|
|
_possible_kwargs = _args + _kwargs
|
|
|
|
whisper_kwargs = {k: v for k, v in kwargs.items() if k in _possible_kwargs}
|
|
|
|
if (task := kwargs.get("task")):
|
|
whisper_kwargs["task"] = task
|
|
|
|
if (language := kwargs.get("language")):
|
|
whisper_kwargs["language"] = language
|
|
|
|
return whisper_kwargs
|
|
|
|
|
|
class WhisperXTranscriber(Transcriber):
|
|
def __init__(self, model: whisper, model_name: str) -> None:
|
|
super().__init__(model, model_name)
|
|
|
|
def transcribe(self, audio: Union[str, Tensor, ndarray],
|
|
*args, **kwargs) -> str:
|
|
"""
|
|
Transcribe an audio file.
|
|
|
|
Args:
|
|
audio (Union[str, Tensor, nparray]): The audio file to transcribe.
|
|
*args: Additional arguments.
|
|
**kwargs: Additional keyword arguments,
|
|
such as the language of the audio file.
|
|
|
|
Returns:
|
|
str: The transcript as a string.
|
|
"""
|
|
kwargs = self._get_whisper_kwargs(**kwargs)
|
|
|
|
if isinstance(audio, Tensor):
|
|
audio = audio.cpu().numpy()
|
|
result = self.model.transcribe(audio, *args, **kwargs)
|
|
text = ""
|
|
for seg in result['segments']:
|
|
text += seg['text']
|
|
return text
|
|
|
|
|
|
@classmethod
|
|
def load_model(cls,
|
|
model: str = "medium",
|
|
download_root: str = WHISPER_DEFAULT_PATH,
|
|
device: Optional[Union[str, device]] = None,
|
|
*args, **kwargs
|
|
) -> 'Transcriber':
|
|
"""
|
|
Load whisper model.
|
|
|
|
Args:
|
|
model (str): Whisper model. Available models include:
|
|
- 'tiny.en'
|
|
- 'tiny'
|
|
- 'base.en'
|
|
- 'base'
|
|
- 'small.en'
|
|
- 'small'
|
|
- 'medium.en'
|
|
- 'medium'
|
|
- 'large-v1'
|
|
- 'large-v2'
|
|
- 'large-v3'
|
|
- 'large'
|
|
|
|
download_root (str, optional): Path to download the model.
|
|
Defaults to WHISPER_DEFAULT_PATH.
|
|
|
|
device (Optional[Union[str, torch.device]], optional):
|
|
Device to load model on. Defaults to None.
|
|
in_memory (bool, optional): Whether to load model in memory.
|
|
Defaults to False.
|
|
args: Additional arguments only to avoid errors.
|
|
kwargs: Additional keyword arguments only to avoid errors.
|
|
|
|
Returns:
|
|
Transcriber: A Transcriber object initialized with the specified model.
|
|
"""
|
|
if not isinstance(device, str):
|
|
device = str(device)
|
|
_model = whisperx_load_model(model, download_root=download_root,
|
|
device=device)
|
|
|
|
return cls(_model, model_name=model)
|
|
|
|
@staticmethod
|
|
def _get_whisper_kwargs(**kwargs) -> dict:
|
|
"""
|
|
Get kwargs for whisper model. Ensure that kwargs are valid.
|
|
|
|
Returns:
|
|
dict: Keyword arguments for whisper model.
|
|
"""
|
|
# _possible_kwargs = WhisperModel.transcribe.__code__.co_varnames
|
|
_args = getfullargspec(WhisperModel.transcribe).args
|
|
_kwargs = getfullargspec(WhisperModel.transcribe).kwonlyargs
|
|
_possible_kwargs = _args + _kwargs
|
|
|
|
whisper_kwargs = {k: v for k, v in kwargs.items() if k in _possible_kwargs}
|
|
|
|
if (task := kwargs.get("task")):
|
|
whisper_kwargs["task"] = task
|
|
|
|
if (language := kwargs.get("language")):
|
|
whisper_kwargs["language"] = language
|
|
|
|
return whisper_kwargs
|