Views
No views yet
ps)Wav2Vec2ForEmotionClassification class to load this model.1import torch
2import torch.nn as nn
3from transformers import Wav2Vec2Processor, Wav2Vec2Model, Wav2Vec2Config
4
5# Define the Custom Model Class
6class Wav2Vec2ForEmotionClassification(nn.Module):
7 def __init__(self, config):
8 super().__init__()
9 self.wav2vec2 = Wav2Vec2Model(config)
10 self.classifier = nn.Sequential(
11 nn.Linear(config.hidden_size, 256),
12 nn.ReLU(),
13 nn.Dropout(0.3),
14 nn.Linear(256, 128),
15 nn.ReLU(),
16 nn.Dropout(0.3),
17 nn.Linear(128, config.num_labels),
18 )
19
20 def forward(self, input_values, attention_mask=None, labels=None, **kwargs):
21 outputs = self.wav2vec2(input_values, attention_mask=attention_mask)
22 hidden_states = outputs.last_hidden_state
23 pooled_output = torch.mean(hidden_states, dim=1)
24 logits = self.classifier(pooled_output)
25
26 loss = None
27 if labels is not None:
28 loss_fct = nn.CrossEntropyLoss()
29 loss = loss_fct(logits, labels.view(-1))
30
31 return {
32 "loss": loss,
33 "logits": logits
34 }
35
36# Load Model
37model_id = "dynann/emotion-speech-recognition"
38config = Wav2Vec2Config.from_pretrained(model_id)
39model = Wav2Vec2ForEmotionClassification(config)
40model.load_state_dict(torch.hub.load_state_dict_from_url(f"https://huggingface.co/{model_id}/resolve/main/pytorch_model.bin"))
41
42processor = Wav2Vec2Processor.from_pretrained(model_id)