| """Transformers-compatible wrapper for harpertoken/tone. |
| |
| The model is a plain two-conv CNN over log-mel spectrograms, not a |
| transformer. This subclasses PreTrainedModel so the weights serialize honestly |
| as config.json plus model.safetensors and load through AutoModel. Importing |
| this module registers the architecture. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel |
|
|
|
|
| class ToneConfig(PretrainedConfig): |
| model_type = "tone_cnn" |
|
|
| def __init__(self, num_labels=10, dropout=0.25, **kwargs): |
| super().__init__(**kwargs) |
| self.num_labels = num_labels |
| self.dropout = dropout |
|
|
|
|
| LABELS = ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go"] |
|
|
|
|
| class ToneCNN(PreTrainedModel): |
| config_class = ToneConfig |
| base_model_prefix = "cnn" |
|
|
| |
| |
| |
| all_tied_weights_keys = {} |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.cnn = nn.ModuleDict( |
| { |
| "conv1": nn.Conv2d(1, 32, 3, padding=1), |
| "conv2": nn.Conv2d(32, 64, 3, padding=1), |
| "fc1": nn.Linear(64 * 8 * 8, 128), |
| "fc2": nn.Linear(128, config.num_labels), |
| } |
| ) |
| self.pool = nn.MaxPool2d(2, 2) |
| self.adapt = nn.AdaptiveAvgPool2d((8, 8)) |
| self.drop = nn.Dropout(config.dropout) |
|
|
| def forward(self, mel): |
| x = self.pool(torch.relu(self.cnn["conv1"](mel))) |
| x = self.pool(torch.relu(self.cnn["conv2"](x))) |
| x = self.adapt(x) |
| x = x.view(x.size(0), -1) |
| return self.cnn["fc2"](self.drop(torch.relu(self.cnn["fc1"](x)))) |
|
|
| def predict_label(self, waveform): |
| """waveform: 16,000 float samples at 16 kHz. Returns the command.""" |
| import torchaudio |
|
|
| self.eval() |
| x = torch.as_tensor(waveform, dtype=torch.float32) |
| mel = torchaudio.transforms.MelSpectrogram( |
| sample_rate=16000, n_fft=512, win_length=400, hop_length=160, |
| n_mels=64, |
| )(x.unsqueeze(0)) |
| logmel = torch.log10(mel + 1e-10).unsqueeze(1) |
| with torch.no_grad(): |
| return LABELS[int(self.forward(logmel).argmax(1)[0])] |
|
|
|
|
| def register(): |
| AutoConfig.register(ToneConfig.model_type, ToneConfig, exist_ok=True) |
| AutoModel.register(ToneConfig, ToneCNN, exist_ok=True) |
|
|
|
|
| register() |
|
|
