wehub-resource-sync
7254f7b4d1
Build / Build (macos-latest) (push) Has been cancelled
Build / Build (ubuntu-latest) (push) Has been cancelled
Build / Build (windows-latest) (push) Has been cancelled
Build / Analyze (javascript) (push) Has been cancelled
Build / Analyze (python) (push) Has been cancelled
438 lines
18 KiB
JavaScript
438 lines
18 KiB
JavaScript
|
|
// Experimental
|
|
|
|
const pickle = {};
|
|
|
|
pickle.ModelFactory = class {
|
|
|
|
async match(context) {
|
|
const stream = context.stream;
|
|
const signature = [0x80, undefined, 0x8a, 0x0a, 0x6c, 0xfc, 0x9c, 0x46, 0xf9, 0x20, 0x6a, 0xa8, 0x50, 0x19];
|
|
if (stream && signature.length <= stream.length && stream.peek(signature.length).every((value, index) => signature[index] === undefined || signature[index] === value)) {
|
|
// Reject PyTorch models with .pkl file extension.
|
|
return null;
|
|
}
|
|
const obj = await context.peek('pkl');
|
|
if (obj !== undefined) {
|
|
const name = obj && obj.__class__ && obj.__class__.__module__ && obj.__class__.__name__ ? `${obj.__class__.__module__}.${obj.__class__.__name__}` : '';
|
|
if (!name.startsWith('__torch__.') &&
|
|
!name.startsWith('catboost.') &&
|
|
!name.startsWith('autogluon.tabular.models.catboost.')) {
|
|
return context.set('pickle', obj);
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
async open(context) {
|
|
let format = 'Pickle';
|
|
const obj = context.value;
|
|
if (obj === null || obj === undefined) {
|
|
context.error(new pickle.Error("Unsupported Pickle null object."));
|
|
} else if (obj instanceof Error) {
|
|
throw obj;
|
|
} else if (!Array.isArray(obj) && obj && obj.__class__) {
|
|
const formats = new Map([
|
|
['cuml.ensemble.randomforestclassifier.RandomForestClassifier', 'cuML'],
|
|
['shap.explainers._linear.LinearExplainer', 'SHAP'],
|
|
['gensim.models.word2vec.Word2Vec', 'Gensim'],
|
|
['ray.rllib.algorithms.ppo.ppo.PPOConfig', 'Ray RLlib'],
|
|
['builtins.bytearray', 'Pickle'],
|
|
['builtins.dict', 'Pickle'],
|
|
['collections.OrderedDict', 'Pickle'],
|
|
['numpy.ndarray', 'NumPy NDArray'],
|
|
]);
|
|
const type = `${obj.__class__.__module__}.${obj.__class__.__name__}`;
|
|
if (formats.has(type)) {
|
|
format = formats.get(type);
|
|
} else {
|
|
context.error(new pickle.Error(`Unsupported Pickle type '${type}'.`));
|
|
}
|
|
}
|
|
return new pickle.Model(obj, format);
|
|
}
|
|
};
|
|
|
|
pickle.Model = class {
|
|
|
|
constructor(value, format) {
|
|
this.format = format;
|
|
this.modules = [new pickle.Module(null, value)];
|
|
}
|
|
};
|
|
|
|
pickle.Module = class {
|
|
|
|
constructor(type = '', obj = null) {
|
|
this.type = type;
|
|
this.inputs = [];
|
|
this.outputs = [];
|
|
this.nodes = [];
|
|
const weights = this.type === 'weights' ? obj : pickle.Utility.weights(obj);
|
|
if (weights) {
|
|
for (const [name, module] of weights) {
|
|
const node = new pickle.Node(module, name, 'Weights');
|
|
this.nodes.push(node);
|
|
}
|
|
} else if (pickle.Utility.isTensor(obj)) {
|
|
const type = `${obj.__class__.__module__}.${obj.__class__.__name__}`;
|
|
const node = new pickle.Node({ value: obj }, null, type);
|
|
this.nodes.push(node);
|
|
} else if (Array.isArray(obj) && (obj.every((item) => item.__class__) || (obj.every((item) => Array.isArray(item))))) {
|
|
for (const item of obj) {
|
|
this.nodes.push(new pickle.Node(item));
|
|
}
|
|
} else if (obj && obj.__class__) {
|
|
this.nodes.push(new pickle.Node(obj));
|
|
} else if (obj && Object(obj) === obj) {
|
|
this.nodes.push(new pickle.Node(obj));
|
|
}
|
|
}
|
|
};
|
|
|
|
pickle.Node = class {
|
|
|
|
constructor(obj, name, type, stack) {
|
|
if (typeof type !== 'string') {
|
|
type = obj.__class__ ? `${obj.__class__.__module__}.${obj.__class__.__name__}` : 'builtins.object';
|
|
}
|
|
this.type = { name: type };
|
|
this.name = name || '';
|
|
this.inputs = [];
|
|
if (type === 'builtins.bytearray') {
|
|
const argument = new pickle.Argument('value', Array.from(obj), 'byte[]');
|
|
this.inputs.push(argument);
|
|
return;
|
|
}
|
|
const weights = pickle.Utility.weights(obj);
|
|
if (weights) {
|
|
const type = this.type.name;
|
|
this.type = new pickle.Module('weights', weights);
|
|
this.type.name = type;
|
|
return;
|
|
}
|
|
const entries = obj instanceof Map ? Array.from(obj) : Object.entries(obj);
|
|
for (const [name, value] of entries) {
|
|
if (name === '__class__') {
|
|
continue;
|
|
} else if (value && pickle.Utility.isTensor(value)) {
|
|
const identifier = value.__name__ || '';
|
|
const tensor = new pickle.Tensor(value);
|
|
const values = [new pickle.Value(identifier, null, tensor)];
|
|
const argument = new pickle.Argument(name, values);
|
|
this.inputs.push(argument);
|
|
} else if (Array.isArray(value) && value.length > 0 && value.every((obj) => pickle.Utility.isTensor(obj))) {
|
|
const values = value.map((obj) => new pickle.Value(obj.__name__ || '', null, new pickle.Tensor(obj)));
|
|
const argument = new pickle.Argument(name, values);
|
|
this.inputs.push(argument);
|
|
} else if (value && value.__class__ && value.__class__.__module__ === 'builtins' && (value.__class__.__name__ === 'function' || value.__class__.__name__ === 'type')) {
|
|
const obj = {};
|
|
obj.__class__ = value;
|
|
const node = new pickle.Node(obj, null, null, stack);
|
|
const argument = new pickle.Argument(name, node, 'object');
|
|
this.inputs.push(argument);
|
|
} else if (pickle.Utility.isByteArray(value)) {
|
|
const argument = new pickle.Argument(name, Array.from(value), 'byte[]');
|
|
this.inputs.push(argument);
|
|
} else {
|
|
stack = stack || new Set();
|
|
if (value && Array.isArray(value) && value.every((obj) => typeof obj === 'string')) {
|
|
const argument = new pickle.Argument(name, value, 'string[]');
|
|
this.inputs.push(argument);
|
|
} else if (value && Array.isArray(value) && value.every((obj) => typeof obj === 'number')) {
|
|
const argument = new pickle.Argument(name, value, 'attribute');
|
|
this.inputs.push(argument);
|
|
} else if (value && Array.isArray(value) && value.length > 0 && value.every((obj) => obj && (obj.__class__ || obj === Object(obj)))) {
|
|
const chain = stack;
|
|
const values = value.filter((value) => !chain.has(value));
|
|
const nodes = values.map((value) => {
|
|
chain.add(value);
|
|
const node = new pickle.Node(value, null, null, chain);
|
|
chain.delete(value);
|
|
return node;
|
|
});
|
|
const argument = new pickle.Argument(name, nodes, 'object[]');
|
|
this.inputs.push(argument);
|
|
} else if (value && (value.__class__ || pickle.Utility.isObject(value)) && !stack.has(value)) {
|
|
stack.add(value);
|
|
const node = new pickle.Node(value, null, null, stack);
|
|
const visible = name !== '_metadata' || !pickle.Utility.isMetadataObject(value);
|
|
const argument = new pickle.Argument(name, node, 'object', visible);
|
|
this.inputs.push(argument);
|
|
stack.delete(value);
|
|
} else {
|
|
const argument = new pickle.Argument(name, value, 'attribute');
|
|
this.inputs.push(argument);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
pickle.Argument = class {
|
|
|
|
constructor(name, value, type = null, visible = true) {
|
|
this.name = name.toString();
|
|
this.value = value;
|
|
this.type = type;
|
|
this.visible = visible;
|
|
}
|
|
};
|
|
|
|
pickle.Value = class {
|
|
|
|
constructor(name, type, initializer = null) {
|
|
if (typeof name !== 'string') {
|
|
throw new pickle.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
|
|
}
|
|
this.name = name;
|
|
this.type = initializer && initializer.type ? initializer.type : type || null;
|
|
this.initializer = initializer;
|
|
}
|
|
};
|
|
|
|
pickle.Tensor = class {
|
|
|
|
constructor(obj) {
|
|
if (obj.__class__ && (obj.__class__.__module__ === 'torch' || obj.__class__.__module__ === 'torch.nn.parameter')) {
|
|
// PyTorch tensor
|
|
const tensor = obj.__class__.__module__ === 'torch.nn.parameter' && obj.__class__.__name__ === 'Parameter' ? obj.data : obj;
|
|
const layout = tensor.layout ? tensor.layout.__str__() : null;
|
|
const storage = tensor.storage();
|
|
const size = tensor.size() || [];
|
|
if (!layout || layout === 'torch.strided') {
|
|
this.type = new pickle.TensorType(storage.dtype.__reduce__(), new pickle.TensorShape(size));
|
|
this.values = storage.data;
|
|
this.encoding = '<';
|
|
this.indices = null;
|
|
this.stride = tensor.stride();
|
|
const stride = this.stride;
|
|
const offset = tensor.storage_offset();
|
|
let length = 0;
|
|
if (!Array.isArray(stride)) {
|
|
length = storage.size();
|
|
} else if (size.every((v) => v !== 0)) {
|
|
length = size.reduce((a, v, i) => a + stride[i] * (v - 1), 1);
|
|
}
|
|
if (this.values !== undefined) {
|
|
if (offset !== 0 || length !== storage.size()) {
|
|
const itemsize = storage.dtype.itemsize();
|
|
const stream = this.values;
|
|
const position = stream.position;
|
|
stream.seek(itemsize * offset);
|
|
this.values = stream.peek(itemsize * length);
|
|
stream.seek(position);
|
|
} else if (this.values) {
|
|
this.values = this.values.peek();
|
|
}
|
|
}
|
|
} else {
|
|
throw new pickle.Error(`Unsupported tensor layout '${layout}'.`);
|
|
}
|
|
} else {
|
|
// NumPy array
|
|
const array = obj;
|
|
this.type = new pickle.TensorType(array.dtype.__name__, new pickle.TensorShape(array.shape));
|
|
this.stride = Array.isArray(array.strides) ? array.strides.map((stride) => stride / array.itemsize) : null;
|
|
this.encoding = this.type.dataType === 'string' || this.type.dataType === 'object' || this.type.dataType === 'datetime' ? '|' : array.dtype.byteorder;
|
|
this.values = this.type.dataType === 'string' || this.type.dataType === 'object' || this.type.dataType === 'datetime' || this.type.dataType === 'void' ? array.flatten().tolist() : array.tobytes();
|
|
}
|
|
}
|
|
};
|
|
|
|
pickle.TensorType = class {
|
|
|
|
constructor(dataType, shape) {
|
|
this.dataType = dataType;
|
|
this.shape = shape;
|
|
}
|
|
|
|
toString() {
|
|
return this.dataType + this.shape.toString();
|
|
}
|
|
};
|
|
|
|
pickle.TensorShape = class {
|
|
|
|
constructor(dimensions) {
|
|
this.dimensions = dimensions;
|
|
}
|
|
|
|
toString() {
|
|
return this.dimensions ? (`[${this.dimensions.map((dimension) => dimension.toString()).join(',')}]`) : '';
|
|
}
|
|
};
|
|
|
|
pickle.Utility = class {
|
|
|
|
static isSubclass(value, name) {
|
|
if (value && value.__module__ && value.__name__) {
|
|
return name === `${value.__module__}.${value.__name__}`;
|
|
} else if (value && value.__bases__) {
|
|
return value.__bases__.some((obj) => pickle.Utility.isSubclass(obj, name));
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static isInstance(value, name) {
|
|
return value && value.__class__ ? pickle.Utility.isSubclass(value.__class__, name) : false;
|
|
}
|
|
|
|
static isMetadataObject(obj) {
|
|
if (pickle.Utility.isInstance(obj, 'collections.OrderedDict')) {
|
|
for (const value of obj.values()) {
|
|
if (pickle.Utility.isInstance(value, 'builtins.dict')) {
|
|
const entries = Array.from(value);
|
|
if (entries.length !== 1 && entries[0] !== 'version' && entries[1] !== 1) {
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static isByteArray(obj) {
|
|
return obj && obj.__class__ && obj.__class__.__module__ === 'builtins' && obj.__class__.__name__ === 'bytearray';
|
|
}
|
|
|
|
static isObject(obj) {
|
|
if (obj && typeof obj === 'object') {
|
|
const proto = Object.getPrototypeOf(obj);
|
|
return proto === Object.prototype || proto === null;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static isTensor(obj) {
|
|
return obj && obj.__class__ && obj.__class__.__name__ &&
|
|
((obj.__class__.__module__ === 'numpy' && obj.__class__.__name__ === 'ndarray') ||
|
|
(obj.__class__.__module__ === 'numpy' && obj.__class__.__name__ === 'matrix') ||
|
|
(obj.__class__.__module__ === 'jax' && obj.__class__.__name__ === 'Array') ||
|
|
(obj.__class__.__module__ === 'torch.nn.parameter' && obj.__class__.__name__ === 'Parameter') ||
|
|
(obj.__class__.__module__ === 'torch' && obj.__class__.__name__.endsWith('Tensor')));
|
|
}
|
|
|
|
static weights(obj) {
|
|
const type = obj && obj.__class__ && obj.__class__.__module__ && obj.__class__.__name__ ? `${obj.__class__.__module__}.${obj.__class__.__name__}` : null;
|
|
if (type && type !== 'builtins.dict' && type !== 'builtins.object' && type !== 'collections.OrderedDict' && type !== 'torch.nn.modules.module.Module') {
|
|
return null;
|
|
}
|
|
if (pickle.Utility.isTensor(obj)) {
|
|
return null;
|
|
}
|
|
if (obj instanceof Map === false && obj && !Array.isArray(obj) && Object(obj) === obj) {
|
|
const entries = Object.entries(obj);
|
|
const named = entries.filter(([name, value]) => (typeof name === 'string' && (name.indexOf('.') !== -1 || name.indexOf('|') !== -1)) && pickle.Utility.isTensor(value));
|
|
if (named.length > 0 && (named.length / entries.length) >= 0.8) {
|
|
obj = new Map(entries);
|
|
}
|
|
}
|
|
if (obj instanceof Map) {
|
|
const entries = Array.from(obj).filter(([name]) => name !== '_metadata');
|
|
let dot = 0;
|
|
let pipe = 0;
|
|
let underscore = 0;
|
|
let count = 0;
|
|
let valid = true;
|
|
for (const [name, value] of entries) {
|
|
if (typeof name === 'string') {
|
|
count++;
|
|
dot += name.indexOf('.') !== -1;
|
|
pipe += name.indexOf('|') !== -1;
|
|
underscore += name.endsWith('_w') || name.endsWith('_b') || name.endsWith('_bn_s');
|
|
}
|
|
if (pickle.Utility.isInstance(value, 'builtins.dict') && !Array.from(value.values()).every((value) => !pickle.Utility.isTensor(value))) {
|
|
valid = false;
|
|
}
|
|
}
|
|
if (valid && count > 1 && (dot >= count || pipe >= count || underscore >= count) && (count / entries.length) >= 0.8) {
|
|
let separator = null;
|
|
if (dot >= pipe && dot >= underscore) {
|
|
separator = '.';
|
|
} else if (pipe >= underscore) {
|
|
separator = '|';
|
|
} else {
|
|
separator = '_';
|
|
}
|
|
const modules = new Map();
|
|
for (const [name, value] of entries) {
|
|
let c = separator;
|
|
if (!c) {
|
|
c = name.indexOf('.') === -1 && name.indexOf('|') !== -1 ? '|' : '.';
|
|
}
|
|
const path = name.split(c);
|
|
let property = path.pop();
|
|
if (path.length > 1 && path[path.length - 1] === '_packed_params') {
|
|
property = `${path.pop()}.${property}`;
|
|
}
|
|
const key = path.join(separator);
|
|
if (!modules.has(key)) {
|
|
modules.set(key, {});
|
|
}
|
|
const module = modules.get(key);
|
|
if (pickle.Utility.isTensor(value)) {
|
|
value.__name__ = name;
|
|
}
|
|
module[property] = value;
|
|
}
|
|
return modules;
|
|
}
|
|
}
|
|
if (obj && !Array.isArray(obj) && Object(obj) === obj) {
|
|
const modules = new Map();
|
|
const entries = obj instanceof Map ? Array.from(obj) : Object.entries(obj);
|
|
if (entries.length > 0 && entries) {
|
|
for (const [key, value] of entries) {
|
|
const name = key.toString();
|
|
if (!value || Object(value) !== value || pickle.Utility.isTensor(value) || ArrayBuffer.isView(value)) {
|
|
return null;
|
|
}
|
|
if (!modules.has(name)) {
|
|
modules.set(name, {});
|
|
}
|
|
const module = modules.get(name);
|
|
let tensor = false;
|
|
const entries = value instanceof Map ? value : new Map(Object.entries(value));
|
|
for (const [name, value] of entries) {
|
|
if (typeof name !== 'string') {
|
|
return null;
|
|
}
|
|
if (name.indexOf('.') !== -1) {
|
|
return null;
|
|
}
|
|
if (name === '_metadata') {
|
|
continue;
|
|
}
|
|
if (typeof value === 'string' || typeof value === 'number') {
|
|
module[name] = value;
|
|
continue;
|
|
}
|
|
if (pickle.Utility.isTensor(value)) {
|
|
value.__name__ = name;
|
|
module[name] = value;
|
|
tensor = true;
|
|
}
|
|
}
|
|
if (!tensor) {
|
|
return null;
|
|
}
|
|
}
|
|
return modules;
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
};
|
|
|
|
pickle.Error = class extends Error {
|
|
|
|
constructor(message) {
|
|
super(message);
|
|
this.name = 'Error loading Pickle model.';
|
|
}
|
|
};
|
|
|
|
export const ModelFactory = pickle.ModelFactory;
|