File size: 3,105 Bytes
00f5885 | 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 80 81 82 83 84 85 86 87 88 |
import json
import torch
from safetensors.torch import load_file
from transformers import AutoModel, AutoTokenizer
def _dequantize(sd, cfg):
"""Dekuantisasi tensor int8/int4 (per-group) kembali ke float16."""
bits = cfg["bits"]
group_size = cfg["group_size"]
shapes = cfg["shapes"]
out = {}
for name, shape in shapes.items():
n_orig = int(torch.tensor(shape).prod())
n_pad = ((n_orig + group_size - 1) // group_size) * group_size
scale = sd[name + ".scale"].float()
zero = sd[name + ".zero"].float()
n_groups = n_pad // group_size
scale_e = scale.reshape(n_groups, 1).expand(n_groups, group_size).reshape(-1)
zero_e = zero.reshape(n_groups, 1).expand(n_groups, group_size).reshape(-1)
if bits == 8:
q = sd[name].float()
w = (q - zero_e) * scale_e
else:
q = sd[name].to(torch.uint8).cpu()
lo = (q & 0x0F).to(torch.int8)
hi = (q >> 4).to(torch.int8)
vals = torch.stack([hi, lo], dim=1).reshape(-1)
vals = vals - 8
w = (vals.float() - zero_e) * scale_e
w = w[:n_orig].reshape(shape)
out[name] = w.half()
return out
def load_quantized_model(repo_id, bits, device=None):
"""Muat model semantic-lite-2 versi terkuantisasi (8 atau 4 bit).
Args:
repo_id: id repo Hugging Face (misal "ukung/semantic-lite-2")
bits: 8 atau 4
device: "cuda" atau "cpu" (default otomatis)
"""
assert bits in (8, 4), "bits harus 8 atau 4"
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
# 1. Muat model fp16 dasar (struktur + bias/norm kecil)
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
model = model.to(device)
# 2. Muat tensor terkuantisasi
from huggingface_hub import hf_hub_download
q_path = hf_hub_download(repo_id, f"model_q{bits}.safetensors")
c_path = hf_hub_download(repo_id, f"model_q{bits}_config.json")
sd = load_file(q_path)
cfg = json.load(open(c_path))
deq = _dequantize(sd, cfg)
# 3. Ganti weight 2D dengan versi dekuantisasi
with torch.no_grad():
for name, p in model.named_parameters():
if name in deq:
p.data = deq[name].to(device)
return model
def encode(model, tokenizer, texts, device=None):
"""Encode teks jadi vektor embedding (sudah ternormalisasi L2)."""
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
tok = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
tok = {k: v.to(device) for k, v in tok.items()}
with torch.no_grad():
out = model(**tok)
return out
if __name__ == "__main__":
import sys
bits = int(sys.argv[1]) if len(sys.argv) > 1 else 8
repo = "ukung/semantic-lite-2"
model = load_quantized_model(repo, bits)
tok = AutoTokenizer.from_pretrained(repo)
v = encode(model, tok, ["contoh kalimat bahasa Indonesia"])
print(f"Q{bits} loaded, output shape: {tuple(v.shape)}")
|