Spaces:
Running
Running
File size: 2,412 Bytes
2eec02e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | from __future__ import annotations
from pathlib import Path
import gradio as gr
import numpy as np
import plotly.graph_objects as go
import torch
from model import BitMLP
from packing import load_binary_model
from safetensors.torch import load_file
from sklearn.datasets import load_digits
ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "bitforge-1bit"
FP32 = BitMLP("fp32")
FP32.load_state_dict(load_file(ARTIFACT_DIR / "fp32.safetensors"))
BINARY = load_binary_model(ARTIFACT_DIR / "binary_weights.npz")
TERNARY = BitMLP("ternary")
TERNARY.load_state_dict(load_file(ARTIFACT_DIR / "ternary_qat.safetensors"))
for model in [FP32, BINARY, TERNARY]:
model.eval()
DIGITS = load_digits()
IMAGES = (DIGITS.images / 16.0).astype(np.float32)
LABELS = DIGITS.target
def compare(index: int) -> tuple[go.Figure, dict]:
index = int(index) % len(IMAGES)
image = torch.from_numpy(IMAGES[index][None])
predictions = {}
with torch.inference_mode():
for name, model in [
("FP32", FP32),
("Binary packed", BINARY),
("Ternary", TERNARY),
]:
predictions[name] = torch.softmax(model(image), dim=1).numpy()[0]
figure = go.Figure()
for name, probabilities in predictions.items():
figure.add_trace(
go.Bar(
name=name,
x=list(range(10)),
y=probabilities,
)
)
figure.update_layout(
title=f"Digit {int(LABELS[index])}: precision variants",
xaxis_title="Predicted class",
yaxis_title="Probability",
barmode="group",
template="plotly_dark",
)
return figure, {
name: {
"prediction": int(values.argmax()),
"confidence": round(float(values.max()), 4),
}
for name, values in predictions.items()
}
with gr.Blocks(title="BitForge 1-bit") as demo:
gr.Markdown(
"# BitForge 1-bit\n"
"Compare full-precision, packed one-bit matrix weights, and ternary weights "
"on the same digit."
)
index = gr.Slider(0, len(IMAGES) - 1, 0, step=1, label="Digit sample")
run = gr.Button("Compare precision", variant="primary")
chart = gr.Plot()
results = gr.JSON()
run.click(compare, index, [chart, results])
demo.load(compare, index, [chart, results])
if __name__ == "__main__":
demo.launch()
|