iRecite-MVP-API / step5_wavlm_test.py
didodev
Deploy iRecite MVP API (Docker + FastAPI)
4ca6263
import torch
import librosa
from transformers import AutoFeatureExtractor, AutoModel
MODEL_ID = "microsoft/wavlm-base"
def load_audio(path: str, target_sr: int = 16000):
audio, sr = librosa.load(path, sr=target_sr, mono=True)
return audio, sr
def main():
print("Loading model:", MODEL_ID)
feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_ID)
model = AutoModel.from_pretrained(MODEL_ID)
model.eval()
audio, sr = load_audio("sample.wav")
print("Audio length (sec):", round(len(audio) / sr, 2))
inputs = feature_extractor(audio, sampling_rate=sr, return_tensors="pt")
with torch.no_grad():
out = model(**inputs)
x = out.last_hidden_state # [batch, frames, hidden]
print("OK ✅ WavLM ran on CPU")
print("Embedding tensor shape:", tuple(x.shape))
if __name__ == "__main__":
main()