File size: 2,870 Bytes
e630855 | 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 75 76 77 78 79 80 81 82 | """Gender detection using a fine-tuned wav2vec2 classifier.
Model: prithivMLmods/Common-Voice-Gender-Detection
https://huggingface.co/prithivMLmods/Common-Voice-Gender-Detection
A `facebook/wav2vec2-base-960h` model fine-tuned for binary (female/male)
speaker-gender classification. We feed it a 16 kHz mono waveform and read the
softmax probabilities. If the top probability is below a confidence floor (or
there isn't enough audio) we report "uncertain" instead of guessing.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from functools import lru_cache
import numpy as np
import torch
from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification
MODEL_NAME = "prithivMLmods/Common-Voice-Gender-Detection"
TARGET_SR = 16000
MIN_SAMPLES = int(0.3 * TARGET_SR) # need ~0.3s of audio to bother
CONFIDENCE_FLOOR = 0.6 # below this we say "uncertain"
@dataclass
class GenderResult:
label: str # "male" | "female" | "uncertain"
confidence: float # top-class softmax probability (0..1)
probs: dict = field(default_factory=dict) # {"male": p, "female": p}
audio_seconds: float = 0.0
@lru_cache(maxsize=1)
def _load():
"""Load (and cache) the model + feature extractor once per process."""
model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_NAME)
extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_NAME)
model.eval()
id2label = {int(k): str(v).lower() for k, v in model.config.id2label.items()}
return model, extractor, id2label
def warmup() -> None:
"""Eagerly load the model (e.g. at server startup) to avoid a cold first request."""
_load()
def classify(samples: np.ndarray, sr: int) -> GenderResult:
"""Classify gender from a mono waveform (float32)."""
audio_seconds = float(len(samples) / sr) if sr else 0.0
if samples is None or samples.size < MIN_SAMPLES:
return GenderResult("uncertain", 0.0, {}, round(audio_seconds, 2))
samples = np.asarray(samples, dtype=np.float32)
if sr != TARGET_SR:
import librosa
samples = librosa.resample(samples, orig_sr=sr, target_sr=TARGET_SR)
sr = TARGET_SR
model, extractor, id2label = _load()
inputs = extractor(samples, sampling_rate=sr, return_tensors="pt", padding=True)
with torch.no_grad():
logits = model(**inputs).logits
probs = torch.softmax(logits, dim=1).squeeze(0).tolist()
prob_map = {id2label[i]: float(probs[i]) for i in range(len(probs))}
label = max(prob_map, key=prob_map.get)
confidence = prob_map[label]
if confidence < CONFIDENCE_FLOOR:
label = "uncertain"
return GenderResult(
label=label,
confidence=round(confidence, 3),
probs={k: round(v, 3) for k, v in prob_map.items()},
audio_seconds=round(audio_seconds, 2),
)
|