#!/usr/bin/env python3 """ Convert model weights to UNARY (base-1) thermometer encoding. True unary: magnitude N = N consecutive 1-bits across N bitplanes. Each bitplane contributes equally (value=1), NOT binary powers. Weight 0.3 with scale -> magnitude 5 -> planes 0,1,2,3,4 have bit set Weight -0.1 with scale -> magnitude 2, sign=neg -> planes 0,1 set + sign bit More precision than ternary (N+1 levels vs 3), still no multiplication. (c) 2026 OpenTransformers Ltd / Scott Bisset """ import os import json import numpy as np from pathlib import Path import time def load_safetensors(model_dir): """Load all tensors from safetensors files.""" import torch from safetensors.torch import load_file tensors = {} for f in sorted(Path(model_dir).glob("*.safetensors")): print(f"Loading {f.name}...") state = load_file(str(f)) for key, val in state.items(): tensors[key] = val.float().numpy() return tensors def quantize_matrix_unary(weight, n_planes=7): """Quantize weight matrix to unary thermometer encoding. n_planes determines max magnitude (and precision levels = n_planes + 1). n_planes=7 gives 8 levels: {0,1,2,3,4,5,6,7} * sign = 15 distinct values. Returns: sign_bits, mag_planes, scales, sparsity """ w = weight.astype(np.float32) out_dim, in_dim = w.shape chunks = ((in_dim + 63) // 64) padded = chunks * 64 # Per-row quantization row_max = np.max(np.abs(w), axis=1, keepdims=True) row_max = np.where(row_max == 0, 1.0, row_max) # Scale to [0, n_planes] range per row scales = (row_max.flatten() / n_planes).astype(np.float32) # Quantize to integer magnitudes w_scaled = w / scales[:, None] # Now in [-n_planes, +n_planes] magnitudes = np.round(np.abs(w_scaled)).astype(np.int32) magnitudes = np.clip(magnitudes, 0, n_planes) signs = (w < 0) # True = negative # Sparsity (magnitude 0) sparsity = np.mean(magnitudes == 0) # Pad to multiple of 64 if in_dim < padded: magnitudes = np.concatenate([magnitudes, np.zeros((out_dim, padded - in_dim), dtype=np.int32)], axis=1) signs = np.concatenate([signs, np.zeros((out_dim, padded - in_dim), dtype=bool)], axis=1) # Pack sign bits - vectorized bit_positions = (np.uint64(1) << np.arange(64, dtype=np.uint64)) signs_r = signs.reshape(out_dim, chunks, 64).astype(np.uint64) sign_bits = np.bitwise_or.reduce(signs_r * bit_positions, axis=2) # [out_dim, chunks] # Pack magnitude planes - thermometer encoding # Plane p has bit set where magnitude > p (i.e., magnitude >= p+1) mag_planes = np.zeros((n_planes, out_dim, chunks), dtype=np.uint64) for p in range(n_planes): active = (magnitudes >= (p + 1)) # [out_dim, padded] active_r = active.reshape(out_dim, chunks, 64).astype(np.uint64) mag_planes[p] = np.bitwise_or.reduce(active_r * bit_positions, axis=2) return sign_bits, mag_planes, scales, sparsity def save_unary_model(tensors, output_dir, n_planes=7): """Convert and save full model to unary format.""" os.makedirs(output_dir, exist_ok=True) config = { "hidden_size": 1536, "intermediate_size": 8960, "num_attention_heads": 12, "num_key_value_heads": 2, "num_hidden_layers": 28, "vocab_size": 151936, "head_dim": 128, "rope_theta": 1000000.0, "rms_norm_eps": 1e-6, "n_planes": n_planes, "quant_type": "unary", } ternary_keys = [] keep_keys = [] for key in tensors: if any(p in key for p in ['q_proj.weight', 'k_proj.weight', 'v_proj.weight', 'o_proj.weight', 'gate_proj.weight', 'up_proj.weight', 'down_proj.weight']): ternary_keys.append(key) else: keep_keys.append(key) print(f"\nUnary layers: {len(ternary_keys)} (n_planes={n_planes}, levels={n_planes+1})") print(f"FP16 layers: {len(keep_keys)}") with open(os.path.join(output_dir, "config.json"), "w") as f: json.dump(config, f, indent=2) total_unary_bytes = 0 total_original_bytes = 0 for key in ternary_keys: w = tensors[key] out_dim, in_dim = w.shape total_original_bytes += w.nbytes t0 = time.time() sign_bits, mag_planes, scales, sparsity = quantize_matrix_unary(w, n_planes) dt = time.time() - t0 prefix = os.path.join(output_dir, key.replace(".", "_")) sign_bits.tofile(prefix + ".sign") mag_planes.tofile(prefix + ".planes") scales.tofile(prefix + ".scales") unary_bytes = sign_bits.nbytes + mag_planes.nbytes + scales.nbytes total_unary_bytes += unary_bytes ratio = w.nbytes / unary_bytes # Calculate effective bits per weight bpw = (unary_bytes * 8) / (out_dim * in_dim) print(f" {key}: {w.shape} -> unary ({unary_bytes/1024:.0f}KB, " f"{ratio:.1f}x compress, {bpw:.2f} bpw, {sparsity:.1%} sparse, {dt:.1f}s)") total_fp16_bytes = 0 for key in keep_keys: w = tensors[key].astype(np.float16) prefix = os.path.join(output_dir, key.replace(".", "_")) w.tofile(prefix + ".fp16") total_fp16_bytes += w.nbytes print(f" {key}: {w.shape} -> fp16 ({w.nbytes/1024:.0f}KB)") manifest = { "unary": {k: list(tensors[k].shape) for k in ternary_keys}, "fp16": {k: list(tensors[k].shape) for k in keep_keys}, } with open(os.path.join(output_dir, "manifest.json"), "w") as f: json.dump(manifest, f, indent=2) total_bytes = total_unary_bytes + total_fp16_bytes avg_bpw = (total_unary_bytes * 8) / sum(np.prod(tensors[k].shape) for k in ternary_keys) print(f"\n=== Summary ===") print(f"Original FP32 linear weights: {total_original_bytes/1024/1024:.1f} MB") print(f"Unary linear weights: {total_unary_bytes/1024/1024:.1f} MB") print(f"FP16 other weights: {total_fp16_bytes/1024/1024:.1f} MB") print(f"Total model size: {total_bytes/1024/1024:.1f} MB") print(f"Average bits per weight (linear): {avg_bpw:.2f}") print(f"Compression vs FP32: {(total_original_bytes + total_fp16_bytes)/total_bytes:.1f}x") print(f"Precision levels: {n_planes + 1} (vs ternary=3, INT4=16)") if __name__ == "__main__": import sys model_dir = sys.argv[1] if len(sys.argv) > 1 else "deepseek-r1-1.5b-hf" output_dir = sys.argv[2] if len(sys.argv) > 2 else "deepseek-r1-1.5b-unary" n_planes = int(sys.argv[3]) if len(sys.argv) > 3 else 7 print(f"Loading model from {model_dir}...") tensors = load_safetensors(model_dir) print(f"Converting to unary (n_planes={n_planes})...") save_unary_model(tensors, output_dir, n_planes) print("Done!")