| """Live speech-gated MLX captions using actual causal CPU VAD decisions.""" |
| import hashlib |
| import json |
| from pathlib import Path |
| import struct |
| import sys |
| import time |
| import mlx.core as mx |
| from .silero_input import GatedPCM,Silero,MODEL_SHA256,POLICY |
| from .stream import TextDecoder,require |
| from .weights import sha256 |
|
|
|
|
| def run(args,model,frontend,profile,load_start,warmup_seconds,warmup_calls): |
| from .cli import emit,memory |
| import importlib.metadata |
| vad=Silero(args.silero_vad) |
| |
| |
| session=model.session(delay_tokens=3) |
| tokenizer=TextDecoder(args.tokenizer) |
| mx.synchronize() |
| stream=sys.stdin.buffer if args.input=='-' else open(args.input,'rb') |
| emit('caption_stream_ready',backend='mlx_metal',mlx=importlib.metadata.version('mlx'), |
| weights=model.weights.provenance,weight_resident_bytes=model.weights.nbytes, |
| config_sha256=sha256(args.config),frontend=frontend.provenance,tokenizer_sha256=tokenizer.table_sha256, |
| source_sha256={p.name:sha256(p) for p in sorted(Path(__file__).parent.glob('*.py'))}, |
| silero_sha256=MODEL_SHA256,silero_runtime='onnxruntime_cpu',silero_runtime_version=vad.version,silence_policy=POLICY, |
| startup='realtime18',delay_ms=240,batch=args.batch_windows,first_batch=1,dtype=args.dtype,cache_dtype=args.cache_dtype, |
| math_mode=args.math_mode,fuse_projections=args.fuse_projections,projection_warmup_seconds=warmup_seconds, |
| projection_warmup_calls=warmup_calls,input_mode='stdin' if args.input=='-' else ('paced_file' if args.pace_audio else 'unpaced_file'), |
| load_seconds=time.monotonic()-load_start,memory=memory(),production_accuracy_approved=False) |
| started=time.monotonic();source=GatedPCM(stream,vad.probability,started if args.pace_audio else None) |
| ids=hashlib.sha256();emissions=0;total_work=total_reset=0.;text_bytes=0;peak_cache=0 |
| prior_visible=segment_visible=False;active=None;local_index=0;previous=None |
| def text_event(delta,segment,**fields): |
| nonlocal prior_visible,segment_visible,text_bytes |
| if not delta:return |
| seconds=time.monotonic()-started |
| if not segment_visible and prior_visible: |
| emit('caption_text_delta',segment=segment,index=emissions,seconds=seconds,text='\n',segment_separator=True) |
| text_bytes+=1 |
| segment_visible=prior_visible=True;text_bytes+=len(delta.encode('utf8')) |
| emit('caption_text_delta',segment=segment,index=emissions,seconds=seconds,text=delta,**fields) |
| try: |
| events=iter(source);pending=None |
| while True: |
| event=pending if pending is not None else next(events) |
| pending=None |
| if event['kind']=='start': |
| s=event['segment'];require(active is None,'overlapping MLX caption session') |
| active=s['id'];local_index=0;previous=None;segment_visible=False |
| began=time.monotonic() |
| if active: |
| session=None;mx.clear_cache();session=model.session(delay_tokens=3) |
| tokenizer=tokenizer.fresh() |
| reset=time.monotonic()-began;total_reset+=reset |
| require(session.position==0 and all(c.length==0 for c in session.encoder+session.decoder),'fresh MLX caches') |
| emit('caption_segment_start',segment=active,start_sample=s['start_sample'],decision_sample=s['decision_sample'], |
| seconds=time.monotonic()-started,reset_seconds=reset) |
| elif event['kind']=='window': |
| require(active==event['segment']['id'] and local_index==event['index'],'MLX caption source clock') |
| inputs=profile.prompt() if local_index==0 else [previous] |
| capacity=session.batch_capacity(1 if local_index==0 else args.batch_windows,len(inputs)) |
| group=[event] |
| while len(group)<capacity: |
| following=next(events) |
| if following['kind']!='window': |
| pending=following;break |
| require(following['segment']['id']==active and following['index']==local_index+len(group),'batch crosses segment') |
| group.append(following) |
| began=time.monotonic();available=max(w['available_sample'] for w in group);ready=available/16000 |
| queue=max(0.,began-started-ready) if args.pace_audio or args.input=='-' else None |
| first_local=local_index;first_global=emissions;last_callback=began |
| def token_ready(offset,token): |
| nonlocal previous,peak_cache,emissions,local_index,last_callback |
| require(type(token) is int and 0<=token<tokenizer.model_count and token!=151645,'MLX greedy token contract') |
| require(offset==local_index-first_local and emissions==first_global+offset,'MLX batch callback clock') |
| require(session.position+session.trims*38==18+local_index,'MLX absolute position differs') |
| now=time.monotonic();interval=now-last_callback;last_callback=now |
| previous=token;ids.update(struct.pack('<I',token));peak_cache=max(peak_cache,session.cache_bytes) |
| text_event(tokenizer.push(token),active) |
| w=group[offset] |
| emit('caption_window',index=emissions,segment=active,segment_index=local_index,token_id=token, |
| seconds=time.monotonic()-started,source_dependency_sample=w['nominal_source_end_sample'], |
| available_source_sample=available,eof_drain=w['nominal_source_end_sample']>available,queue_seconds=queue, |
| work_seconds=interval,cache_bytes=session.cache_bytes,rolling_trims=session.trims, |
| batch_first_index=first_global,batch_windows=len(group),batch_offset=offset) |
| emissions+=1;local_index+=1 |
| session.step_batch([w['audio'] for w in group],inputs,on_token=token_ready) |
| require(local_index-first_local==len(group),'MLX missing batch callback') |
| total_work+=time.monotonic()-began |
| elif event['kind']=='end': |
| require(active==event['segment']['id'] and local_index==event['windows'],'MLX caption end clock') |
| text_event(tokenizer.finish(),active,utf8_final_flush=True) |
| emit('caption_segment_end',segment=active,start_sample=event['segment']['start_sample'], |
| end_sample=event['end_sample'],final_eof=event['final_eof'],unconsumed_samples=event['unconsumed_samples'], |
| windows=local_index,seconds=time.monotonic()-started) |
| active=None |
| elif event['kind']=='complete':break |
| require(active is None and source.reader.eof,'MLX caption incomplete EOF') |
| audio_seconds=source.reader.samples/16000 |
| emit('caption_stream_complete',status='complete',audio_samples=source.reader.samples,audio_seconds=audio_seconds, |
| segments=source.segment_count,selected_samples=source.selected_samples,unconsumed_selected_samples=source.unconsumed_samples, |
| emissions=emissions,text_bytes=text_bytes,model_seconds=total_work,reset_seconds=total_reset,vad_seconds=source.vad_seconds, |
| read_seconds=source.read_seconds,sleep_seconds=source.sleep_seconds,elapsed_seconds=time.monotonic()-started, |
| work_rtf=(total_work+total_reset+source.vad_seconds)/audio_seconds if audio_seconds else None, |
| peak_cache_bytes=peak_cache,max_retained_input_samples=source.reader.max_retained_samples, |
| input_sha256=source.reader.digest.hexdigest(),generated_ids_sha256=ids.hexdigest(),memory=memory(), |
| all_source_audio_observed=True,measured_word_aligned_latency=False,production_accuracy_approved=False) |
| finally: |
| if args.input!='-':stream.close() |
|
|
