Views
No views yet
model.safetensors — PyTorch weights (safetensors)config.json — model architecturepreprocessor_config.json — audio feature extraction settingslabel_mapping.json — index → label1import torch, json, numpy as np, librosa
2from safetensors.torch import load_file as load_safetensors
3
4# Load config
5import json, os
6model_dir = "./hf/models/shanghai-binary"
7cfg = json.load(open(os.path.join(model_dir, "config.json")))
8pp = json.load(open(os.path.join(model_dir, "preprocessor_config.json")))
9lm = json.load(open(os.path.join(model_dir, "label_mapping.json")))
10
11# Define the model class you trained (LanNetBinary)
12# (Same as in your training notebook)
13class LanNetBinary(torch.nn.Module):
14 def __init__(self, input_dim=40, hidden_dim=512, num_layers=2):
15 super().__init__()
16 self.gru = torch.nn.GRU(input_dim, hidden_dim, num_layers=num_layers, batch_first=True)
17 self.linear2 = torch.nn.Linear(hidden_dim, 192)
18 self.linear3 = torch.nn.Linear(192, 2)
19 def forward(self, x):
20 out, _ = self.gru(x)
21 last = out[:, -1, :]
22 x = self.linear2(last)
23 x = self.linear3(x)
24 return x
25
26# Load weights
27model = LanNetBinary(cfg["input_dim"], cfg["hidden_dim"], cfg["num_layers"])
28sd = load_safetensors(os.path.join(model_dir, "model.safetensors"))
29model.load_state_dict(sd, strict=True)
30model.eval()
31
32# Feature extraction should match preprocessor_config.json
33def fbanks_from_array(y, sr=pp["sampling_rate"], n_mels=pp["n_mels"], n_fft=pp["n_fft"], hop_length=pp["hop_length"], max_len=pp["max_len_frames"]):
34 mel = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length, power=2.0)
35 fbanks = librosa.power_to_db(mel).T
36 T = fbanks.shape[0]
37 if T < max_len:
38 import numpy as np
39 fbanks = np.pad(fbanks, ((0, max_len - T), (0, 0)), mode="constant")
40 else:
41 fbanks = fbanks[:max_len, :]
42 return torch.tensor(fbanks, dtype=torch.float32).unsqueeze(0) # (1, T, F)
43
44# Example: predict from a waveform array "y" at 16kHz
45# y, _ = librosa.load("example.wav", sr=pp["sampling_rate"])
46# x = fbanks_from_array(y)
47# with torch.no_grad():
48# logits = model(x)
49# pred = int(torch.argmax(logits, dim=1))
50# print(lm[str(pred)])