xiaobailing-turn-detection/experiments/soulx-duplug/run_experiment.py

"""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')