358 lines
18 KiB
Python
358 lines
18 KiB
Python
"""Run only on 2033lab04ztr; point scoring without reply/timeout gates.
|
|
|
|
Primary score: softmax(complete, incomplete)[complete] at the latest causal
|
|
chunk. It is a conditional two-token score, not a calibrated probability of
|
|
replying. We also retain the native five-state argmax and official corrected
|
|
state. ASR is the model's own Paraformer, never the annotation transcript.
|
|
"""
|
|
import argparse
|
|
from collections import Counter
|
|
import csv
|
|
from datetime import datetime, timezone
|
|
import hashlib
|
|
import importlib.metadata
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import platform
|
|
import random
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
ROOT = Path('/mnt/58462F4C462F29E8/xiaobailing-turn-detection-experiments')
|
|
REPO = ROOT / 'SoulX-Duplug'
|
|
OUT = ROOT / 'results/work-data-can-reply'
|
|
for name, value in {
|
|
'HF_ENDPOINT': 'https://hf-mirror.com',
|
|
'HF_HOME': str(ROOT / 'cache/huggingface'),
|
|
'MODELSCOPE_CACHE': str(ROOT / 'cache/modelscope'),
|
|
'XDG_CACHE_HOME': str(ROOT / 'cache/xdg'),
|
|
'TMPDIR': str(ROOT / 'cache/tmp'),
|
|
'TOKENIZERS_PARALLELISM': 'false',
|
|
'TQDM_DISABLE': '1',
|
|
}.items():
|
|
os.environ[name] = value
|
|
os.chdir(REPO)
|
|
sys.path.insert(0, str(REPO))
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import soxr
|
|
import torch
|
|
from service.model import load_turn_model
|
|
from model.model import State_Prediction_Model
|
|
|
|
TOKENS = ('<|user_complete|>', '<|user_incomplete|>', '<|user_idle|>',
|
|
'<|user_nonidle|>', '<|user_backchannel|>')
|
|
|
|
|
|
def sha(path):
|
|
digest = hashlib.sha256()
|
|
with Path(path).open('rb') as f:
|
|
for b in iter(lambda: f.read(8*1024*1024), b''):
|
|
digest.update(b)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def read(path):
|
|
with Path(path).open(encoding='utf-8-sig', newline='') as f:
|
|
return list(csv.DictReader(f))
|
|
|
|
|
|
def write(path, rows):
|
|
if not rows:
|
|
return
|
|
Path(path).parent.mkdir(parents=True, exist_ok=True)
|
|
with Path(path).open('w', encoding='utf-8-sig', newline='') as f:
|
|
w = csv.DictWriter(f, fieldnames=list(rows[0]))
|
|
w.writeheader()
|
|
w.writerows(rows)
|
|
|
|
|
|
def save(path, data):
|
|
Path(path).parent.mkdir(parents=True, exist_ok=True)
|
|
Path(path).write_text(json.dumps(data, ensure_ascii=False, indent=2)+'\n')
|
|
|
|
|
|
def metrics(rows, threshold=.5, native=False):
|
|
c = Counter((int(r['can_reply_now']), int(r['raw_state'] == TOKENS[0])
|
|
if native else int(r['probability'] > threshold)) for r in rows)
|
|
tn, fp, fn, tp = c[0,0], c[0,1], c[1,0], c[1,1]
|
|
ratio=lambda n,d:n/d if d else None
|
|
return dict(n=tn+fp+fn+tp, tn=tn, fp=fp, fn=fn, tp=tp,
|
|
accuracy=ratio(tn+tp,tn+fp+fn+tp), false_reply_rate=ratio(fp,fp+tn),
|
|
false_hold_rate=ratio(fn,fn+tp))
|
|
|
|
|
|
def stats(values):
|
|
a = np.asarray(values, dtype=float)
|
|
return dict(n=len(a), mean_ms=float(a.mean()), median_ms=float(np.median(a)),
|
|
p95_ms=float(np.percentile(a,95)), p99_ms=float(np.percentile(a,99)),
|
|
min_ms=float(a.min()), max_ms=float(a.max())) if len(a) else None
|
|
|
|
|
|
class Instrument:
|
|
def __init__(self, model):
|
|
self.model = model
|
|
self.phase = ''
|
|
self.ids = [model.model.tokenizer.convert_tokens_to_ids(t) for t in TOKENS]
|
|
assert len(set(self.ids)) == 5
|
|
self.latest = {}
|
|
self.parts = {}
|
|
self.asr_calls = 0
|
|
model.model.llm.register_forward_hook(self.capture)
|
|
for method, phase in (('_audio_to_tokens','audio_tokens'), ('_asr','asr_stage'),
|
|
('_state_predict','state_model')):
|
|
original = getattr(model, method)
|
|
def wrapped(*args, _fn=original, _phase=phase, **kwargs):
|
|
previous = self.phase
|
|
self.phase = _phase
|
|
torch.cuda.synchronize()
|
|
begin = time.perf_counter()
|
|
try:
|
|
return _fn(*args, **kwargs)
|
|
finally:
|
|
torch.cuda.synchronize()
|
|
self.parts[_phase+'_ms'] = (time.perf_counter()-begin)*1000
|
|
self.phase = previous
|
|
setattr(model, method, wrapped)
|
|
original_asr = model.cascade_asr.recognize
|
|
def recognize(*args, **kwargs):
|
|
torch.cuda.synchronize()
|
|
began = time.perf_counter()
|
|
result = original_asr(*args, **kwargs)
|
|
torch.cuda.synchronize()
|
|
self.parts['cascade_asr_ms'] += (time.perf_counter()-began)*1000
|
|
self.asr_calls += 1
|
|
return result
|
|
model.cascade_asr.recognize = recognize
|
|
|
|
def capture(self, _module, _args, output):
|
|
if self.phase != 'state_model':
|
|
return
|
|
full = output.logits[0,-1].detach()
|
|
selected = full[self.ids].float().cpu().numpy()
|
|
pair = torch.softmax(full[self.ids[:2]].float(), dim=0).cpu().numpy()
|
|
self.latest = dict(probability=float(pair[0]),
|
|
raw_state=self.model.model.tokenizer.decode(int(full.argmax().item())),
|
|
state_logits=dict(zip(TOKENS, map(float,selected))))
|
|
|
|
def reset(self):
|
|
np.random.seed(42)
|
|
random.seed(42)
|
|
torch.manual_seed(42)
|
|
torch.cuda.manual_seed_all(42)
|
|
self.model.reset()
|
|
|
|
def infer(self, current, back, ahead):
|
|
self.parts = dict(cascade_asr_ms=0.)
|
|
self.asr_calls = 0
|
|
torch.cuda.synchronize()
|
|
begin = time.perf_counter()
|
|
# Direct infer retains native KV/state prediction. It bypasses the
|
|
# service's far-field filter, timeout replies, complete-trigger resets,
|
|
# and final utterance ASR, none of which define this point classifier.
|
|
state, delta, text = self.model.infer(current, back, ahead)
|
|
torch.cuda.synchronize()
|
|
total = (time.perf_counter()-begin)*1000
|
|
return dict(self.latest, state=state, asr_delta=delta, asr_buffer=text,
|
|
total_ms=total, asr_calls=self.asr_calls, **self.parts)
|
|
|
|
|
|
def load():
|
|
torch.set_num_threads(4)
|
|
torch.set_num_interop_threads(1)
|
|
torch.cuda.reset_peak_memory_stats()
|
|
began = time.perf_counter()
|
|
loading_diagnostics=[]
|
|
original_load=State_Prediction_Model.load_state_dict
|
|
def checked_load(instance,state_dict,*args,**kwargs):
|
|
result=original_load(instance,state_dict,*args,**kwargs)
|
|
assert not result.missing_keys, f'Missing trained model weights: {result.missing_keys}'
|
|
assert set(result.unexpected_keys) <= {'embed_tokens_func.weight'}, result.unexpected_keys
|
|
if 'embed_tokens_func.weight' in result.unexpected_keys:
|
|
alias=state_dict['embed_tokens_func.weight']
|
|
matching=[v for k,v in state_dict.items() if k.endswith('embed_tokens.weight')
|
|
and v.shape==alias.shape and torch.equal(v,alias)]
|
|
assert matching, 'Unexpected embedding weight is not a duplicate alias'
|
|
loading_diagnostics.append(dict(missing=result.missing_keys,unexpected=result.unexpected_keys))
|
|
return result
|
|
State_Prediction_Model.load_state_dict=checked_load
|
|
model = load_turn_model(str(REPO/'config/config.yaml'))
|
|
State_Prediction_Model.load_state_dict=original_load
|
|
torch.cuda.synchronize()
|
|
meta = dict(host='2033lab04ztr', hostname=platform.node(), python=sys.version,
|
|
started_at_utc=datetime.now(timezone.utc).isoformat(),
|
|
gpu_driver=subprocess.check_output(['nvidia-smi','--query-gpu=driver_version',
|
|
'--format=csv,noheader'],text=True).strip(),
|
|
gpu=torch.cuda.get_device_name(0), cpu_threads=4, torch=torch.__version__,
|
|
cuda=torch.version.cuda, load_ms=(time.perf_counter()-began)*1000,
|
|
parameter_dtypes=sorted({str(p.dtype) for p in model.model.parameters()}),
|
|
declared_precision=model.config.infer_config.precision,
|
|
actual_precision_note='Official loader does not cast to bf16; no autocast added.',
|
|
checkpoint_loading=loading_diagnostics,
|
|
repo_commit=subprocess.check_output(['git','-c',f'safe.directory={REPO}',
|
|
'rev-parse','HEAD'],cwd=REPO,text=True).strip(),
|
|
model_revision='61701c4ab8193cc1ee2220d3848872ac6c720142',
|
|
model_checksums=json.loads((ROOT/'results/model-download.json').read_text()),
|
|
label_sha256=sha(ROOT/'data/labels/eot_labels.csv'),
|
|
script_sha256=sha(__file__),
|
|
dependencies={k:importlib.metadata.version(k) for k in ('transformers',
|
|
'funasr','modelscope','peft','pytorch-lightning','numpy','soxr')})
|
|
return Instrument(model), meta
|
|
|
|
|
|
def point_audio(pcm, sr, decision):
|
|
# Target chunk ends 40ms before D; its 40ms lookahead ends exactly at D.
|
|
# <=20s total real audio. Cropping precedes resampling, including last chunk.
|
|
frames = (min(20000, decision)-40)//160
|
|
if frames < 1:
|
|
raise ValueError('Not enough past audio for a causal chunk')
|
|
origin = decision-40-frames*160
|
|
start, end = origin*sr//1000, decision*sr//1000
|
|
assert 0 <= start < end <= len(pcm)
|
|
began = time.perf_counter()
|
|
audio = np.ascontiguousarray(soxr.resample(pcm[start:end,0],sr,16000),dtype=np.float32)
|
|
resample_ms = (time.perf_counter()-began)*1000
|
|
assert len(audio) == frames*2560+640
|
|
return audio, origin, frames, resample_ms
|
|
|
|
|
|
def run_points(instrument, meta, smoke=False):
|
|
labels = read(ROOT/'data/labels/eot_labels.csv')
|
|
if smoke:
|
|
# A real speech checkpoint, not silence alone.
|
|
labels = [next(r for r in labels if r['call_id'].startswith('872a1e67')
|
|
and int(r['decision_ms']) == 17153)]
|
|
audio_cache, rows, traces = {}, [], []
|
|
began = time.perf_counter()
|
|
for index, label in enumerate(labels):
|
|
call, decision = label['call_id'], int(label['decision_ms'])
|
|
if call not in audio_cache:
|
|
path = ROOT/'data/audio'/f'{call}.mp3'
|
|
pcm, sr = sf.read(path,dtype='float32',always_2d=True)
|
|
assert pcm.shape[1] == 2
|
|
audio_cache[call] = pcm,sr
|
|
pcm,sr = audio_cache[call]
|
|
audio, origin, frames, resample_ms = point_audio(pcm,sr,decision)
|
|
instrument.reset()
|
|
# Official reset's tiny noise fills unavailable initial backward context.
|
|
back_initial = instrument.model.buffer[:15360].astype(np.float32)
|
|
padded = np.concatenate([back_initial,audio])
|
|
replay_ms = 0.
|
|
for j in range(frames):
|
|
start = 15360+j*2560
|
|
result = instrument.infer(padded[start:start+2560],
|
|
padded[start-15360:start],padded[start+2560:start+3200])
|
|
replay_ms += result['total_ms']
|
|
traces.append(dict(call_id=call,decision_ms=decision,frame=j,
|
|
current_end_ms=origin+(j+1)*160,
|
|
input_available_ms=origin+(j+1)*160+40,
|
|
state=result['state'],probability=result['probability'],
|
|
total_ms=result['total_ms'],audio_tokens_ms=result['audio_tokens_ms'],
|
|
asr_stage_ms=result['asr_stage_ms'],state_model_ms=result['state_model_ms'],
|
|
cascade_asr_ms=result['cascade_asr_ms'],asr_calls=result['asr_calls']))
|
|
prediction = int(result['probability'] > .5)
|
|
expected = int(label['can_reply_now'])
|
|
r = dict(label, probability=result['probability'], prediction=prediction,
|
|
correct=int(prediction==expected),
|
|
error_type='correct' if prediction==expected else ('false_reply' if prediction else 'false_hold'),
|
|
raw_state=result['raw_state'], corrected_state=result['state'],
|
|
state_logits=json.dumps(result['state_logits']), asr_text=result['asr_buffer'],
|
|
audio_start_ms=origin,audio_end_ms=decision,current_chunk_end_ms=decision-40,
|
|
lookahead_ms=40,frame_count=frames,context_ms=decision-origin,
|
|
replay_inference_ms=replay_ms,last_incremental_ms=result['total_ms'],
|
|
last_cascade_asr_ms=result['cascade_asr_ms'],resample_ms=resample_ms,
|
|
audio_channel=0,label_status='draft' if label['annotator']=='xlsx_draft' else 'unconfirmed_manual')
|
|
rows.append(r)
|
|
print(f"POINT {index+1}/{len(labels)} {call[:8]} {decision} label={expected} "
|
|
f"p_complete={r['probability']:.5f} raw={r['raw_state']} "
|
|
f"ASR={r['asr_text']!r} replay_ms={replay_ms:.1f}",flush=True)
|
|
# Partial results survive interrupted runs and are clearly incomplete.
|
|
destination = ROOT/'results/smoke' if smoke else OUT
|
|
write(destination/'predictions.csv',rows)
|
|
thresholds = [dict(threshold=t,**metrics(rows,t)) for t in (.1,.2,.3,.4,.5,.6,.7,.8,.9,.95,.99)]
|
|
summary = dict(total_points=len(rows),complete=len(rows)==(1 if smoke else 44),
|
|
metrics=metrics(rows),native_state_metrics=metrics(rows,native=True),
|
|
corrected_state_metrics=metrics([dict(r,raw_state=r['corrected_state']) for r in rows],native=True),
|
|
thresholds=thresholds,metadata=meta,
|
|
grouped={k:{v:metrics([r for r in rows if r[k]==v]) for v in sorted({r[k] for r in rows})}
|
|
for k in ('call_id','label_status','reason')},
|
|
elapsed_seconds=time.perf_counter()-began,
|
|
peak_gpu_allocated_bytes=torch.cuda.max_memory_allocated(),
|
|
peak_gpu_reserved_bytes=torch.cuda.max_memory_reserved(),
|
|
audio_sha256={call:sha(ROOT/'data/audio'/f'{call}.mp3') for call in audio_cache},
|
|
config=dict(sample_rate=16000,chunk_ms=160,back_ms=960,ahead_ms=40,
|
|
max_real_context_ms=20000,channel=0,threshold=.5,
|
|
rule='softmax(logit_complete,logit_incomplete)[complete] > threshold',
|
|
state_source='direct infer/_state_predict, no service action/timeout/reset gates'),
|
|
method='Each point resets native KV state, replays <=20s past user audio with built-in Paraformer. Final process chunk ends D-40ms and lookahead ends D. No label transcript input, no future audio, no fallback reply or complete-trigger reset.',
|
|
limitations=['Conditional two-token score is not a calibrated can_reply probability; raw five-state argmax retained as a diagnostic.',
|
|
'29 draft labels and 15 labels without verification metadata; selected-point exploratory evaluation.',
|
|
'Input channels/context and hardware differ across models; this is a configured-system comparison.',
|
|
'Independent point restart differs from continuous server state; ASR is part of SoulX inference.',
|
|
'Known ASR/label alignment issues at cb3238dc 13540 and e304f740 11531 remain unchanged.'])
|
|
save(destination/'summary.json',summary)
|
|
write(destination/'errors.csv',[r for r in rows if not r['correct']])
|
|
write(destination/'thresholds.csv',thresholds)
|
|
write(destination/'frame-times.csv',traces)
|
|
print(json.dumps(summary['metrics']),flush=True)
|
|
|
|
|
|
def benchmark(instrument,meta):
|
|
labels = read(ROOT/'data/labels/eot_labels.csv')
|
|
rows = []
|
|
for call in sorted({r['call_id'] for r in labels}):
|
|
decision = max(int(r['decision_ms']) for r in labels if r['call_id']==call)
|
|
pcm,sr = sf.read(ROOT/'data/audio'/f'{call}.mp3',dtype='float32',always_2d=True)
|
|
# No audio beyond the last checkpoint. Real continuous frame grid begins
|
|
# at call time 0; the final incomplete frame is not padded into the future.
|
|
audio = np.ascontiguousarray(soxr.resample(pcm[:decision*sr//1000,0],sr,16000),dtype=np.float32)
|
|
count = (len(audio)-640)//2560
|
|
instrument.reset()
|
|
padded = np.concatenate([instrument.model.buffer[:15360].astype(np.float32),audio])
|
|
for j in range(count):
|
|
p = 15360+j*2560
|
|
r = instrument.infer(padded[p:p+2560],padded[p-15360:p],padded[p+2560:p+3200])
|
|
rows.append(dict(call_id=call,frame=j,current_end_ms=(j+1)*160,
|
|
steady=int((j+1)*160>20000),total_ms=r['total_ms'],
|
|
audio_tokens_ms=r['audio_tokens_ms'],asr_stage_ms=r['asr_stage_ms'],
|
|
state_model_ms=r['state_model_ms'],cascade_asr_ms=r['cascade_asr_ms'],
|
|
non_cascade_ms=r['total_ms']-r['cascade_asr_ms'],asr_calls=r['asr_calls']))
|
|
if (j+1)%100==0:
|
|
print(f'BENCH {call[:8]} {j+1}/{count}',flush=True)
|
|
write(ROOT/'results/latency-frames.csv',rows)
|
|
fields=('total_ms','audio_tokens_ms','asr_stage_ms','state_model_ms','cascade_asr_ms','non_cascade_ms')
|
|
groups={'all':rows,'steady':[r for r in rows if r['steady']],
|
|
'asr_called':[r for r in rows if r['asr_calls']],
|
|
'no_asr':[r for r in rows if not r['asr_calls']]}
|
|
data=dict(metadata=meta,complete=True,update_interval_ms=160,lookahead_ms=40,
|
|
summary={g:{f:stats([r[f] for r in data]) for f in fields} for g,data in groups.items()},
|
|
realtime_factor=sum(r['total_ms'] for r in rows)/(len(rows)*160),
|
|
fraction_over_160ms=sum(r['total_ms']>160 for r in rows)/len(rows),
|
|
peak_gpu_allocated_bytes=torch.cuda.max_memory_allocated(),
|
|
peak_gpu_reserved_bytes=torch.cuda.max_memory_reserved(),
|
|
note='CUDA synchronized wall time. Includes waveform features, audio tokenizer, LLM state, cascade ASR and instrumentation. Excludes imports/model loading, file decode/resample, network, reply policy. Official FP32 loading retained; BF16 config does not cast weights.')
|
|
save(ROOT/'results/latency.json',data)
|
|
print(json.dumps(data['summary']['steady']['total_ms']),flush=True)
|
|
|
|
|
|
if __name__=='__main__':
|
|
p=argparse.ArgumentParser()
|
|
p.add_argument('mode',choices=('smoke','evaluate','benchmark','all'))
|
|
args=p.parse_args()
|
|
instrument,meta=load()
|
|
save(ROOT/'results'/f'{args.mode}-environment.json',meta)
|
|
if args.mode=='all':
|
|
run_points(instrument,meta,smoke=True)
|
|
print('SMOKE_COMPLETE',flush=True)
|
|
run_points(instrument,meta)
|
|
print('EVALUATION_COMPLETE',flush=True)
|
|
benchmark(instrument,meta)
|
|
print('ALL_COMPLETE',flush=True)
|
|
elif args.mode=='benchmark':
|
|
benchmark(instrument,meta)
|
|
else:
|
|
run_points(instrument,meta,smoke=args.mode=='smoke')
|