File size: 5,129 Bytes
fa9d230 | 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 | """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))
|