tensorflowjs-string-weight-dos-poc / source /tf-layers-node-loadlayersmodel-decode-before-match.txt
hacnho's picture
Upload folder using huggingface_hub
32810b6 verified
Raw
History Blame Contribute Delete
4.01 kB
25762 function loadLayersModelFromIOHandler(handler, customObjects, options) {
25763 return __awaiter(this, void 0, void 0, function () {
25764 var artifacts, modelTopology, strict, fastWeightInit, model, trainingConfig, _a, modelWeights, optimizerWeights;
25765 return __generator(this, function (_b) {
25766 switch (_b.label) {
25767 case 0:
25768 if (options == null) {
25769 options = {};
25770 }
25771 if (handler.load == null) {
25772 throw new ValueError('Cannot proceed with model loading because the IOHandler provided ' +
25773 'does not have the `load` method implemented.');
25774 }
25775 return [4 /*yield*/, handler.load()];
25776 case 1:
25777 artifacts = _b.sent();
25778 modelTopology = artifacts.modelTopology;
25779 if (modelTopology['model_config'] != null) {
25780 modelTopology = modelTopology['model_config'];
25781 }
25782 strict = options.strict == null ? true : options.strict;
25783 fastWeightInit = artifacts.weightData != null && artifacts.weightSpecs != null && strict;
25784 model = deserialize(convertPythonicToTs(modelTopology), customObjects, fastWeightInit);
25785 trainingConfig = artifacts.trainingConfig;
25786 if (trainingConfig != null) {
25787 model.loadTrainingConfig(trainingConfig);
25788 }
25789 if (artifacts.userDefinedMetadata != null) {
25790 model.setUserDefinedMetadata(artifacts.userDefinedMetadata);
25791 }
25792 if (!(artifacts.weightData != null)) return [3 /*break*/, 4];
25793 // Loading weights requires weightSpecs.
25794 if (artifacts.weightSpecs == null) {
25795 throw new ValueError('LayersModel artifacts contains weight data, but not weight specs. ' +
25796 'Therefore loading of weights cannot proceed.');
25797 }
25798 _a = decodeModelAndOptimizerWeights(artifacts.weightData, artifacts.weightSpecs), modelWeights = _a.modelWeights, optimizerWeights = _a.optimizerWeights;
25799 model.loadWeights(modelWeights, strict);
25800 if (!(model.optimizer != null && optimizerWeights.length > 0)) return [3 /*break*/, 3];
25801 return [4 /*yield*/, model.optimizer.setWeights(optimizerWeights)];
25802 case 2:
25803 _b.sent();
25804 _b.label = 3;
25805 case 3:
25806 // Dispose temporary weight values.
25807 tfc.dispose(modelWeights);
25808 tfc.dispose(optimizerWeights.map(function (w) { return w.tensor; }));
25809 _b.label = 4;
25810 case 4: return [2 /*return*/, model];
25811 }
25812 });
25813 });
25814 }
25815 function decodeModelAndOptimizerWeights(weightData, specs) {
25816 var name2Tensor = tfc.io.decodeWeights(weightData, specs);
25817 var modelWeights = {};
25818 var optimizerWeights = [];
25819 specs.forEach(function (spec) {
25820 if (spec.group === 'optimizer') {
25821 optimizerWeights.push({ name: spec.name, tensor: name2Tensor[spec.name] });
25822 }
25823 else {
25824 modelWeights[spec.name] = name2Tensor[spec.name];
25825 }
25826 });
25827 return { modelWeights: modelWeights, optimizerWeights: optimizerWeights };
25828 }