Repository: SuperDuperDB/superduperdb
File: superduperdb/vector_search/in_memory.py (Python)

Code before:
        measure: t.Union[str, t.Callable] = 'cosine',
    ):
        self.identifier = identifier
        self.dimensions = dimensions
        self._cache: t.Sequence[VectorItem] = []
        self._CACHE_SIZE = 10000

        self.measure = measure
        if isinstance(measure, str):
            self.measure = measures[measure]

        if h is not None:
            assert index is not None
            self._setup(h, index)
        else:
            self.h = None
            self.index = None
            self.lookup = None

        self.identifier = identifier

    def __len__(self):
        return self.h.shape[0]

    def _setup(self, h, index):
        h = numpy.array(h) if not isinstance(h, numpy.ndarray) else h

        if self.measure == 'cosine':
            # Normalization is required for cosine, hence preparing
            # all vectors in advance.
            h = h / numpy.linalg.norm(h, axis=1)[:, None]
        self.h = h
        self.index = index
        self.lookup = dict(zip(index, range(len(index))))

    def find_nearest_from_id(self, _id, n=100):
        self.post_create()
        return self.find_nearest_from_array(self.h[self.lookup[_id]], n=n)

    def find_nearest_from_array(self, h, n=100, within_ids=None):
        self.post_create()
        h = self.to_numpy(h)[None, :]
        if within_ids:
            ix = list(map(self.lookup.__getitem__, within_ids))
            similarities = self.measure(h, self.h[ix, :])  # mypy: ignore
        else:
            similarities = self.measure(h, self.h)  # mypy: ignore
        similarities = similarities[0, :]
        logging.debug(similarities)
        scores = -numpy.sort(-similarities)
        ## different ways of handling

Review comment (suggestion):
Please use `self.h is not None`, because this will raise an error when self.h is `numpy.array`

Code after:
        measure: t.Union[str, t.Callable] = 'cosine',
    ):
        self.identifier = identifier
        self.dimensions = dimensions
        self._cache: t.Sequence[VectorItem] = []
        self._CACHE_SIZE = 10000

        self.measure = measure
        if isinstance(measure, str):
            self.measure = measures[measure]

        if h is not None:
            assert index is not None
            self._setup(h, index)
        else:
            self.h = None
            self.index = None
            self.lookup = None

        self.identifier = identifier

    def __len__(self):
        if self.h is not None:
            return self.h.shape[0]
        else:
            return 0

    def _setup(self, h, index):
        h = numpy.array(h) if not isinstance(h, numpy.ndarray) else h

        if self.measure == 'cosine':
            # Normalization is required for cosine, hence preparing
            # all vectors in advance.
            h = h / numpy.linalg.norm(h, axis=1)[:, None]
        self.h = h
        self.index = index
        self.lookup = dict(zip(index, range(len(index))))

    def find_nearest_from_id(self, _id, n=100):
        self.post_create()
        return self.find_nearest_from_array(self.h[self.lookup[_id]], n=n)

    def find_nearest_from_array(self, h, n=100, within_ids=None):
        self.post_create()

        if self.h is None:
            logging.error(
                'Tried to search on an empty vector database',
                'Vectors are not yet loaded in vector database.',
                '\nPlease check if model outputs are ready.',
            )