| """Hybrid prototype: real CoreML encoder -> host array -> MLX decoder.""" |
| import argparse,json,time,sys,wave |
| from pathlib import Path |
| import numpy as np |
| import coremltools as ct |
| import mlx.core as mx |
| from decoder import Decoder |
| from frontend import features |
|
|
| p=argparse.ArgumentParser();p.add_argument('--root',required=True);p.add_argument('--weights');p.add_argument('--fixture');p.add_argument('--audio');p.add_argument('--language',default='hi');p.add_argument('--mixed',action='store_true');p.add_argument('--runs',type=int,default=3);p.add_argument('--report',default='hybrid_mlx');a=p.parse_args() |
| root=Path(a.root);weights=Path(a.weights) if a.weights else root/'experiments/mlx-decoder' |
| tokenizer=json.loads((root/'native-assets/tokenizer.json').read_text()) |
| fixture_path=Path(a.fixture) if a.fixture else root/'baseline/fixture_0.npz' |
| fixture=np.load(fixture_path) if fixture_path.exists() else None |
| if fixture is None and not a.audio:raise ValueError('Provide --audio (16 kHz float32 or PCM16 WAV) or a validation --fixture') |
| (root/'reports').mkdir(exist_ok=True) |
| start=time.perf_counter();decoder=Decoder(weights);decoder_load=time.perf_counter()-start |
| |
| parity=None |
| if fixture is not None: |
| source_enc=mx.array(fixture['encoder'],dtype=mx.float16);cross=decoder.cross_cache(source_enc) |
| logits,cache=decoder.step(mx.array(fixture['prompt'],dtype=mx.int32),0,cross,None);mx.eval(logits,cache) |
| logits=np.array(logits).astype(np.float32);reference=fixture['logits'];delta=np.abs(logits-reference) |
| parity={'logitsMaxAbs':float(delta.max()),'logitsMeanAbs':float(delta.mean()),'lastTokenEqual':int(logits[0,-1].argmax())==int(reference[0,-1].argmax())} |
| print('source_decoder_parity',parity,flush=True) |
| if not parity['lastTokenEqual'] or not np.isfinite(logits).all():raise RuntimeError('Decoder parity check failed') |
| del source_enc,cross,cache,logits |
| start=time.perf_counter() |
| compiled=root/'coreml/encoder.mlmodelc' |
| if not compiled.exists(): |
| ct.models.utils.compile_model(str(root/'coreml/encoder.mlpackage'),destination_path=str(compiled)) |
| encoder=ct.models.CompiledMLModel(str(compiled),compute_units=ct.ComputeUnit.CPU_AND_GPU) |
| encoder_load=time.perf_counter()-start |
| if a.audio: |
| audio=Path(a.audio) |
| if audio.suffix.lower()=='.wav': |
| with wave.open(str(audio)) as w: |
| assert w.getframerate()==16000 and w.getnchannels()==1 and w.getsampwidth()==2,'Use 16 kHz mono PCM16 WAV' |
| samples=np.frombuffer(w.readframes(w.getnframes()),dtype='<i2').astype(np.float32)/32768 |
| else:samples=np.fromfile(audio,dtype='<f4') |
| else:samples=fixture['audio'].reshape(-1) |
| if not 0<len(samples)<=30*16000:raise ValueError('Prototype accepts 0–30 seconds per call') |
| results=[] |
| for _ in range(a.runs): |
| start=time.perf_counter();mel=features(samples,root/'native-assets/frontend.bin');frontend_seconds=time.perf_counter()-start |
| start=time.perf_counter();out=encoder.predict({'features':mel,'mask':np.ones((1,mel.shape[-1]),dtype=np.int32)});encoder_seconds=time.perf_counter()-start |
| length=int(np.asarray(out['lengths']).reshape(-1)[0]);acoustic=np.asarray(out['encoder'])[:,:length,:] |
| start=time.perf_counter();acoustic_mlx=mx.array(acoustic,dtype=mx.float16);mx.eval(acoustic_mlx);transfer_seconds=time.perf_counter()-start |
| result=decoder.generate(acoustic_mlx,tokenizer,a.language,a.mixed) |
| result.update(frontendSeconds=frontend_seconds,encoderSeconds=encoder_seconds,transferSeconds=transfer_seconds,audioSeconds=len(samples)/16000,encoderPolicy='cpuAndGPU',decoderRuntime='mlx-fp16',peakMLXBytes=mx.get_peak_memory()) |
| result['totalSeconds']=sum(result[k] for k in ['frontendSeconds','encoderSeconds','transferSeconds','crossSeconds','decodeSeconds']) |
| results.append(result);print(json.dumps(result,ensure_ascii=False),flush=True) |
| report={'decoderLoadSeconds':decoder_load,'encoderLoadSeconds':encoder_load,'sourceDecoderParity':parity,'runs':results} |
| if not a.audio:report['frontendMaxAbs']=float(np.abs(mel-fixture['features']).max()) |
| (root/'reports'/f'{a.report}.json').write_text(json.dumps(report,indent=2,ensure_ascii=False)) |
|
|
