license import { env } from '../environment' ; import { NamedTensorMap } from '../tensor_types' ; import * as util from '../util' ; import { decodeWeights } from './io_utils' ; import { monitorPromisesProgress } from './progress' ; import { DTYPE_VALUE_SIZE_MAP , LoadOptions , WeightsManifestConfig , WeightsManifestEntry } from './types' ; param fetchURLs param requestOptions param fetchFunc param onProgress returns export async function loadWeightsAsArrayBuffer ( fetchURLs , loadOptions ? ) { if ( loadOptions == null ) { loadOptions = { } ; } const fetchFunc = loadOptions . fetchFunc == null ? env ( ) . platform . fetch : loadOptions . fetchFunc ; const requests = fetchURLs . map ( fetchURL => fetchFunc ( fetchURL , loadOptions . requestInit , { isBinary : true } ) ) ; const fetchStartFraction = 0 ; const fetchEndFraction = 0.5 ; const responses = loadOptions . onProgress == null ? await Promise . all ( requests ) : await monitorPromisesProgress ( requests , loadOptions . onProgress , fetchStartFraction , fetchEndFraction ) ; const bufferPromises = responses . map ( response => response . arrayBuffer ( ) ) ; const bufferStartFraction = 0.5 ; const bufferEndFraction = 1 ; const buffers = loadOptions . onProgress == null ? await Promise . all ( bufferPromises ) : await monitorPromisesProgress ( bufferPromises , loadOptions . onProgress , bufferStartFraction , bufferEndFraction ) ; return buffers ; } param manifest param filePathPrefix param weightNames export async function loadWeights ( manifest , filePathPrefix = '' , weightNames ? , requestInit ? ) { const fetchWeights = ( fetchUrls ) => loadWeightsAsArrayBuffer ( fetchUrls , { requestInit } ) ; const loadWeights = weightsLoaderFactory ( fetchWeights ) ; return loadWeights ( manifest , filePathPrefix , weightNames ) ; } param fetchWeightsFunction returns export function weightsLoaderFactory ( fetchWeightsFunction ) { return async ( manifest , filePathPrefix = '' , weightNames ? ) : Promise < NamedTensorMap > => { const groupIndicesToFetchMap = manifest . map ( ( ) => false ) ; const groupWeightsToFetch = { } ; const weightsFound = weightNames != null ? weightNames . map ( ( ) => false ) : [ ] ; const allManifestWeightNames = [ ] ; manifest . forEach ( ( manifestGroupConfig , groupIndex ) => { let groupOffset = 0 ; manifestGroupConfig . weights . forEach ( weightsEntry => { const rawDtype = ( 'quantization' in weightsEntry ) ? weightsEntry . quantization . dtype : weightsEntry . dtype ; const weightsBytes = DTYPE_VALUE_SIZE_MAP [ rawDtype ] * util . sizeFromShape ( weightsEntry . shape ) ; const enqueueWeightsForFetchingFn = ( ) => { groupIndicesToFetchMap [ groupIndex ] = true ; if ( groupWeightsToFetch [ groupIndex ] == null ) { groupWeightsToFetch [ groupIndex ] = [ ] ; } groupWeightsToFetch [ groupIndex ] . push ( { manifestEntry : weightsEntry , groupOffset , sizeBytes : weightsBytes } ) ; } ; if ( weightNames != null ) { weightNames . forEach ( ( weightName , weightIndex ) => { if ( weightName === weightsEntry . name ) { enqueueWeightsForFetchingFn ( ) ; weightsFound [ weightIndex ] = true ; } } ) ; } else { enqueueWeightsForFetchingFn ( ) ; } allManifestWeightNames . push ( weightsEntry . name ) ; groupOffset += weightsBytes ; } ) ; } ) ; if ( ! weightsFound . every ( found => found ) ) { const weightsNotFound = weightNames . filter ( ( _ , i ) => ! weightsFound [ i ] ) ; throw new Error ( `Could not find weights in manifest with names: ` + `${ weightsNotFound . join ( ', ' ) }. \n` + `Manifest JSON has weights with names: ` + `${ allManifestWeightNames . join ( ', ' ) }.` ) ; } const groupIndicesToFetch = groupIndicesToFetchMap . reduce ( ( accumulator , shouldFetch , i ) => { if ( shouldFetch ) { accumulator . push ( i ) ; } return accumulator ; } , [ ] ) ; const fetchUrls = [ ] ; groupIndicesToFetch . forEach ( i => { manifest [ i ] . paths . forEach ( filepath => { const fetchUrl = filePathPrefix + ( ! filePathPrefix . endsWith ( '/' ) ? '/' : '' ) + filepath ; fetchUrls . push ( fetchUrl ) ; } ) ; } ) ; const buffers = await fetchWeightsFunction ( fetchUrls ) ; const weightsTensorMap = { } ; let bufferIndexOffset = 0 ; groupIndicesToFetch . forEach ( i => { const numBuffers = manifest [ i ] . paths . length ; let groupBytes = 0 ; for ( let i = 0 ; i < numBuffers ; i ++ ) { groupBytes += buffers [ bufferIndexOffset + i ] . byteLength ; } const groupBuffer = new ArrayBuffer ( groupBytes ) ; const groupByteBuffer = new Uint8Array ( groupBuffer ) ; let groupBufferOffset = 0 ; for ( let i = 0 ; i < numBuffers ; i ++ ) { const buffer = new Uint8Array ( buffers [ bufferIndexOffset + i ] ) ; groupByteBuffer . set ( buffer , groupBufferOffset ) ; groupBufferOffset += buffer . byteLength ; } const weightsEntries = groupWeightsToFetch [ i ] ; weightsEntries . forEach ( weightsEntry => { const byteBuffer = groupBuffer . slice ( weightsEntry . groupOffset , weightsEntry . groupOffset + weightsEntry . sizeBytes ) ; const nameToTensorMap = decodeWeights ( byteBuffer , [ weightsEntry . manifestEntry ] ) ; for ( const name in nameToTensorMap ) { weightsTensorMap [ name ] = nameToTensorMap [ name ] ; } } ) ; bufferIndexOffset += numBuffers ; } ) ; return weightsTensorMap ; } ; }