| """Reference-based pronunciation scorer (2-stage RAPS). |
| |
| Inputs: user audio + reference (standard) audio + canonical text. |
| Stage 1 (RAPS, end-to-end fine-tuned WavLM + reference cross-attention): |
| user/ref audio → V21 forced-align → per-phone user-vs-ref comparison → per-phone engine-style score |
| Stage 2 (Transformer aggregator): |
| per-phone scores → utterance accuracy + integrity (engine-calibrated) |
| |
| Holdout results (260520, no leakage): |
| word : acc P<10 80.2% / P<15 90.2% ; integ P<10 98.6% |
| sentence: acc P<10 77.3% / P<15 89.0% ; integ P<10 82.4% / P<15 93.3% |
| all : acc P<10 78.6% / P<15 89.3% ; integ P<10 87.8% / P<15 94.2% |
| |
| Models: |
| exp_raps_xattn/best.pth (RAPS stage-1, ~324M) |
| exp_stage2_agg/best.pth (Transformer aggregator, 3.3MB) |
| sentence_scorer_v32v3_compact/v21_xlsr_phoneme.pt (V21 forced align) |
| sentence_scorer_v32v3_compact/phone_vocab.json |
| gpu_service/custom_phone_dict.json (G2P) |
| |
| Usage: |
| python3 inference_reference.py --user user.wav --ref ref.wav --text "I am going to school" |
| """ |
| import os, sys, json, re, argparse |
| from pathlib import Path |
| import numpy as np |
| import torch, torch.nn as nn |
| import torch.nn.functional as F |
| import soundfile as sf |
| from transformers import WavLMModel, Wav2Vec2ForCTC |
|
|
| REPO = Path(__file__).resolve().parent |
| sys.path.insert(0, str(REPO)) |
| sys.path.insert(0, str(REPO/'sentence_scorer_v32v3_compact')) |
| SR = 16000; MAX_PHONES = 50 |
|
|
| |
| _G2P=None; _CD=None |
| def g2p(): |
| global _G2P |
| if _G2P is None: |
| import g2p_en; _G2P=g2p_en.G2p() |
| return _G2P |
| def custom_dict(): |
| global _CD |
| if _CD is None: |
| p=REPO/'gpu_service/custom_phone_dict.json'; _CD=json.load(open(p)) if p.exists() else {} |
| return _CD |
| def text_to_phones(text, phone_to_id): |
| """Return phone ids + word_id-per-phone (for fluency word grouping).""" |
| text=re.sub(r"[^a-z'\s]"," ",text.lower()); text=re.sub(r"\s+"," ",text).strip() |
| ids=[]; wids=[] |
| for wi,w in enumerate(text.split()): |
| cd=custom_dict() |
| phs=cd[w] if w in cd else [re.sub(r"\d+$","",t).lower() for t in g2p()(w) if t!=" "] |
| for p in phs: |
| ids.append(phone_to_id.get(p, phone_to_id.get('<unk>',1))); wids.append(wi) |
| return ids, wids |
|
|
| def load_audio(path, max_sec=10.0): |
| w,sr=sf.read(path,dtype='float32') |
| if w.ndim>1: w=w.mean(1) |
| if sr!=SR: |
| idx=np.linspace(0,len(w)-1,int(len(w)*SR/sr)).astype(np.int64); w=w[idx] |
| w=w[:int(max_sec*SR)]; w=np.pad(w,int(0.2*SR)).astype(np.float32) |
| return (w-w.mean())/(w.std()+1e-7) |
|
|
| |
| from inference_sentence import V21Aligner |
|
|
| |
| class RAPS(nn.Module): |
| def __init__(s, n_phon=60, embed=256): |
| super().__init__() |
| s.w=WavLMModel.from_pretrained('microsoft/wavlm-large'); H=s.w.config.hidden_size |
| s.proj=nn.Linear(H*4,embed); s.pe=nn.Embedding(n_phon+1,embed,padding_idx=0) |
| s.pos=nn.Parameter(torch.zeros(1,MAX_PHONES,embed)) |
| s.blocks=nn.ModuleList([nn.TransformerEncoderLayer(embed,8,embed*4,0.1,batch_first=True) for _ in range(4)]) |
| s.xattn=nn.MultiheadAttention(H,8,batch_first=True) |
| s.phn=nn.Linear(embed,1); s.pherr=nn.Linear(embed,1) |
| s.utt=nn.Sequential(nn.Linear(embed+7,embed),nn.GELU(),nn.Linear(embed,5)) |
| def enc(s,w,dev): |
| wt=torch.from_numpy(w).unsqueeze(0).to(dev) |
| with torch.amp.autocast('cuda',dtype=torch.bfloat16): |
| o=s.w(wt,attention_mask=torch.ones(1,len(w),device=dev)).last_hidden_state[0] |
| return o.float(), s.w._get_feat_extract_output_lengths(torch.tensor([len(w)])).item() |
|
|
| class AGG(nn.Module): |
| """Unified 4-metric aggregator: phone scores + fluency feats → acc/integ/flu/overall.""" |
| def __init__(s,emb=128,n_phon=60,nflu=8): |
| super().__init__(); s.inp=nn.Linear(2,emb); s.pe=nn.Embedding(n_phon+1,emb,padding_idx=0) |
| s.pos=nn.Parameter(torch.zeros(1,MAX_PHONES,emb)) |
| s.bl=nn.ModuleList([nn.TransformerEncoderLayer(emb,8,emb*4,0.1,batch_first=True) for _ in range(4)]) |
| s.h=nn.Sequential(nn.Linear(emb+nflu,emb),nn.GELU(),nn.Linear(emb,4)) |
| def forward(s,x,pc,f): |
| vm=(pc>0); h=s.inp(x)+s.pe(pc)+s.pos[:,:x.size(1)] |
| for b in s.bl: h=b(h,src_key_padding_mask=~vm) |
| vmf=vm.float().unsqueeze(-1); pooled=(h*vmf).sum(1)/vmf.sum(1).clamp(min=1) |
| return s.h(torch.cat([pooled,f],-1)) |
|
|
| def fluency_feats(bu, wids, n): |
| """8-dim fluency feats from V21 user spans + word ids — matches unified_aggregator training.""" |
| s=np.array([bu[i][0] for i in range(n)]); e=np.array([bu[i][1] for i in range(n)]); w=np.array(wids[:n]) |
| total=e.max()-s.min()+1 if n>0 else 1 |
| wd=[]; gaps=[]; prev=None |
| for wi in sorted(set(w.tolist())): |
| ph=np.where(w==wi)[0]; ws=s[ph].min(); we=e[ph].max(); wd.append(we-ws+1) |
| if prev is not None: gaps.append(max(0,ws-prev)) |
| prev=we |
| wd=np.array(wd) if wd else np.array([1.]); gaps=np.array(gaps) if gaps else np.array([0.]) |
| return np.array([n/max(total,1)*50, total/50.0, len(set(w.tolist())), |
| wd.mean()/50.0, wd.std()/50.0, gaps.mean()/50.0, gaps.max()/50.0, |
| float((gaps>10).sum())], np.float32) |
|
|
| def mapb(b,Tv,Tw): |
| a=int(b[0]*Tw/max(Tv,1)); bb=max(a,int((b[1]+1)*Tw/max(Tv,1))-1) |
| a=max(0,min(a,Tw-1)); bb=max(a,min(bb,Tw-1)); return a,bb |
|
|
| class ReferenceScorer: |
| def __init__(s, device='cuda:0', |
| raps='exp_raps_clean/best.pth', agg='exp_unified_agg/best.pth', |
| v21='sentence_scorer_v32v3_compact/v21_xlsr_phoneme.pt', |
| vocab='sentence_scorer_v32v3_compact/phone_vocab.json'): |
| s.dev=device; torch.cuda.set_device(int(device.split(':')[1])) |
| pv=json.load(open(REPO/vocab)); s.vocab=pv['phone_vocab']; s.p2id=pv['phone_to_id'] |
| s.aligner=V21Aligner(str(REPO/v21), device) |
| s.raps=RAPS().to(device); s.raps.load_state_dict(torch.load(REPO/raps,map_location='cpu',weights_only=False)['model_state'],strict=False); s.raps.eval() |
| s.agg=AGG().to(device); s.agg.load_state_dict(torch.load(REPO/agg,map_location='cpu',weights_only=False)['model_state']); s.agg.eval() |
|
|
| @torch.no_grad() |
| def score(s, user_path, ref_path, text): |
| ids,wids=text_to_phones(text, s.p2id) |
| ids=ids[:MAX_PHONES]; wids=wids[:MAX_PHONES] |
| n=len(ids) |
| if n==0: return {'error':'empty text'} |
| wu=load_audio(user_path); wr=load_audio(ref_path) |
| bu,Tu,_=s.aligner.align_and_posteriors(wu,ids) |
| br,Tr,_=s.aligner.align_and_posteriors(wr,ids) |
| Hu,Lu=s.raps.enc(wu,s.dev); Hr,Lr=s.raps.enc(wr,s.dev) |
| toks=[] |
| for p in range(n): |
| a,b=mapb(bu[p],Tu,Lu); c,d=mapb(br[p],Tr,Lr) |
| U=Hu[a:b+1]; R=Hr[c:d+1] |
| A,_=s.raps.xattn(U.unsqueeze(0),R.unsqueeze(0),R.unsqueeze(0),need_weights=False); A=A[0] |
| u=U.mean(0); am=A.mean(0); mx=(U-A).abs().max(0).values |
| toks.append(s.raps.proj(torch.cat([u,am,u-am,mx]))) |
| x=torch.stack(toks).unsqueeze(0) |
| pidt=torch.tensor([[i+1 for i in ids]],device=s.dev) |
| x=x+s.raps.pe(pidt)+s.raps.pos[:,:n] |
| for blk in s.raps.blocks: x=blk(x) |
| phn=(s.raps.phn(x).squeeze(-1)[0]).cpu().numpy() |
| phone_scores=(phn*10).clip(0,100) |
| |
| sc=np.zeros((1,MAX_PHONES),np.float32); pc=np.zeros((1,MAX_PHONES),np.int64); per=np.zeros((1,MAX_PHONES),np.float32) |
| sc[0,:n]=phn; pc[0,:n]=[i+1 for i in ids] |
| flu=fluency_feats(bu, wids, n)[None] |
| xin=np.stack([sc,per],-1) |
| out=s.agg(torch.from_numpy(xin).to(s.dev), torch.from_numpy(pc).to(s.dev), |
| torch.from_numpy(flu).to(s.dev))[0].cpu().numpy() |
| clip=lambda v: round(float(np.clip(v,0,100)),1) |
| return { |
| 'accuracy': clip(out[0]), |
| 'integrity': clip(out[1]), |
| 'fluency': clip(out[2]), |
| 'overall': clip(out[3]), |
| 'phones': [s.vocab[i] for i in ids], |
| 'phone_scores': [round(float(x),1) for x in phone_scores], |
| } |
|
|
| if __name__=='__main__': |
| ap=argparse.ArgumentParser() |
| ap.add_argument('--user',required=True); ap.add_argument('--ref',required=True) |
| ap.add_argument('--text',required=True); ap.add_argument('--gpu',type=int,default=0) |
| a=ap.parse_args() |
| sc=ReferenceScorer(device=f'cuda:{a.gpu}') |
| print(json.dumps(sc.score(a.user,a.ref,a.text),indent=2,ensure_ascii=False)) |
|
|
