devasheeshG commited on
Commit
0838193
1 Parent(s): 9c5bcaf

Delete __init__.py

Browse files
Files changed (1) hide show
  1. __init__.py +0 -113
__init__.py DELETED
@@ -1,113 +0,0 @@
1
- from transformers import (
2
- WhisperForConditionalGeneration, WhisperProcessor, WhisperConfig
3
- )
4
- import torch
5
- import ffmpeg
6
- import torch
7
- import torch.nn.functional as F
8
- import numpy as np
9
- import os
10
-
11
- # load_audio and pad_or_trim functions
12
- SAMPLE_RATE = 16000
13
- CHUNK_LENGTH = 30 # 30-second chunks
14
- N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000 samples in a 30-second chunk
15
-
16
- class Model:
17
- def __init__(self,
18
- model_name_or_path: str,
19
- cuda_visible_device: str = "0",
20
- device: str = 'cuda' # torch.device("cuda" if torch.cuda.is_available() else "cpu")
21
- ):
22
-
23
- os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_device
24
- self.DEVICE = device
25
-
26
- self.processor = WhisperProcessor.from_pretrained(model_name_or_path)
27
- self.tokenizer = self.processor.tokenizer
28
-
29
- self.config = WhisperConfig.from_pretrained(model_name_or_path)
30
-
31
- self.model = WhisperForConditionalGeneration(
32
- config=self.config
33
- ).from_pretrained(
34
- pretrained_model_name_or_path = model_name_or_path,
35
- torch_dtype = self.config.torch_dtype,
36
- # device_map=DEVICE, # 'balanced', 'balanced_low_0', 'sequential', 'cuda', 'cpu'
37
- low_cpu_mem_usage = True,
38
- )
39
-
40
- # Move model to GPU
41
- if self.model.device.type != self.DEVICE:
42
- print(f'Moving model to {self.DEVICE}')
43
- self.model = self.model.to(self.DEVICE)
44
- self.model.eval()
45
-
46
- else:
47
- print(f'Model is already on {self.DEVICE}')
48
- self.model.eval()
49
-
50
- print('dtype of model acc to config: ', self.config.torch_dtype)
51
- print('dtype of loaded model: ', self.model.dtype)
52
-
53
-
54
- # audio = whisper.load_audio('test.wav')
55
- def load_audio(self, file: str, sr: int = SAMPLE_RATE, start_time: int = 0, dtype=np.float16):
56
- try:
57
- # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
58
- # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
59
- out, _ = (
60
- ffmpeg.input(file, ss=start_time, threads=0)
61
- .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
62
- .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
63
- )
64
- except ffmpeg.Error as e:
65
- raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
66
-
67
- # return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
68
- return np.frombuffer(out, np.int16).flatten().astype(dtype) / 32768.0
69
-
70
-
71
- # audio = whisper.pad_or_trim(audio)
72
- def _pad_or_trim(self, array, length: int = N_SAMPLES, *, axis: int = -1):
73
- """
74
- Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
75
- """
76
- if torch.is_tensor(array):
77
- if array.shape[axis] > length:
78
- array = array.index_select(
79
- dim=axis, index=torch.arange(length, device=array.device)
80
- )
81
-
82
- if array.shape[axis] < length:
83
- pad_widths = [(0, 0)] * array.ndim
84
- pad_widths[axis] = (0, length - array.shape[axis])
85
- array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
86
- else:
87
- if array.shape[axis] > length:
88
- array = array.take(indices=range(length), axis=axis)
89
-
90
- if array.shape[axis] < length:
91
- pad_widths = [(0, 0)] * array.ndim
92
- pad_widths[axis] = (0, length - array.shape[axis])
93
- array = np.pad(array, pad_widths)
94
-
95
- return array
96
-
97
- def transcribe(self, audio: np.ndarray, language: str = "english"):
98
- # audio = load_audio(audio)
99
- audio = self._pad_or_trim(audio)
100
- input_features = self.processor(audio, sampling_rate=SAMPLE_RATE, return_tensors="pt").input_features.half().to(self.DEVICE)
101
- with torch.no_grad():
102
- predicted_ids = self.model.generate(
103
- input_features,
104
- num_beams = 1,
105
- language=language,
106
- task="transcribe",
107
- use_cache=True,
108
- is_multilingual=True,
109
- return_timestamps=True,
110
- )
111
-
112
- transcription = self.tokenizer.batch_decode(predicted_ids, skip_special_tokens=True)[0]
113
- return transcription.strip()