Create embeddings.py

embeddings.py

@@ -0,0 +1,150 @@
+"""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))