license /**
 * IOHandlers related to files, such as browser-triggered file downloads,
 * user-selected files in browser.
 */ import '../flags' ; import { env } from '../environment' ; import { basename , concatenateArrayBuffers , getModelArtifactsForJSON , getModelArtifactsInfoForJSON , getModelJSONForModelArtifacts } from './io_utils' ; import { IORouter , IORouterRegistry } from './router_registry' ; import { IOHandler , ModelArtifacts , ModelJSON , SaveResult , WeightsManifestConfig , WeightsManifestEntry } from './types' ; const DEFAULT_FILE_NAME_PREFIX = 'model' ; const DEFAULT_JSON_EXTENSION_NAME = '.json' ; const DEFAULT_WEIGHT_DATA_EXTENSION_NAME = '.weights.bin' ; function defer < T > ( f ) { return new Promise ( resolve => setTimeout ( resolve ) ) . then ( f ) ; } export class BrowserDownloads implements IOHandler { private readonly modelJsonFileName : string ; private readonly weightDataFileName : string ; private readonly modelJsonAnchor : HTMLAnchorElement ; private readonly weightDataAnchor : HTMLAnchorElement ; static readonly URL_SCHEME = 'downloads://' ; constructor ( fileNamePrefix ? ) { if ( ! env ( ) . getBool ( 'IS_BROWSER' ) ) { throw new Error ( 'browserDownloads() cannot proceed because the current environment ' + 'is not a browser.' ) ; } if ( fileNamePrefix . startsWith ( BrowserDownloads . URL_SCHEME ) ) { fileNamePrefix = fileNamePrefix . slice ( BrowserDownloads . URL_SCHEME . length ) ; } if ( fileNamePrefix == null || fileNamePrefix . length === 0 ) { fileNamePrefix = DEFAULT_FILE_NAME_PREFIX ; } this . modelJsonFileName = fileNamePrefix + DEFAULT_JSON_EXTENSION_NAME ; this . weightDataFileName = fileNamePrefix + DEFAULT_WEIGHT_DATA_EXTENSION_NAME ; } async save ( modelArtifacts ) { if ( typeof ( document ) === 'undefined' ) { throw new Error ( 'Browser downloads are not supported in ' + 'this environment since `document` is not present' ) ; } const weightsURL = window . URL . createObjectURL ( new Blob ( [ modelArtifacts . weightData ] , { type : 'application/octet-stream' } ) ) ; if ( modelArtifacts . modelTopology instanceof ArrayBuffer ) { throw new Error ( 'BrowserDownloads.save() does not support saving model topology ' + 'in binary formats yet.' ) ; } else { const weightsManifest = [ { paths : [ './' + this . weightDataFileName ] , weights : modelArtifacts . weightSpecs } ] ; const modelJSON = getModelJSONForModelArtifacts ( modelArtifacts , weightsManifest ) ; const modelJsonURL = window . URL . createObjectURL ( new Blob ( [ JSON . stringify ( modelJSON ) ] , { type : 'application/json' } ) ) ; const jsonAnchor = this . modelJsonAnchor == null ? document . createElement ( 'a' ) : this . modelJsonAnchor ; jsonAnchor . download = this . modelJsonFileName ; jsonAnchor . href = modelJsonURL ; await defer ( ( ) => jsonAnchor . dispatchEvent ( new MouseEvent ( 'click' ) ) ) ; if ( modelArtifacts . weightData != null ) { const weightDataAnchor = this . weightDataAnchor == null ? document . createElement ( 'a' ) : this . weightDataAnchor ; weightDataAnchor . download = this . weightDataFileName ; weightDataAnchor . href = weightsURL ; await defer ( ( ) => weightDataAnchor . dispatchEvent ( new MouseEvent ( 'click' ) ) ) ; } return { modelArtifactsInfo : getModelArtifactsInfoForJSON ( modelArtifacts ) } ; } } } class BrowserFiles implements IOHandler { private readonly jsonFile : File ; private readonly weightsFiles : File [ ] ; constructor ( files ) { if ( files == null || files . length < 1 ) { throw new Error ( `When calling browserFiles, at least 1 file is required, ` + `but received ${ files }` ) ; } this . jsonFile = files [ 0 ] ; this . weightsFiles = files . slice ( 1 ) ; } async load ( ) { return new Promise ( ( resolve , reject ) => { const jsonReader = new FileReader ( ) ; jsonReader . onload = ( event ) => { const modelJSON = JSON . parse ( ( event . target as any ) . result ) as ModelJSON ; const modelTopology = modelJSON . modelTopology ; if ( modelTopology == null ) { reject ( new Error ( `modelTopology field is missing from file ${ this . jsonFile . name }` ) ) ; return ; } const weightsManifest = modelJSON . weightsManifest ; if ( weightsManifest == null ) { reject ( new Error ( `weightManifest field is missing from file ${ this . jsonFile . name }` ) ) ; return ; } if ( this . weightsFiles . length === 0 ) { resolve ( { modelTopology } ) ; return ; } const modelArtifactsPromise = getModelArtifactsForJSON ( modelJSON , ( weightsManifest ) => this . loadWeights ( weightsManifest ) ) ; resolve ( modelArtifactsPromise ) ; } ; jsonReader . onerror = error => reject ( `Failed to read model topology and weights manifest JSON ` + `from file '${ this . jsonFile . name }'. BrowserFiles supports loading ` + `Keras-style tf.Model artifacts only.` ) ; jsonReader . readAsText ( this . jsonFile ) ; } ) ; } private loadWeights ( weightsManifest ) { const weightSpecs = [ ] ; const paths = [ ] ; for ( const entry of weightsManifest ) { weightSpecs . push ( ... entry . weights ) ; paths . push ( ... entry . paths ) ; } const pathToFile = this . checkManifestAndWeightFiles ( weightsManifest ) ; const promises = paths . map ( path => this . loadWeightsFile ( path , pathToFile [ path ] ) ) ; return Promise . all ( promises ) . then ( buffers => [ weightSpecs , concatenateArrayBuffers ( buffers ) ] ) ; } private loadWeightsFile ( path , file ) { return new Promise ( ( resolve , reject ) => { const weightFileReader = new FileReader ( ) ; weightFileReader . onload = ( event ) => { const weightData = ( event . target as any ) . result as ArrayBuffer ; resolve ( weightData ) ; } ; weightFileReader . onerror = error => reject ( `Failed to weights data from file of path '${ path }'.` ) ; weightFileReader . readAsArrayBuffer ( file ) ; } ) ; } /**
     * Check the compatibility between weights manifest and weight files.
     */ private checkManifestAndWeightFiles ( manifest ) { const basenames = [ ] ; const fileNames = this . weightsFiles . map ( file => basename ( file . name ) ) ; const pathToFile = { } ; for ( const group of manifest ) { group . paths . forEach ( path => { const pathBasename = basename ( path ) ; if ( basenames . indexOf ( pathBasename ) !== - 1 ) { throw new Error ( `Duplicate file basename found in weights manifest: ` + `'${ pathBasename }'` ) ; } basenames . push ( pathBasename ) ; if ( fileNames . indexOf ( pathBasename ) === - 1 ) { throw new Error ( `Weight file with basename '${ pathBasename }' is not provided.` ) ; } else { pathToFile [ path ] = this . weightsFiles [ fileNames . indexOf ( pathBasename ) ] ; } } ) ; } if ( basenames . length !== this . weightsFiles . length ) { throw new Error ( `Mismatch in the number of files in weights manifest ` + `(${ basenames . length }) and the number of weight files provided ` + `(${ this . weightsFiles . length }).` ) ; } return pathToFile ; } } export const browserDownloadsRouter = ( url ) => { if ( ! env ( ) . getBool ( 'IS_BROWSER' ) ) { return null ; } else { if ( ! Array . isArray ( url ) && url . startsWith ( BrowserDownloads . URL_SCHEME ) ) { return browserDownloads ( url . slice ( BrowserDownloads . URL_SCHEME . length ) ) ; } else { return null ; } } } ; IORouterRegistry . registerSaveRouter ( browserDownloadsRouter ) ; param fileNamePrefix param config returns doc export function browserDownloads ( fileNamePrefix = 'model' ) { return new BrowserDownloads ( fileNamePrefix ) ; } param files returns doc export function browserFiles ( files ) { return new BrowserFiles ( files ) ; }