Spaces:
Runtime error
Runtime error
File size: 2,408 Bytes
a93df0f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import torch
import torch.nn as nn
from transformers import Wav2Vec2Processor, Wav2Vec2Model, Wav2Vec2ForCTC, HubertModel, HubertForCTC
#import whisper
class WhisperModel(nn.Module):
def __init__(self, model_type="small.en", n_class=14):
super().__init__()
self.encoder = whisper.load_model(model_type).encoder
for param in self.encoder.parameters():
param.requires_grad = True
feature_dim = 768
# 512 = tiny.en,
# 768 = small.en
self.intent_classifier = nn.Sequential(
nn.Linear(feature_dim, n_class)
)
def forward(self, x):
x = self.encoder(x)
x = torch.mean(x, dim=1)
intent = self.intent_classifier(x)
return intent
class Wav2VecModel(nn.Module):
def __init__(self, ):
super().__init__()
self.processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h")
self.encoder = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-large-960h")
for param in self.encoder.parameters():
param.requires_grad = False
for param in self.encoder.encoder.parameters():
param.requires_grad = True
self.intent_classifier = nn.Sequential(
nn.Linear(1024, 14),
)
def forward(self, x):
x = self.processor(x, sampling_rate=16000, return_tensors="pt")["input_values"].squeeze(0).to("cuda")
x = self.encoder(x).last_hidden_state
x = torch.mean(x, dim=1)
logits = self.intent_classifier(x)
return logits
class HubertSSLModel(nn.Module):
def __init__(self, ):
super().__init__()
self.processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
self.encoder = HubertModel.from_pretrained("facebook/hubert-large-ll60k")
for param in self.encoder.parameters():
param.requires_grad = False
for param in self.encoder.encoder.parameters():
param.requires_grad = True
self.intent_classifier = nn.Sequential(
nn.Linear(1024, 14),
)
def forward(self, x):
x = self.processor(x, sampling_rate=16000, return_tensors="pt")["input_values"].squeeze(0).to("cuda")
x = self.encoder(x).last_hidden_state
x = torch.mean(x, dim=1)
logits = self.intent_classifier(x)
return logits |