| """Portable 560ms explicit-cache EN reference runtime for bundled LiteRT graphs.""" |
| from pathlib import Path |
| import json,time,hashlib |
| import numpy as np |
|
|
| class Graph: |
| def __init__(self,path,threads=4): |
| from ai_edge_litert.interpreter import Interpreter |
| self.interpreter=Interpreter(model_path=str(path),num_threads=threads) |
| self.runner=self.interpreter.get_signature_runner('serving_default') |
| def __call__(self,*args): |
| result=self.runner(**{f'args_{i}':a for i,a in enumerate(args)}) |
| outputs=tuple(result[k] for k in sorted(result,key=lambda k:int(k.rsplit('_',1)[1]))) |
| for out in outputs: |
| if not np.isfinite(out).all():raise ValueError('Nonfinite graph output') |
| return outputs |
|
|
| class Runtime: |
| def __init__(self,bundle,threads=4,allow_experimental=False,validate_manifest=True): |
| bundle=Path(bundle) |
| if validate_manifest: |
| manifest=json.loads((bundle/'nemotron-manifest.json').read_text()) |
| if manifest['adapter_id']!='nemotron-en-hf560-v1' or manifest['blank_id']!=1024:raise ValueError('Wrong model adapter') |
| if manifest.get('evaluation_only') and not allow_experimental:raise ValueError('Experimental bundle: inspect parity report and explicitly allow evaluation') |
| for entry in manifest['files']: |
| path=bundle/entry['path'] |
| if path.parent!=bundle or not path.is_file():raise ValueError('Invalid or missing bundle file') |
| if path.stat().st_size!=entry['bytes'] or hashlib.sha256(path.read_bytes()).hexdigest()!=entry['sha256']:raise ValueError('Bundle checksum mismatch: '+path.name) |
| self.graphs={p.stem:Graph(p,threads) for p in Path(bundle).glob('*.tflite')} |
| required={'frontend_initial','frontend','subsampling_initial','subsampling','projection','predictor','joint'}|{f'block_{i:02}' for i in range(24)} |
| if set(self.graphs)!=required:raise ValueError('Incomplete or unexpected EN component inventory') |
| def transcribe(self,pcm): |
| if pcm.ndim!=1 or not np.isfinite(pcm).all():raise ValueError('Expected finite mono float32 PCM at 16kHz') |
| z=lambda *s:np.zeros(s,np.float32) |
| states=[z(1,1,1,131),z(1,256,1,68),z(1,256,1,36)] |
| blocks=[[z(1,8,70,128),z(1,8,70,128),z(1,1024,8)] for _ in range(24)] |
| g=self.graphs;tokens=[];emissions=[];timings={};step=0;offset=0 |
| def run(name,*args): |
| begun=time.monotonic();out=g[name](*args);timings[name]=timings.get(name,0)+time.monotonic()-begun;return out |
| pred,h,c=run('predictor',np.array([[1024]],np.int64),z(2,1,640),z(2,1,640)) |
| while offset<len(pcm)//160: |
| frames=49 if step==0 else 56;length=frames*160+353;begin=offset*160-257 |
| wave=z(1,length);a=max(0,begin);b=min(len(pcm),begin+length) |
| if b>a:wave[0,a-begin:b-begin]=pcm[a:b] |
| valid=np.array([min(frames,max(0,len(pcm)//160-offset))],np.int32) |
| samples=np.array([min(length-1,max(0,len(pcm)-begin-1))],np.int32) |
| mel,=run('frontend_initial' if step==0 else 'frontend',wave,valid,samples) |
| x,*states=run('subsampling_initial' if step==0 else 'subsampling',mel,*states) |
| for i in range(24):x,*blocks[i]=run(f'block_{i:02}',x,*blocks[i],np.array([min(step*7,70)],np.int32)) |
| projected,=run('projection',x) |
| for f in range(7): |
| for symbol in range(10): |
| logits,=run('joint',projected[:,f:f+1],pred);token=int(logits.reshape(-1).argmax()) |
| if token==1024:break |
| tokens.append(token);emissions.append((step*7+f)*.08) |
| if len(tokens)>100000:raise ValueError('Token limit') |
| pred,h,c=run('predictor',np.array([[token]],np.int64),h,c) |
| else:raise ValueError('No blank within 10 evaluations on one encoder frame') |
| offset+=frames;step+=1 |
| return dict(tokens=tokens,emission_seconds=emissions,compute_seconds=sum(timings.values()),timings=timings,steps=step) |
|
|
|
|
| def decode(tokens,vocabulary): |
| text=''.join(vocabulary['pieces'][t] for t in tokens if t not in vocabulary['special_ids']).replace('▁',' ') |
| return text[1:] if text.startswith(' ') else text |
|
|
|
|
| if __name__=='__main__': |
| import argparse,math |
| import soundfile as sf |
| from scipy.signal import resample_poly |
| parser=argparse.ArgumentParser();parser.add_argument('bundle',type=Path);parser.add_argument('audio',type=Path);parser.add_argument('--threads',type=int,default=4);parser.add_argument('--allow-experimental',action='store_true');args=parser.parse_args() |
| wave,sr=sf.read(args.audio,dtype='float32',always_2d=True);divisor=math.gcd(sr,16000) |
| pcm=resample_poly(wave.mean(axis=1),16000//divisor,sr//divisor).astype('float32') |
| runtime=Runtime(args.bundle,args.threads,args.allow_experimental) |
| result=runtime.transcribe(pcm);vocabulary=json.loads((args.bundle/'vocabulary.json').read_text()) |
| result['text']=decode(result['tokens'],vocabulary) |
| print(json.dumps(result,ensure_ascii=False,indent=2)) |
|
|
