"""Read individual embedding rows without loading an entire LLM into RAM.""" import json import struct from pathlib import Path import numpy as np def embedding_name(name): return any(s in name for s in ('token_embd.weight', 'embed_tokens.weight', 'word_embeddings.weight', 'wte.weight', 'shared.weight', 'tok_embeddings.weight', 'token_embedding.weight')) class Model: def __init__(self, location): self.path = Path(location).expanduser().resolve() self.tensors = {} self.vocab = None self.reader = None if self.path.suffix.lower() == '.gguf': from gguf import GGUFReader self.reader = GGUFReader(str(self.path)) for t in self.reader.tensors: if len(t.shape) == 2: self.tensors[t.name] = ('gguf', t, tuple(int(v) for v in reversed(t.shape))) field = self.reader.fields.get('tokenizer.ggml.tokens') if field is not None: self.vocab = field.contents() else: files = sorted(self.path.rglob('*.safetensors')) if self.path.is_dir() else [self.path] if not files: raise ValueError('No safetensors files found in that folder.') for file in files: if file.suffix != '.safetensors': raise ValueError('Choose a GGUF file, safetensors file, or model folder.') with file.open('rb') as f: raw = f.read(8) if len(raw) != 8: raise ValueError('Invalid safetensors header.') size = struct.unpack(' min(file.stat().st_size - 8, 100_000_000): raise ValueError('Invalid safetensors header size.') header = json.loads(f.read(size)) for name, info in header.items(): if name != '__metadata__' and len(info['shape']) == 2: self.tensors[name] = ('safe', (file, 8 + size, info), tuple(info['shape'])) if not self.tensors: raise ValueError('No two-dimensional weight tensors found.') def close(self): """Release the GGUF mapping when this session unloads its model.""" reader = self.reader self.reader = None self.tensors.clear() self.vocab = None if reader is not None: reader.tensors.clear() reader.fields.clear() reader.data._mmap.close() def row(self, name, token_id): return self.rows(name, token_id, token_id + 1)[0] def rows(self, name, start, stop): """Read a bounded contiguous batch, dequantizing only these rows.""" kind, source, shape = self.tensors[name] if not 0 <= start < stop <= shape[0]: raise ValueError(f'Token range must be within 0 to {shape[0] - 1}.') if kind == 'gguf': from gguf.quants import dequantize values = dequantize(source.data[start:stop], source.tensor_type) return np.asarray(values, dtype=np.float32).reshape(stop - start, shape[1]).copy() file, base, info = source types = {'F64': ' base + info['data_offsets'][1] or end > file.stat().st_size: raise ValueError('Tensor data is truncated or its shape is invalid.') values = np.memmap(file, mode='r', dtype=dtype, offset=offset, shape=(stop - start, shape[1])) if info['dtype'] == 'BF16': return (values.astype(np.uint32) << 16).view(np.float32) return np.array(values, dtype=np.float32) def nearest_tokens(read_rows, row_count, token_id, count=10, batch_size=2048, progress=None): """Exact cosine search over all rows; skip the query, zero and nonfinite rows. Memory is one embedding batch plus one scalar score per vocabulary row. The callback receives (rows_scanned, total_rows). """ if not 0 <= token_id < row_count: raise ValueError('Source token ID is outside the vocabulary.') query = read_rows(token_id, token_id + 1)[0] return search_embeddings(read_rows, row_count, query, count=count, exclude_ids=[token_id], batch_size=batch_size, progress=progress)['neighbours'] def search_embeddings(read_rows, row_count, query, count=10, exclude_ids=(), report_id=None, batch_size=2048, progress=None): """Search a supplied vector, optionally reporting a target's full-vocabulary rank.""" if count < 1 or batch_size < 1: raise ValueError('Neighbour count and batch size must be positive.') query = np.asarray(query, dtype=np.float64) if query.ndim != 1: raise ValueError('The query must be a one-dimensional vector.') norm = np.linalg.norm(query) if not np.isfinite(norm) or norm == 0: raise ValueError('The query has a zero or nonfinite embedding, so cosine similarity is undefined. Choose another vector.') query = query / norm scores = np.full(row_count, -np.inf, dtype=np.float64) for start in range(0, row_count, batch_size): stop = min(start + batch_size, row_count) values = np.asarray(read_rows(start, stop)) if values.shape != (stop - start, len(query)): raise ValueError('Query dimensions do not match the embedding tensor.') with np.errstate(invalid='ignore', over='ignore', divide='ignore'): norms = np.sqrt(np.einsum('ij,ij->i', values, values, dtype=np.float64)) valid = np.isfinite(norms) & (norms > 0) dot = np.einsum('ij,j->i', values, query, dtype=np.float64) batch_scores = np.divide(dot, norms, out=np.full(stop - start, -np.inf), where=valid) scores[start:stop] = np.where(valid & np.isfinite(batch_scores), np.clip(batch_scores, -1, 1), -np.inf) if progress is not None: progress(stop, row_count) for i in exclude_ids: if not isinstance(i, (int, np.integer)) or not 0 <= i < row_count: raise ValueError('An excluded token ID is outside the vocabulary.') scores[i] = -np.inf candidates = np.flatnonzero(np.isfinite(scores)) # Stable tie-breaking by ID makes duplicate-vector results reproducible. ordered = candidates[np.lexsort((candidates, -scores[candidates]))] rank = None if report_id is not None: positions = np.flatnonzero(ordered == report_id) rank = int(positions[0]) + 1 if len(positions) else None return {'neighbours': [(int(i), float(scores[i])) for i in ordered[:count]], 'expected_rank': rank, 'expected_score': float(scores[report_id]) if rank is not None else None, 'valid_count': len(ordered)} def load_tokenizer(path): from tokenizers import Tokenizer p = Path(path).expanduser() if p.is_dir(): p = p / 'tokenizer.json' return Tokenizer.from_file(str(p))