Blame for embeddings.py

"""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))