Spaces:
Running
Running
| """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('<Q', raw)[0] | |
| if size > 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': '<f8', 'F32': '<f4', 'F16': '<f2', 'BF16': '<u2'} | |
| if info['dtype'] not in types: | |
| raise ValueError(f"Unsupported safetensors dtype: {info['dtype']}. Packed quantized safetensors need conversion first.") | |
| dtype = np.dtype(types[info['dtype']]) | |
| offset = base + info['data_offsets'][0] + start * shape[1] * dtype.itemsize | |
| end = offset + (stop - start) * shape[1] * dtype.itemsize | |
| if end > 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)) | |