license import { complex } from '../ops/complex' ; import { tensor } from '../ops/tensor' ; import { NamedTensor , NamedTensorMap } from '../tensor_types' ; import { TypedArray } from '../types' ; import { sizeFromShape } from '../util' ; import { DTYPE_VALUE_SIZE_MAP , ModelArtifacts , ModelArtifactsInfo , ModelJSON , WeightGroup , WeightsManifestConfig , WeightsManifestEntry } from './types' ; /** Number of bytes reserved for the length of the string. (32bit integer). */ const NUM_BYTES_STRING_LENGTH = 4 ; param tensors param group returns throws export async function encodeWeights ( tensors , group ? ) { const specs = [ ] ; const dataPromises = [ ] ; const names = Array . isArray ( tensors ) ? tensors . map ( tensor => tensor . name ) : Object . keys ( tensors ) ; for ( let i = 0 ; i < names . length ; ++ i ) { const name = names [ i ] ; const t = Array . isArray ( tensors ) ? tensors [ i ] . tensor : tensors [ name ] ; if ( t . dtype !== 'float32' && t . dtype !== 'int32' && t . dtype !== 'bool' && t . dtype !== 'string' && t . dtype !== 'complex64' ) { throw new Error ( `Unsupported dtype in weight '${ name }': ${ t . dtype }` ) ; } const spec = { name , shape : t . shape , dtype : t . dtype } ; if ( t . dtype === 'string' ) { const utf8bytes = new Promise < TypedArray > ( async ( resolve ) => { const vals = await t . bytes ( ) as Uint8Array [ ] ; const totalNumBytes = vals . reduce ( ( p , c ) => p + c . length , 0 ) + NUM_BYTES_STRING_LENGTH * vals . length ; const bytes = new Uint8Array ( totalNumBytes ) ; let offset = 0 ; for ( let i = 0 ; i < vals . length ; i ++ ) { const val = vals [ i ] ; const bytesOfLength = new Uint8Array ( new Uint32Array ( [ val . length ] ) . buffer ) ; bytes . set ( bytesOfLength , offset ) ; offset += NUM_BYTES_STRING_LENGTH ; bytes . set ( val , offset ) ; offset += val . length ; } resolve ( bytes ) ; } ) ; dataPromises . push ( utf8bytes ) ; } else { dataPromises . push ( t . data ( ) ) ; } if ( group != null ) { spec . group = group ; } specs . push ( spec ) ; } const tensorValues = await Promise . all ( dataPromises ) ; return { data : concatenateTypedArrays ( tensorValues ) , specs } ; } param buffer param specs return throws export function decodeWeights ( buffer , specs ) { const out = { } ; let float16Decode ; let offset = 0 ; for ( const spec of specs ) { const name = spec . name ; const dtype = spec . dtype ; const shape = spec . shape ; const size = sizeFromShape ( shape ) ; let values ; if ( 'quantization' in spec ) { const quantization = spec . quantization ; if ( quantization . dtype === 'uint8' || quantization . dtype === 'uint16' ) { if ( ! ( 'min' in quantization && 'scale' in quantization ) ) { throw new Error ( `Weight ${ spec . name } with quantization ${ quantization . dtype } ` + `doesn't have corresponding metadata min and scale.` ) ; } } else if ( quantization . dtype === 'float16' ) { if ( dtype !== 'float32' ) { throw new Error ( `Weight ${ spec . name } is quantized with ${ quantization . dtype } ` + `which only supports weights of type float32 not ${ dtype }.` ) ; } } else { throw new Error ( `Weight ${ spec . name } has unknown ` + `quantization dtype ${ quantization . dtype }. ` + `Supported quantization dtypes are: ` + `'uint8', 'uint16', and 'float16'.` ) ; } const quantizationSizeFactor = DTYPE_VALUE_SIZE_MAP [ quantization . dtype ] ; const byteBuffer = buffer . slice ( offset , offset + size * quantizationSizeFactor ) ; const quantizedArray = ( quantization . dtype === 'uint8' ) ? new Uint8Array ( byteBuffer ) : new Uint16Array ( byteBuffer ) ; if ( dtype === 'float32' ) { if ( quantization . dtype === 'uint8' || quantization . dtype === 'uint16' ) { values = new Float32Array ( quantizedArray . length ) ; for ( let i = 0 ; i < quantizedArray . length ; i ++ ) { const v = quantizedArray [ i ] ; values [ i ] = v * quantization . scale + quantization . min ; } } else if ( quantization . dtype === 'float16' ) { if ( float16Decode === undefined ) { float16Decode = getFloat16Decoder ( ) ; } values = float16Decode ( quantizedArray as Uint16Array ) ; } else { throw new Error ( `Unsupported quantization type ${ quantization . dtype } ` + `for weight type float32.` ) ; } } else if ( dtype === 'int32' ) { if ( quantization . dtype !== 'uint8' && quantization . dtype !== 'uint16' ) { throw new Error ( `Unsupported quantization type ${ quantization . dtype } ` + `for weight type int32.` ) ; } values = new Int32Array ( quantizedArray . length ) ; for ( let i = 0 ; i < quantizedArray . length ; i ++ ) { const v = quantizedArray [ i ] ; values [ i ] = Math . round ( v * quantization . scale + quantization . min ) ; } } else { throw new Error ( `Unsupported dtype in weight '${ name }': ${ dtype }` ) ; } offset += size * quantizationSizeFactor ; } else if ( dtype === 'string' ) { const size = sizeFromShape ( spec . shape ) ; values = [ ] ; for ( let i = 0 ; i < size ; i ++ ) { const byteLength = new Uint32Array ( buffer . slice ( offset , offset + NUM_BYTES_STRING_LENGTH ) ) [ 0 ] ; offset += NUM_BYTES_STRING_LENGTH ; const bytes = new Uint8Array ( buffer . slice ( offset , offset + byteLength ) ) ; ( values as Uint8Array [ ] ) . push ( bytes ) ; offset += byteLength ; } } else { const dtypeFactor = DTYPE_VALUE_SIZE_MAP [ dtype ] ; const byteBuffer = buffer . slice ( offset , offset + size * dtypeFactor ) ; if ( dtype === 'float32' ) { values = new Float32Array ( byteBuffer ) ; } else if ( dtype === 'int32' ) { values = new Int32Array ( byteBuffer ) ; } else if ( dtype === 'bool' ) { values = new Uint8Array ( byteBuffer ) ; } else if ( dtype === 'complex64' ) { values = new Float32Array ( byteBuffer ) ; const real = new Float32Array ( values . length / 2 ) ; const image = new Float32Array ( values . length / 2 ) ; for ( let i = 0 ; i < real . length ; i ++ ) { real [ i ] = values [ i * 2 ] ; image [ i ] = values [ i * 2 + 1 ] ; } const realTensor = tensor ( real , shape , 'float32' ) ; const imageTensor = tensor ( image , shape , 'float32' ) ; out [ name ] = complex ( realTensor , imageTensor ) ; realTensor . dispose ( ) ; imageTensor . dispose ( ) ; } else { throw new Error ( `Unsupported dtype in weight '${ name }': ${ dtype }` ) ; } offset += size * dtypeFactor ; } if ( dtype !== 'complex64' ) { out [ name ] = tensor ( values , shape , dtype ) ; } } return out ; } /**
 * Concatenate TypedArrays into an ArrayBuffer.
 */ export function concatenateTypedArrays ( xs ) { if ( xs === null ) { throw new Error ( `Invalid input value: ${ JSON . stringify ( xs ) }` ) ; } let totalByteLength = 0 ; const normalizedXs = [ ] ; xs . forEach ( ( x ) => { totalByteLength += x . byteLength ; normalizedXs . push ( x . byteLength === x . buffer . byteLength ? x : new ( x . constructor as any ) ( x ) ) ; if ( ! ( x as any instanceof Float32Array || x as any instanceof Int32Array || x as any instanceof Uint8Array ) ) { throw new Error ( `Unsupported TypedArray subtype: ${ x . constructor . name }` ) ; } } ) ; const y = new Uint8Array ( totalByteLength ) ; let offset = 0 ; normalizedXs . forEach ( ( x ) => { y . set ( new Uint8Array ( x . buffer ) , offset ) ; offset += x . byteLength ; } ) ; return y . buffer ; } const useNodeBuffer = typeof Buffer !== 'undefined' && ( typeof Blob === 'undefined' || typeof atob === 'undefined' || typeof btoa === 'undefined' ) ; param str returns export function stringByteLength ( str ) { if ( useNodeBuffer ) { return Buffer . byteLength ( str ) ; } return new Blob ( [ str ] ) . size ; } param buffer returns export function arrayBufferToBase64String ( buffer ) { if ( useNodeBuffer ) { return Buffer . from ( buffer ) . toString ( 'base64' ) ; } const buf = new Uint8Array ( buffer ) ; let s = '' ; for ( let i = 0 , l = buf . length ; i < l ; i ++ ) { s += String . fromCharCode ( buf [ i ] ) ; } return btoa ( s ) ; } param str returns export function base64StringToArrayBuffer ( str ) { if ( useNodeBuffer ) { const buf = Buffer . from ( str , 'base64' ) ; return buf . buffer . slice ( buf . byteOffset , buf . byteOffset + buf . byteLength ) ; } const s = atob ( str ) ; const buffer = new Uint8Array ( s . length ) ; for ( let i = 0 ; i < s . length ; ++ i ) { buffer . set ( [ s . charCodeAt ( i ) ] , i ) ; } return buffer . buffer ; } param buffers returns export function concatenateArrayBuffers ( buffers ) { if ( buffers . length === 1 ) { return buffers [ 0 ] ; } let totalByteLength = 0 ; buffers . forEach ( ( buffer ) => { totalByteLength += buffer . byteLength ; } ) ; const temp = new Uint8Array ( totalByteLength ) ; let offset = 0 ; buffers . forEach ( ( buffer ) => { temp . set ( new Uint8Array ( buffer ) , offset ) ; offset += buffer . byteLength ; } ) ; return temp . buffer ; } param path export function basename ( path ) { const SEPARATOR = '/' ; path = path . trim ( ) ; while ( path . endsWith ( SEPARATOR ) ) { path = path . slice ( 0 , path . length - 1 ) ; } const items = path . split ( SEPARATOR ) ; return items [ items . length - 1 ] ; } param artifacts param manifest returns export function getModelJSONForModelArtifacts ( artifacts , manifest ) { const result = { modelTopology : artifacts . modelTopology , format : artifacts . format , generatedBy : artifacts . generatedBy , convertedBy : artifacts . convertedBy , weightsManifest : manifest } ; if ( artifacts . signature != null ) { result . signature = artifacts . signature ; } if ( artifacts . userDefinedMetadata != null ) { result . userDefinedMetadata = artifacts . userDefinedMetadata ; } if ( artifacts . modelInitializer != null ) { result . modelInitializer = artifacts . modelInitializer ; } if ( artifacts . trainingConfig != null ) { result . trainingConfig = artifacts . trainingConfig ; } return result ; } param modelJSON param loadWeights returns export async function getModelArtifactsForJSON ( modelJSON , loadWeights ) { const modelArtifacts = { modelTopology : modelJSON . modelTopology , format : modelJSON . format , generatedBy : modelJSON . generatedBy , convertedBy : modelJSON . convertedBy } ; if ( modelJSON . trainingConfig != null ) { modelArtifacts . trainingConfig = modelJSON . trainingConfig ; } if ( modelJSON . weightsManifest != null ) { const [ weightSpecs , weightData ] = await loadWeights ( modelJSON . weightsManifest ) ; modelArtifacts . weightSpecs = weightSpecs ; modelArtifacts . weightData = weightData ; } if ( modelJSON . signature != null ) { modelArtifacts . signature = modelJSON . signature ; } if ( modelJSON . userDefinedMetadata != null ) { modelArtifacts . userDefinedMetadata = modelJSON . userDefinedMetadata ; } if ( modelJSON . modelInitializer != null ) { modelArtifacts . modelInitializer = modelJSON . modelInitializer ; } return modelArtifacts ; } param modelArtifacts returns export function getModelArtifactsInfoForJSON ( modelArtifacts ) { if ( modelArtifacts . modelTopology instanceof ArrayBuffer ) { throw new Error ( 'Expected JSON model topology, received ArrayBuffer.' ) ; } return { dateSaved : new Date ( ) , modelTopologyType : 'JSON' , modelTopologyBytes : modelArtifacts . modelTopology == null ? 0 : stringByteLength ( JSON . stringify ( modelArtifacts . modelTopology ) ) , weightSpecsBytes : modelArtifacts . weightSpecs == null ? 0 : stringByteLength ( JSON . stringify ( modelArtifacts . weightSpecs ) ) , weightDataBytes : modelArtifacts . weightData == null ? 0 : modelArtifacts . weightData . byteLength , } ; } returns function computeFloat16MantisaTable ( ) { const convertMantissa = ( i ) : number => { let m = i << 13 ; let e = 0 ; while ( ( m & 0x00800000 ) === 0 ) { e -= 0x00800000 ; m <<= 1 ; } m &= ~ 0x00800000 ; e += 0x38800000 ; return m | e ; } ; const mantisaTable = new Uint32Array ( 2048 ) ; mantisaTable [ 0 ] = 0 ; for ( let i = 1 ; i < 1024 ; i ++ ) { mantisaTable [ i ] = convertMantissa ( i ) ; } for ( let i = 1024 ; i < 2048 ; i ++ ) { mantisaTable [ i ] = 0x38000000 + ( ( i - 1024 ) << 13 ) ; } return mantisaTable ; } returns function computeFloat16ExponentTable ( ) { const exponentTable = new Uint32Array ( 64 ) ; exponentTable [ 0 ] = 0 ; exponentTable [ 31 ] = 0x47800000 ; exponentTable [ 32 ] = 0x80000000 ; exponentTable [ 63 ] = 0xc7800000 ; for ( let i = 1 ; i < 31 ; i ++ ) { exponentTable [ i ] = i << 23 ; } for ( let i = 33 ; i < 63 ; i ++ ) { exponentTable [ i ] = 0x80000000 + ( ( i - 32 ) << 23 ) ; } return exponentTable ; } returns function computeFloat16OffsetTable ( ) { const offsetTable = new Uint32Array ( 64 ) ; for ( let i = 0 ; i < 64 ; i ++ ) { offsetTable [ i ] = 1024 ; } offsetTable [ 0 ] = offsetTable [ 32 ] = 0 ; return offsetTable ; } returns export function getFloat16Decoder ( ) { const mantisaTable = computeFloat16MantisaTable ( ) ; const exponentTable = computeFloat16ExponentTable ( ) ; const offsetTable = computeFloat16OffsetTable ( ) ; return ( quantizedArray ) => { const buffer = new ArrayBuffer ( 4 * quantizedArray . length ) ; const bufferUint32View = new Uint32Array ( buffer ) ; for ( let index = 0 ; index < quantizedArray . length ; index ++ ) { const float16Bits = quantizedArray [ index ] ; const float32Bits = mantisaTable [ offsetTable [ float16Bits >> 10 ] + ( float16Bits & 0x3ff ) ] + exponentTable [ float16Bits >> 10 ] ; bufferUint32View [ index ] = float32Bits ; } return new Float32Array ( buffer ) ; } ; }