Audio Classification
PyTorch
whisper
biology
Inference Endpoints
bnestor commited on
Commit
071b3bf
1 Parent(s): 12fa3b3

Upload handler.py

Browse files
Files changed (1) hide show
  1. handler.py +86 -0
handler.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ handler.py
3
+
4
+ Set up the possibility for an inference endpoint on huggingface.
5
+ """
6
+ from typing import Dict, Any
7
+ import torch
8
+ import torchaudio
9
+ from transformers import WhisperForAudioClassification, WhisperFeatureExtractor
10
+ from transformers.pipelines.audio_utils import ffmpeg_read
11
+ import numpy as np
12
+
13
+ class EndpointHandler():
14
+ """
15
+ This is a wrapper for huggingface models so that they return json objects and consider the same configs as other implementations
16
+ """
17
+ def __init__(self, threshold=0.5):
18
+
19
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
20
+ torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
21
+ model_id = 'DORI-SRKW/whisper-base-mm'
22
+
23
+ # Load the model
24
+ try:
25
+ self.model = WhisperForAudioClassification.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
26
+ except:
27
+ self.model = WhisperForAudioClassification.from_pretrained(model_id, torch_dtype=torch_dtype)
28
+ self.feature_extractor = WhisperFeatureExtractor.from_pretrained(model_id)
29
+
30
+ self.model.eval()
31
+ self.model.to(self.device)
32
+ self.threshold = threshold
33
+
34
+
35
+ def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
36
+ """
37
+ Args:
38
+ data (:obj:):
39
+ includes the input data and the parameters for the inference.
40
+ Return:
41
+ A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
42
+ - "label": A string representing what the label/class is. There can be multiple labels.
43
+ - "score": A score between 0 and 1 describing how confident the model is for this label/class.
44
+ """
45
+
46
+ # step one, get the sampling rate of the audio
47
+ audio = data['audio']
48
+
49
+ fs = data['sampling_rate']
50
+
51
+ # split into 15 second intervals
52
+ audio_np_array = ffmpeg_read(audio, fs)
53
+
54
+ audio = torch.from_numpy(np.asarray(audio_np_array).copy())
55
+ audio = audio.reshape(1, -1)
56
+
57
+ # torchaudio resamples the audio to 32000
58
+ audio = torchaudio.functional.resample(audio, orig_freq=fs, new_freq=32000)
59
+
60
+ # highpass filter 1000 hz
61
+ audio = torchaudio.functional.highpass_biquad(audio, 32000, 1000, 0.707)
62
+
63
+ audio3 = []
64
+ for i in range(0, len(audio[-1]), 32000*15):
65
+ audio3.append(audio[:,i:i+32000*15].squeeze().cpu().data.numpy())
66
+
67
+ data = self.feature_extractor(audio3, sampling_rate = 16000, padding='max_length', max_length=32000*15, return_tensors='pt')
68
+
69
+ try:
70
+ data['input_values'] = data['input_values'].squeeze(0)
71
+ except:
72
+ # it is called input_features for whisper
73
+ data['input_features'] = data['input_features'].squeeze(0)
74
+
75
+ data = {k: v.to(self.device) for k, v in data.items()}
76
+ with torch.amp.autocast(device_type=self.device):
77
+ outputs = []
78
+ for segment in range(data['input_features'].shape[0]):
79
+ # iterate through 15 second segments
80
+ output = self.model(data['input_features'][segment].unsqueeze(0))
81
+
82
+ outputs.append({'logit': torch.softmax(output.logits, dim=1)[0][1].cpu().data.numpy().max(), 'start_time_s': segment*15})
83
+
84
+ outputs = {'logit': max([x['logit'] for x in outputs]), 'classification': 'present' if max([x['logit'] for x in outputs]) >= self.threshold else 'absent'}
85
+ return outputs
86
+