File size: 2,343 Bytes
df41178 | 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 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
bake_raw.py -- 次元を落とさずに焼く
次元圧縮（射影）は、掛け算を減らすためにやっていたが、
実測で意味の差が潰れることが分かった。
そこで圧縮はやめ、int8 に量子化するだけにする。
・品質の損失は量子化ぶんだけ（ごくわずか）
・大きさは 1/2（BF16 → int8）
・全部をメモリに載せず、必要な行だけ mmap で読む
"""
import os, json, math, time, re, sys
from embed_cards import EmbedCards
HERE = os.path.dirname(os.path.abspath(__file__))
def bake(tag, tokfile, tok_kind, out, pattern=None, report=4000):
e = EmbedCards(tag=tag, tokfile=tokfile, tok_kind=tok_kind)
inv = {}
for b, i in e.vocab.items():
try: inv[i] = b.decode("utf-8")
except Exception: pass
targets = []
for tid, w in inv.items():
w = w.replace("▁", "")
if len(w) < 2: continue
if pattern and not pattern.search(w): continue
targets.append((tid, w))
print(f" 対象 {len(targets)} 語 × {e.dim} 次元", flush=True)
idx, seen, t0 = {}, set(), time.time()
with open(out + ".bin", "wb") as f:
for n, (tid, w) in enumerate(targets, 1):
if w in seen: continue
v = e.row(tid)
if not v: continue
nrm = math.sqrt(sum(x * x for x in v)) or 1.0
f.write(bytes(max(0, min(255, int(round(x / nrm * 127)) + 128)) for x in v))
idx[w] = len(seen); seen.add(w)
if n % report == 0:
el = time.time() - t0
print(f" {n}/{len(targets)} ({el:.0f}秒, 残り約{el/n*(len(targets)-n):.0f}秒)",
flush=True)
json.dump({"meta": {"source": tag, "dim": e.dim, "quant": "int8",
"count": len(idx)}, "index": idx},
open(out, "w", encoding="utf-8"), ensure_ascii=False)
print(f" → {len(idx)}語 索引{os.path.getsize(out)/1e6:.1f}MB "
f"+ 本体{os.path.getsize(out+'.bin')/1e6:.0f}MB ({time.time()-t0:.0f}秒)", flush=True)
if __name__ == "__main__":
JP = re.compile(r"[ぁ-んァ-ヴー一-龥]")
bake("llm-jp-3-13b", "llmjp.tokenizer.json", "unigram",
os.path.join(HERE, "jp_cards.json"), pattern=JP)
|