devasheeshG
/

whisper_medium_fp16_transformers

Automatic Speech Recognition

Inference Endpoints

Model card Files Files and versions Community

whisper_medium_fp16_transformers / __init__.py

devasheeshG's picture

added benchmarks

fbc6b69 12 months ago

raw history blame

No virus

4.6 kB

	from transformers import (
	WhisperForConditionalGeneration, WhisperProcessor, WhisperConfig
	)
	import torch
	import ffmpeg
	import torch
	import torch.nn.functional as F
	import numpy as np
	import os

	# load_audio and pad_or_trim functions
	SAMPLE_RATE = 16000
	CHUNK_LENGTH = 30 # 30-second chunks
	N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000 samples in a 30-second chunk

	class Model:
	def __init__(self,
	model_name_or_path: str,
	cuda_visible_device: str = "0",
	device: str = 'cuda' # torch.device("cuda" if torch.cuda.is_available() else "cpu")
	):

	os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_device
	self.DEVICE = device

	self.processor = WhisperProcessor.from_pretrained(model_name_or_path)
	self.tokenizer = self.processor.tokenizer

	self.config = WhisperConfig.from_pretrained(model_name_or_path)

	self.model = WhisperForConditionalGeneration(
	config=self.config
	).from_pretrained(
	pretrained_model_name_or_path = model_name_or_path,
	torch_dtype = self.config.torch_dtype,
	# device_map=DEVICE, # 'balanced', 'balanced_low_0', 'sequential', 'cuda', 'cpu'
	low_cpu_mem_usage = True,
	)

	# Move model to GPU
	if self.model.device.type != self.DEVICE:
	print(f'Moving model to {self.DEVICE}')
	self.model = self.model.to(self.DEVICE)
	self.model.eval()

	else:
	print(f'Model is already on {self.DEVICE}')
	self.model.eval()

	print('dtype of model acc to config: ', self.config.torch_dtype)
	print('dtype of loaded model: ', self.model.dtype)


	# audio = whisper.load_audio('test.wav')
	def load_audio(self, file: str, sr: int = SAMPLE_RATE, start_time: int = 0, dtype=np.float16):
	try:
	# This launches a subprocess to decode audio while down-mixing and resampling as necessary.
	# Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
	out, _ = (
	ffmpeg.input(file, ss=start_time, threads=0)
	.output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
	.run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
	)
	except ffmpeg.Error as e:
	raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e

	# return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
	return np.frombuffer(out, np.int16).flatten().astype(dtype) / 32768.0


	# audio = whisper.pad_or_trim(audio)
	def _pad_or_trim(self, array, length: int = N_SAMPLES, *, axis: int = -1):
	"""
	Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
	"""
	if torch.is_tensor(array):
	if array.shape[axis] > length:
	array = array.index_select(
	dim=axis, index=torch.arange(length, device=array.device)
	)

	if array.shape[axis] < length:
	pad_widths = [(0, 0)] * array.ndim
	pad_widths[axis] = (0, length - array.shape[axis])
	array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
	else:
	if array.shape[axis] > length:
	array = array.take(indices=range(length), axis=axis)

	if array.shape[axis] < length:
	pad_widths = [(0, 0)] * array.ndim
	pad_widths[axis] = (0, length - array.shape[axis])
	array = np.pad(array, pad_widths)

	return array

	def transcribe(self, audio: np.ndarray, language: str = "english"):
	# audio = load_audio(audio)
	audio = self._pad_or_trim(audio)
	input_features = self.processor(audio, sampling_rate=SAMPLE_RATE, return_tensors="pt").input_features.half().to(self.DEVICE)
	with torch.no_grad():
	predicted_ids = self.model.generate(
	input_features,
	num_beams = 1,
	language=language,
	task="transcribe",
	use_cache=True,
	is_multilingual=True,
	return_timestamps=True,
	)

	transcription = self.tokenizer.batch_decode(predicted_ids, skip_special_tokens=True)[0]
	return transcription.strip()