ai-content-maker/.venv/Lib/site-packages/torch/utils/model_dump/code.js

690 lines
19 KiB
JavaScript

import { h, Component, render } from 'https://unpkg.com/preact?module';
import htm from 'https://unpkg.com/htm?module';
const html = htm.bind(h);
const BURNED_IN_MODEL_INFO = null;
// https://stackoverflow.com/a/20732091
function humanFileSize(size) {
if (size == 0) { return "0 B"; }
var i = Math.floor( Math.log(size) / Math.log(1024) );
return (size / Math.pow(1024, i)).toFixed(2) * 1 + ' ' + ['B', 'kB', 'MB', 'GB', 'TB'][i];
}
function caret(down) {
return down ? "\u25BE" : "\u25B8";
}
class Blamer {
constructor() {
this.blame_on_click = false;
this.aux_content_pane = null;
}
setAuxContentPane(pane) {
this.aux_content_pane = pane;
}
readyBlame() {
this.blame_on_click = true;
}
maybeBlame(arg) {
if (!this.blame_on_click) {
return;
}
this.blame_on_click = false;
if (!this.aux_content_pane) {
return;
}
this.aux_content_pane.doBlame(arg);
}
}
let blame = new Blamer();
class Hider extends Component {
constructor() {
super();
this.state = { shown: null };
}
componentDidMount() {
this.setState({ shown: this.props.shown === "true" });
}
render({name, children}, {shown}) {
let my_caret = html`<span class=caret onClick=${() => this.click()} >${caret(shown)}</span>`;
return html`<div data-hider-title=${name} data-shown=${shown}>
<h2>${my_caret} ${name}</h2>
<div>${shown ? this.props.children : []}</div></div>`;
}
click() {
this.setState({shown: !this.state.shown});
}
}
function ModelSizeSection({model: {file_size, zip_files}}) {
let store_size = 0;
let compr_size = 0;
for (const zi of zip_files) {
if (zi.compression === 0) {
// TODO: Maybe check that compressed_size === file_size.
store_size += zi.compressed_size;
} else {
compr_size += zi.compressed_size;
}
}
let zip_overhead = file_size - store_size - compr_size;
// TODO: Better formatting. Right-align this.
return html`
<${Hider} name="Model Size" shown=true>
<pre>.
Model size: ${file_size} (${humanFileSize(file_size)})
Stored files: ${store_size} (${humanFileSize(store_size)})
Compressed files: ${compr_size} (${humanFileSize(compr_size)})
Zip overhead: ${zip_overhead} (${humanFileSize(zip_overhead)})
</pre><//>`;
}
function StructuredDataSection({name, data, shown}) {
return html`
<${Hider} name=${name} shown=${shown}>
<div style="font-family:monospace;">
<${StructuredData} data=${data} indent="" prefix=""/>
</div><//>`;
}
class StructuredData extends Component {
constructor() {
super();
this.state = { shown: false };
this.INLINE_TYPES = new Set(["boolean", "number", "string"])
this.IGNORED_STATE_KEYS = new Set(["training", "_is_full_backward_hook"])
}
click() {
this.setState({shown: !this.state.shown});
}
expando(data) {
if (data === null || this.INLINE_TYPES.has(typeof(data))) {
return false;
}
if (typeof(data) != "object") {
throw new Error("Not an object");
}
if (Array.isArray(data)) {
// TODO: Maybe show simple lists and tuples on one line.
return true;
}
if (data.__tuple_values__) {
// TODO: Maybe show simple lists and tuples on one line.
return true;
}
if (data.__is_dict__) {
// TODO: Maybe show simple (empty?) dicts on one line.
return true;
}
if (data.__module_type__) {
return true;
}
if (data.__tensor_v2__) {
return false;
}
if (data.__qtensor__) {
return false;
}
throw new Error("Can't handle data type.", data);
}
renderHeadline(data) {
if (data === null) {
return "None";
}
if (typeof(data) == "boolean") {
const sd = String(data);
return sd.charAt(0).toUpperCase() + sd.slice(1);
}
if (typeof(data) == "number") {
return JSON.stringify(data);
}
if (typeof(data) == "string") {
return JSON.stringify(data);
}
if (typeof(data) != "object") {
throw new Error("Not an object");
}
if (Array.isArray(data)) {
return "list([";
}
if (data.__tuple_values__) {
return "tuple((";
}
if (data.__is_dict__) {
return "dict({";
}
if (data.__module_type__) {
return data.__module_type__ + "()";
}
if (data.__tensor_v2__) {
const [storage, offset, size, stride, grad] = data.__tensor_v2__;
const [dtype, key, device, numel] = storage;
return this.renderTensor(
"tensor", dtype, key, device, numel, offset, size, stride, grad, []);
}
if (data.__qtensor__) {
const [storage, offset, size, stride, quantizer, grad] = data.__qtensor__;
const [dtype, key, device, numel] = storage;
let extra_parts = [];
if (quantizer[0] == "per_tensor_affine") {
extra_parts.push(`scale=${quantizer[1]}`);
extra_parts.push(`zero_point=${quantizer[2]}`);
} else {
extra_parts.push(`quantizer=${quantizer[0]}`);
}
return this.renderTensor(
"qtensor", dtype, key, device, numel, offset, size, stride, grad, extra_parts);
}
throw new Error("Can't handle data type.", data);
}
renderTensor(
prefix,
dtype,
storage_key,
device,
storage_numel,
offset,
size,
stride,
grad,
extra_parts) {
let parts = [
"(" + size.join(",") + ")",
dtype,
];
parts.push(...extra_parts);
if (device != "cpu") {
parts.push(device);
}
if (grad) {
parts.push("grad");
}
// TODO: Check stride and indicate if the tensor is channels-last or non-contiguous
// TODO: Check size, stride, offset, and numel and indicate if
// the tensor doesn't use all data in storage.
// TODO: Maybe show key?
void(offset);
void(stride);
void(storage_key);
void(storage_numel);
return prefix + "(" + parts.join(", ") + ")";
}
renderBody(indent, data) {
if (data === null || this.INLINE_TYPES.has(typeof(data))) {
throw "Should not reach here."
}
if (typeof(data) != "object") {
throw new Error("Not an object");
}
if (Array.isArray(data)) {
let new_indent = indent + "\u00A0\u00A0";
let parts = [];
for (let idx = 0; idx < data.length; idx++) {
// Does it make sense to put explicit index numbers here?
parts.push(html`<br/><${StructuredData} prefix=${idx + ": "} indent=${new_indent} data=${data[idx]} />`);
}
return parts;
}
if (data.__tuple_values__) {
// Handled the same as lists.
return this.renderBody(indent, data.__tuple_values__);
}
if (data.__is_dict__) {
let new_indent = indent + "\u00A0\u00A0";
let parts = [];
for (let idx = 0; idx < data.keys.length; idx++) {
if (typeof(data.keys[idx]) != "string") {
parts.push(html`<br/>${new_indent}Non-string key`);
} else {
parts.push(html`<br/><${StructuredData} prefix=${data.keys[idx] + ": "} indent=${new_indent} data=${data.values[idx]} />`);
}
}
return parts;
}
if (data.__module_type__) {
const mstate = data.state;
if (mstate === null || typeof(mstate) != "object") {
throw new Error("Bad module state");
}
let new_indent = indent + "\u00A0\u00A0";
let parts = [];
if (mstate.__is_dict__) {
// TODO: Less copy/paste between this and normal dicts.
for (let idx = 0; idx < mstate.keys.length; idx++) {
if (typeof(mstate.keys[idx]) != "string") {
parts.push(html`<br/>${new_indent}Non-string key`);
} else if (this.IGNORED_STATE_KEYS.has(mstate.keys[idx])) {
// Do nothing.
} else {
parts.push(html`<br/><${StructuredData} prefix=${mstate.keys[idx] + ": "} indent=${new_indent} data=${mstate.values[idx]} />`);
}
}
} else if (mstate.__tuple_values__) {
parts.push(html`<br/><${StructuredData} prefix="" indent=${new_indent} data=${mstate} />`);
} else if (mstate.__module_type__) {
// We normally wouldn't have the state of a module be another module,
// but we use "modules" to encode special values (like Unicode decode
// errors) that might be valid states. Just go with it.
parts.push(html`<br/><${StructuredData} prefix="" indent=${new_indent} data=${mstate} />`);
} else {
throw new Error("Bad module state");
}
return parts;
}
if (data.__tensor_v2__) {
throw "Should not reach here."
}
if (data.__qtensor__) {
throw "Should not reach here."
}
throw new Error("Can't handle data type.", data);
}
render({data, indent, prefix}, {shown}) {
const exp = this.expando(data) ? html`<span class=caret onClick=${() => this.click()} >${caret(shown)} </span>` : "";
const headline = this.renderHeadline(data);
const body = shown ? this.renderBody(indent, data) : "";
return html`${indent}${exp}${prefix}${headline}${body}`;
}
}
function ZipContentsSection({model: {zip_files}}) {
// TODO: Add human-readable sizes?
// TODO: Add sorting options?
// TODO: Add hierarchical collapsible tree?
return html`
<${Hider} name="Zip Contents" shown=false>
<table>
<thead>
<tr>
<th>Mode</th>
<th>Size</th>
<th>Compressed</th>
<th>Name</th>
</tr>
</thead>
<tbody style="font-family:monospace;">
${zip_files.map(zf => html`<tr>
<td>${{0: "store", 8: "deflate"}[zf.compression] || zf.compression}</td>
<td>${zf.file_size}</td>
<td>${zf.compressed_size}</td>
<td>${zf.filename}</td>
</tr>`)}
</tbody>
</table><//>`;
}
function CodeSection({model: {code_files}}) {
return html`
<${Hider} name="Code" shown=false>
<div>
${Object.entries(code_files).map(([fn, code]) => html`<${OneCodeSection}
filename=${fn} code=${code} />`)}
</div><//>`;
}
class OneCodeSection extends Component {
constructor() {
super();
this.state = { shown: false };
}
click() {
const shown = !this.state.shown;
this.setState({shown: shown});
}
render({filename, code}, {shown}) {
const header = html`
<h3 style="font-family:monospace;">
<span class=caret onClick=${() => this.click()} >${caret(shown)} </span>
${filename}</h3>
`;
if (!shown) {
return header;
}
return html`
${header}
<pre>${code.map(c => this.renderBlock(c))}</pre>
`;
}
renderBlock([text, ist_file, line, ist_s_text, s_start, s_end]) {
return html`<span
onClick=${() => blame.maybeBlame({ist_file, line, ist_s_text, s_start, s_end})}
>${text}</span>`;
}
}
function ExtraJsonSection({files}) {
return html`
<${Hider} name="Extra files (JSON)" shown=false>
<div>
<p>Use "Log Raw Model Info" for hierarchical view in browser console.</p>
${Object.entries(files).map(([fn, json]) => html`<${OneJsonSection}
filename=${fn} json=${json} />`)}
</div><//>`;
}
class OneJsonSection extends Component {
constructor() {
super();
this.state = { shown: false };
}
click() {
const shown = !this.state.shown;
this.setState({shown: shown});
}
render({filename, json}, {shown}) {
const header = html`
<h3 style="font-family:monospace;">
<span class=caret onClick=${() => this.click()} >${caret(shown)} </span>
${filename}</h3>
`;
if (!shown) {
return header;
}
return html`
${header}
<pre>${JSON.stringify(json, null, 2)}</pre>
`;
}
}
function ExtraPicklesSection({files}) {
return html`
<${Hider} name="Extra Pickles" shown=false>
<div>
${Object.entries(files).map(([fn, content]) => html`<${OnePickleSection}
filename=${fn} content=${content} />`)}
</div><//>`;
}
class OnePickleSection extends Component {
constructor() {
super();
this.state = { shown: false };
}
click() {
const shown = !this.state.shown;
this.setState({shown: shown});
}
render({filename, content}, {shown}) {
const header = html`
<h3 style="font-family:monospace;">
<span class=caret onClick=${() => this.click()} >${caret(shown)} </span>
${filename}</h3>
`;
if (!shown) {
return header;
}
return html`
${header}
<pre>${content}</pre>
`;
}
}
function assertStorageAreEqual(key, lhs, rhs) {
if (lhs.length !== rhs.length ||
!lhs.every((val, idx) => val === rhs[idx])) {
throw new Error("Storage mismatch for key '" + key + "'");
}
}
function computeTensorMemory(numel, dtype) {
const sizes = {
"Byte": 1,
"Char": 1,
"Short": 2,
"Int": 4,
"Long": 8,
"Half": 2,
"Float": 4,
"Double": 8,
"ComplexHalf": 4,
"ComplexFloat": 8,
"ComplexDouble": 16,
"Bool": 1,
"QInt8": 1,
"QUInt8": 1,
"QInt32": 4,
"BFloat16": 2,
};
let dtsize = sizes[dtype];
if (!dtsize) {
throw new Error("Unrecognized dtype: " + dtype);
}
return numel * dtsize;
}
// TODO: Maybe track by dtype as well.
// TODO: Maybe distinguish between visible size and storage size.
function getTensorStorages(data) {
if (data === null) {
return new Map();
}
if (typeof(data) == "boolean") {
return new Map();
}
if (typeof(data) == "number") {
return new Map();
}
if (typeof(data) == "string") {
return new Map();
}
if (typeof(data) != "object") {
throw new Error("Not an object");
}
if (Array.isArray(data)) {
let result = new Map();
for (const item of data) {
const tensors = getTensorStorages(item);
for (const [key, storage] of tensors.entries()) {
if (!result.has(key)) {
result.set(key, storage);
} else {
const old_storage = result.get(key);
assertStorageAreEqual(key, old_storage, storage);
}
}
}
return result;
}
if (data.__tuple_values__) {
return getTensorStorages(data.__tuple_values__);
}
if (data.__is_dict__) {
return getTensorStorages(data.values);
}
if (data.__module_type__) {
return getTensorStorages(data.state);
}
if (data.__tensor_v2__) {
const [storage, offset, size, stride, grad] = data.__tensor_v2__;
const [dtype, key, device, numel] = storage;
return new Map([[key, storage]]);
}
if (data.__qtensor__) {
const [storage, offset, size, stride, quantizer, grad] = data.__qtensor__;
const [dtype, key, device, numel] = storage;
return new Map([[key, storage]]);
}
throw new Error("Can't handle data type.", data);
}
function getTensorMemoryByDevice(pickles) {
let all_tensors = [];
for (const [name, pickle] of pickles) {
const tensors = getTensorStorages(pickle);
all_tensors.push(...tensors.values());
}
let result = {};
for (const storage of all_tensors.values()) {
const [dtype, key, device, numel] = storage;
const size = computeTensorMemory(numel, dtype);
result[device] = (result[device] || 0) + size;
}
return result;
}
// Make this a separate component so it is rendered lazily.
class OpenTensorMemorySection extends Component {
render({model: {model_data, constants}}) {
let sizes = getTensorMemoryByDevice(new Map([
["data", model_data],
["constants", constants],
]));
return html`
<table>
<thead>
<tr>
<th>Device</th>
<th>Bytes</th>
<th>Human</th>
</tr>
</thead>
<tbody style="font-family:monospace;">
${Object.entries(sizes).map(([dev, size]) => html`<tr>
<td>${dev}</td>
<td>${size}</td>
<td>${humanFileSize(size)}</td>
</tr>`)}
</tbody>
</table>`;
}
}
function TensorMemorySection({model}) {
return html`
<${Hider} name="Tensor Memory" shown=false>
<${OpenTensorMemorySection} model=${model} /><//>`;
}
class AuxContentPane extends Component {
constructor() {
super();
this.state = {
blame_info: null,
};
}
doBlame(arg) {
this.setState({...this.state, blame_info: arg});
}
render({model: {interned_strings}}, {blame_info}) {
let blame_content = "";
if (blame_info) {
const {ist_file, line, ist_s_text, s_start, s_end} = blame_info;
let s_text = interned_strings[ist_s_text];
if (s_start != 0 || s_end != s_text.length) {
let prefix = s_text.slice(0, s_start);
let main = s_text.slice(s_start, s_end);
let suffix = s_text.slice(s_end);
s_text = html`${prefix}<strong>${main}</strong>${suffix}`;
}
blame_content = html`
<h3>${interned_strings[ist_file]}:${line}</h3>
<pre>${s_start}:${s_end}</pre>
<pre>${s_text}</pre><br/>
`;
}
return html`
<button onClick=${() => blame.readyBlame()}>Blame Code</button>
<br/>
${blame_content}
`;
}
}
class App extends Component {
constructor() {
super();
this.state = {
err: false,
model: null,
};
}
componentDidMount() {
const app = this;
if (BURNED_IN_MODEL_INFO !== null) {
app.setState({model: BURNED_IN_MODEL_INFO});
} else {
fetch("./model_info.json").then(function(response) {
if (!response.ok) {
throw new Error("Response not ok.");
}
return response.json();
}).then(function(body) {
app.setState({model: body});
}).catch(function(error) {
console.log("Top-level error: ", error);
});
}
}
componentDidCatch(error) {
void(error);
this.setState({...this.state, err: true});
}
render(_, {err}) {
if (this.state.model === null) {
return html`<h1>Loading...</h1>`;
}
const model = this.state.model.model;
let error_msg = "";
if (err) {
error_msg = html`<h2 style="background:red">An error occurred. Check console</h2>`;
}
return html`
${error_msg}
<div id=main_content style="position:absolute;width:99%;height:79%;overflow:scroll">
<h1>TorchScript Model (version ${model.version}): ${model.title}</h1>
<button onClick=${() => console.log(model)}>Log Raw Model Info</button>
<${ModelSizeSection} model=${model}/>
<${StructuredDataSection} name="Model Data" data=${model.model_data} shown=true/>
<${StructuredDataSection} name="Constants" data=${model.constants} shown=false/>
<${ZipContentsSection} model=${model}/>
<${CodeSection} model=${model}/>
<${ExtraJsonSection} files=${model.extra_files_jsons}/>
<${ExtraPicklesSection} files=${model.extra_pickles}/>
<${TensorMemorySection} model=${model}/>
</div>
<div id=aux_content style="position:absolute;width:99%;top:80%;height:20%;overflow:scroll">
<${AuxContentPane}
err=${this.state.error}
model=${model}
ref=${(p) => blame.setAuxContentPane(p)}/>
</div>
`;
}
}
render(h(App), document.body);