Views
No views yet
1import numpy as np
2import torch
3import torch.nn as nn
4from transformers import Wav2Vec2Processor
5from transformers.models.wav2vec2.modeling_wav2vec2 import (
6 Wav2Vec2Model,
7 Wav2Vec2PreTrainedModel,
8)
9
10
11class RegressionHead(nn.Module):
12 r"""Classification head."""
13
14 def __init__(self, config):
15
16 super().__init__()
17
18 self.dense = nn.Linear(config.hidden_size, config.hidden_size)
19 self.dropout = nn.Dropout(config.final_dropout)
20 self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
21
22 def forward(self, features, **kwargs):
23
24 x = features
25 x = self.dropout(x)
26 x = self.dense(x)
27 x = torch.tanh(x)
28 x = self.dropout(x)
29 x = self.out_proj(x)
30
31 return x
32
33
34class EmotionModel(Wav2Vec2PreTrainedModel):
35 r"""Speech emotion classifier."""
36
37 def __init__(self, config):
38
39 super().__init__(config)
40
41 self.config = config
42 self.wav2vec2 = Wav2Vec2Model(config)
43 self.classifier = RegressionHead(config)
44 self.init_weights()
45
46 def forward(
47 self,
48 input_values,
49 ):
50
51 outputs = self.wav2vec2(input_values)
52 hidden_states = outputs[0]
53 hidden_states = torch.mean(hidden_states, dim=1)
54 logits = self.classifier(hidden_states)
55
56 return hidden_states, logits
57
58
59
60# load model from hub
61device = 'cpu'
62model_name = 'audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim'
63processor = Wav2Vec2Processor.from_pretrained(model_name)
64model = EmotionModel.from_pretrained(model_name)
65
66# dummy signal
67sampling_rate = 16000
68signal = np.zeros((1, sampling_rate), dtype=np.float32)
69
70
71def process_func(
72 x: np.ndarray,
73 sampling_rate: int,
74 embeddings: bool = False,
75) -> np.ndarray:
76 r"""Predict emotions or extract embeddings from raw audio signal."""
77
78 # run through processor to normalize signal
79 # always returns a batch, so we just get the first entry
80 # then we put it on the device
81 y = processor(x, sampling_rate=sampling_rate)
82 y = y['input_values'][0]
83 y = y.reshape(1, -1)
84 y = torch.from_numpy(y).to(device)
85
86 # run through model
87 with torch.no_grad():
88 y = model(y)[0 if embeddings else 1]
89
90 # convert to numpy
91 y = y.detach().cpu().numpy()
92
93 return y
94
95
96print(process_func(signal, sampling_rate))
97# Arousal dominance valence
98# [[0.5460754 0.6062266 0.40431657]]
99
100print(process_func(signal, sampling_rate, embeddings=True))
101# Pooled hidden states of last transformer layer
102# [[-0.00752167 0.0065819 -0.00746342 ... 0.00663632 0.00848748
103# 0.00599211]]