adarshadda's picture
initial commit
b0dd51d
raw
history blame contribute delete
No virus
934 Bytes
import torch
from faster_whisper import WhisperModel
class Transcription:
def __init__(self, model_name: str) -> None:
print("Loading whisper model...")
self.whisper_model = self.load_whisper(model_id=model_name)
print("Loaded whisper model")
def load_whisper(self, model_id: str):
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = "float16" if torch.cuda.is_available() else "float32"
return WhisperModel(model_id, device=device, compute_type=torch_dtype)
def transcribe(self, audio_path: str) -> str:
"""Transcribes the given audio data
Args:
audio_path (str): audio path
Returns:
str: text
"""
segments, info = self.whisper_model.transcribe(audio_path, language="en")
return "".join(segment.text for segment in segments)