116 lines
4.8 KiB
Python
116 lines
4.8 KiB
Python
import re
|
|
import threading
|
|
|
|
import torch
|
|
|
|
from core.lens_manager import MASKS_DIR, _vocab_fingerprint
|
|
|
|
# ── EVALUATION TOGGLE ────────────────────────────────────────────────────────
|
|
# Nearest tokens (the "translation" of a non-latin token to readable neighbors):
|
|
# True = keep only ENGLISH words (pure ASCII, no accents) as targets
|
|
# False = any readable latin script (accents included: fr/de/es…)
|
|
# Set to True by default; flip it to compare.
|
|
ENGLISH_ONLY = True
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
# "translation" targets: readable tokens (2+ letter word, apostrophe/hyphen
|
|
# allowed) so the neighbors are interpretable
|
|
_LATIN_RE = re.compile(r"^[ A-Za-zÀ-ɏ'\-]+$")
|
|
_LATIN_LETTERS_RE = re.compile(r"[A-Za-zÀ-ɏ]{2}")
|
|
# english variant: pure ASCII (excludes café, über, naïve… → filters out the
|
|
# other latin-script languages)
|
|
_ENGLISH_RE = re.compile(r"^[ A-Za-z'\-]+$")
|
|
_ENGLISH_LETTERS_RE = re.compile(r"[A-Za-z]{2}")
|
|
|
|
|
|
def _latin_target_mask(tokenizer, vocab_size):
|
|
MASKS_DIR.mkdir(parents=True, exist_ok=True)
|
|
# distinct cache per mode (otherwise a "latin" mask would serve in english mode)
|
|
tag = "english" if ENGLISH_ONLY else "latin"
|
|
word_re = _ENGLISH_RE if ENGLISH_ONLY else _LATIN_RE
|
|
letters_re = _ENGLISH_LETTERS_RE if ENGLISH_ONLY else _LATIN_LETTERS_RE
|
|
path = MASKS_DIR / f"{_vocab_fingerprint(tokenizer)}_{vocab_size}_{tag}.pt"
|
|
if path.exists():
|
|
return torch.load(path, weights_only=True)
|
|
mask = torch.zeros(vocab_size, dtype=torch.bool)
|
|
n_decodable = min(vocab_size, len(tokenizer))
|
|
decoded = tokenizer.batch_decode(
|
|
[[tid] for tid in range(n_decodable)], clean_up_tokenization_spaces=False
|
|
)
|
|
for tid, raw in enumerate(decoded):
|
|
s = raw.strip()
|
|
mask[tid] = bool(
|
|
len(s) >= 2 and word_re.match(s) and letters_re.search(s)
|
|
)
|
|
torch.save(mask, path)
|
|
return mask
|
|
|
|
|
|
class TokenNeighbors:
|
|
"""Approximate local translation: latin tokens whose output direction (row of
|
|
W_U) is closest in cosine to a non-latin token most often carry the same
|
|
meaning (答案 → ' answer')."""
|
|
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
self._key = None
|
|
self._mask = None
|
|
self._norms = None
|
|
self._cache = {}
|
|
|
|
def _prepare(self, jl, tokenizer, model_key):
|
|
if self._key == model_key and self._norms is not None:
|
|
return
|
|
weight = jl._lm_head.weight
|
|
if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
|
raise ValueError("neighbors unavailable on a quantized model")
|
|
vocab_size = weight.shape[0]
|
|
self._mask = _latin_target_mask(tokenizer, vocab_size).to(weight.device)
|
|
norms = torch.empty(vocab_size, dtype=torch.float32, device=weight.device)
|
|
with torch.no_grad():
|
|
for start in range(0, vocab_size, 8192):
|
|
chunk = weight[start:start + 8192].float()
|
|
norms[start:start + 8192] = chunk.norm(dim=1)
|
|
self._norms = norms.clamp_min(1e-8)
|
|
self._cache = {}
|
|
self._key = model_key
|
|
|
|
def lookup(self, jl, tokenizer, model_key, token_ids, k=3):
|
|
with self._lock:
|
|
self._prepare(jl, tokenizer, model_key)
|
|
weight = jl._lm_head.weight
|
|
out = {}
|
|
for tid in token_ids:
|
|
tid = int(tid)
|
|
if tid < 0 or tid >= weight.shape[0]:
|
|
out[tid] = []
|
|
continue
|
|
if tid in self._cache:
|
|
out[tid] = self._cache[tid]
|
|
continue
|
|
with torch.no_grad():
|
|
v = weight[tid]
|
|
sims = (weight @ v).float() / (self._norms * self._norms[tid])
|
|
sims[~self._mask] = float("-inf")
|
|
sims[tid] = float("-inf")
|
|
top = torch.topk(sims, min(k, int(self._mask.sum())))
|
|
entries = [
|
|
{
|
|
"id": int(i),
|
|
"str": tokenizer.decode([int(i)]),
|
|
"sim": round(float(s), 3),
|
|
}
|
|
for s, i in zip(top.values.tolist(), top.indices.tolist())
|
|
if s != float("-inf")
|
|
]
|
|
self._cache[tid] = entries
|
|
out[tid] = entries
|
|
return out
|
|
|
|
def reset(self):
|
|
with self._lock:
|
|
self._key = None
|
|
self._mask = None
|
|
self._norms = None
|
|
self._cache = {}
|