File size: 8,043 Bytes
21fd722 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """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)
# Prepared before Ready, as in the ordinary CLI. Future segment resets are
# included in the audio clock and measured separately from model windows.
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()
|