usearch 0.19.0 → 0.19.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -8,6 +8,7 @@ import os
8
8
  from typing import Optional, Union, NamedTuple, List, Iterable
9
9
 
10
10
  import numpy as np
11
+ from tqdm import tqdm
11
12
 
12
13
  from usearch.compiled import Index as _CompiledIndex, IndexMetadata
13
14
  from usearch.compiled import SparseIndex as _CompiledSetsIndex
@@ -33,32 +34,71 @@ SparseIndex = _CompiledSetsIndex
33
34
  Label = np.uint32
34
35
 
35
36
 
36
- def _normalize_dtype(dtype) -> ScalarKind:
37
- if dtype is None:
38
- return ScalarKind.F32
37
+ def _normalize_dtype(dtype, metric: MetricKind = MetricKind.IP) -> ScalarKind:
38
+ if dtype is None or dtype == "":
39
+ return ScalarKind.B1 if metric in MetricKindBitwise else ScalarKind.F32
40
+
39
41
  if isinstance(dtype, ScalarKind):
40
42
  return dtype
41
43
 
42
44
  if isinstance(dtype, str):
45
+ dtype = dtype.lower()
46
+
47
+ _normalize = {
48
+ "f64": ScalarKind.F64,
49
+ "f32": ScalarKind.F32,
50
+ "f16": ScalarKind.F16,
51
+ "f8": ScalarKind.F8,
52
+ "b1": ScalarKind.B1,
53
+ "b1x8": ScalarKind.B1,
54
+ np.float64: ScalarKind.F64,
55
+ np.float32: ScalarKind.F32,
56
+ np.float16: ScalarKind.F16,
57
+ np.int8: ScalarKind.F8,
58
+ }
59
+ return _normalize[dtype]
60
+
61
+
62
+ def _to_numpy_compatible_dtype(dtype: ScalarKind) -> ScalarKind:
63
+ _normalize = {
64
+ ScalarKind.F64: ScalarKind.F64,
65
+ ScalarKind.F32: ScalarKind.F32,
66
+ ScalarKind.F16: ScalarKind.F16,
67
+ ScalarKind.F8: ScalarKind.F16,
68
+ ScalarKind.B1: ScalarKind.B1,
69
+ }
70
+ return _normalize[dtype]
71
+
72
+
73
+ def _to_numpy_dtype(dtype: ScalarKind):
74
+ _normalize = {
75
+ ScalarKind.F64: np.float64,
76
+ ScalarKind.F32: np.float32,
77
+ ScalarKind.F16: np.float16,
78
+ ScalarKind.F8: np.float16,
79
+ ScalarKind.B1: np.uint8,
80
+ }
81
+ return _normalize[dtype]
82
+
83
+
84
+ def _normalize_metric(metric):
85
+ if metric is None:
86
+ return MetricKind.IP
87
+
88
+ if isinstance(metric, str):
43
89
  _normalize = {
44
- "f64": ScalarKind.F64,
45
- "f32": ScalarKind.F32,
46
- "f16": ScalarKind.F16,
47
- "f8": ScalarKind.F8,
48
- "b1": ScalarKind.B1,
49
- }
50
- return _normalize[dtype.lower()]
51
-
52
- if not isinstance(dtype, ScalarKind):
53
- _normalize = {
54
- np.float64: ScalarKind.F64,
55
- np.float32: ScalarKind.F32,
56
- np.float16: ScalarKind.F16,
57
- np.int8: ScalarKind.F8,
90
+ "cos": MetricKind.Cos,
91
+ "ip": MetricKind.IP,
92
+ "l2_sq": MetricKind.L2sq,
93
+ "haversine": MetricKind.Haversine,
94
+ "perason": MetricKind.Pearson,
95
+ "hamming": MetricKind.Hamming,
96
+ "tanimoto": MetricKind.Tanimoto,
97
+ "sorensen": MetricKind.Sorensen,
58
98
  }
59
- return _normalize[dtype]
99
+ return _normalize[metric.lower()]
60
100
 
61
- return dtype
101
+ return metric
62
102
 
63
103
 
64
104
  class Matches(NamedTuple):
@@ -67,7 +107,7 @@ class Matches(NamedTuple):
67
107
  counts: Union[np.ndarray, int]
68
108
 
69
109
  @property
70
- def is_batch(self) -> bool:
110
+ def is_multiple(self) -> bool:
71
111
  return isinstance(self.counts, np.ndarray)
72
112
 
73
113
  @property
@@ -79,7 +119,7 @@ class Matches(NamedTuple):
79
119
  return np.sum(self.counts)
80
120
 
81
121
  def to_list(self, row: Optional[int] = None) -> Union[List[dict], List[List[dict]]]:
82
- if not self.is_batch:
122
+ if not self.is_multiple:
83
123
  assert row is None, "Exporting a single sequence is only for batch requests"
84
124
  labels = self.labels
85
125
  distances = self.distances
@@ -98,13 +138,13 @@ class Matches(NamedTuple):
98
138
  ]
99
139
 
100
140
  def recall_first(self, expected: np.ndarray) -> float:
101
- best_matches = self.labels if not self.is_batch else self.labels[:, 0]
141
+ best_matches = self.labels if not self.is_multiple else self.labels[:, 0]
102
142
  return np.sum(best_matches == expected) / len(expected)
103
143
 
104
144
  def __repr__(self) -> str:
105
145
  return (
106
146
  f"usearch.Matches({self.total_matches})"
107
- if self.is_batch
147
+ if self.is_multiple
108
148
  else f"usearch.Matches({self.total_matches} across {self.batch_size} queries)"
109
149
  )
110
150
 
@@ -190,28 +230,7 @@ class Index:
190
230
  :type view: bool, optional
191
231
  """
192
232
 
193
- if metric in MetricKindBitwise:
194
- assert dtype is None or dtype == ScalarKind.B1
195
- dtype = ScalarKind.B1
196
- else:
197
- dtype = _normalize_dtype(dtype)
198
-
199
- if metric is None:
200
- metric = MetricKind.IP
201
-
202
- if isinstance(metric, str):
203
- _normalize = {
204
- "cos": MetricKind.Cos,
205
- "ip": MetricKind.IP,
206
- "l2_sq": MetricKind.L2sq,
207
- "haversine": MetricKind.Haversine,
208
- "perason": MetricKind.Pearson,
209
- "hamming": MetricKind.Hamming,
210
- "tanimoto": MetricKind.Tanimoto,
211
- "sorensen": MetricKind.Sorensen,
212
- }
213
- metric = _normalize[metric.lower()]
214
-
233
+ metric = _normalize_metric(metric)
215
234
  if isinstance(metric, MetricKind) and jit:
216
235
  try:
217
236
  from usearch.numba import jit
@@ -224,7 +243,7 @@ class Index:
224
243
  metric = jit(
225
244
  ndim=ndim,
226
245
  metric=metric,
227
- dtype=dtype,
246
+ dtype=_normalize_dtype(dtype, metric),
228
247
  )
229
248
 
230
249
  if isinstance(metric, MetricKind):
@@ -242,6 +261,8 @@ class Index:
242
261
  "The `metric` must be a `CompiledMetric` or a `MetricKind`"
243
262
  )
244
263
 
264
+ # Validate, that the right scalar type is defined
265
+ dtype = _normalize_dtype(dtype, self._metric_kind)
245
266
  self._compiled = _CompiledIndex(
246
267
  ndim=ndim,
247
268
  dtype=dtype,
@@ -285,7 +306,14 @@ class Index:
285
306
  )
286
307
 
287
308
  def add(
288
- self, labels, vectors, *, copy: bool = True, threads: int = 0
309
+ self,
310
+ labels,
311
+ vectors,
312
+ *,
313
+ copy: bool = True,
314
+ threads: int = 0,
315
+ log: Union[str, bool] = False,
316
+ batch_size: int = 0,
289
317
  ) -> Union[int, np.ndarray]:
290
318
  """Inserts one or move vectors into the index.
291
319
 
@@ -315,27 +343,67 @@ class Index:
315
343
  """
316
344
  assert isinstance(vectors, np.ndarray), "Expects a NumPy array"
317
345
  assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector"
318
- is_batch = vectors.ndim == 2
319
- generate_labels = labels is None
346
+ count_vectors = vectors.shape[0] if vectors.ndim == 2 else 1
347
+ is_multiple = count_vectors > 1
348
+ if not is_multiple:
349
+ vectors = vectors.flatten()
320
350
 
321
- # If no `labels` were provided, generate some
351
+ # Validate of generate teh labels
352
+ generate_labels = labels is None
322
353
  if generate_labels:
323
354
  start_id = len(self._compiled)
324
- if is_batch:
325
- labels = np.arange(start_id, start_id + vectors.shape[0])
355
+ if is_multiple:
356
+ labels = np.arange(start_id, start_id + count_vectors, dtype=Label)
326
357
  else:
327
358
  labels = start_id
328
- if isinstance(labels, np.ndarray):
329
- labels = labels.astype(Label)
330
- elif isinstance(labels, Iterable):
331
- labels = np.array(labels, dtype=Label)
359
+ else:
360
+ if isinstance(labels, np.ndarray):
361
+ if not is_multiple:
362
+ labels = int(labels[0])
363
+ else:
364
+ labels = labels.astype(Label)
365
+ count_labels = len(labels) if isinstance(labels, Iterable) else 1
366
+ assert count_labels == count_vectors
367
+
368
+ # Split into batches and log progress, if needed
369
+ if batch_size:
370
+ labels = [
371
+ labels[start_row : start_row + batch_size]
372
+ for start_row in range(0, count_vectors, batch_size)
373
+ ]
374
+ vectors = [
375
+ vectors[start_row : start_row + batch_size, :]
376
+ for start_row in range(0, count_vectors, batch_size)
377
+ ]
378
+ tasks = zip(labels, vectors)
379
+ name = log if isinstance(log, str) else "Add"
380
+ pbar = tqdm(
381
+ tasks,
382
+ desc=name,
383
+ total=count_vectors,
384
+ unit="Vector",
385
+ disable=log is False,
386
+ )
387
+ for labels, vectors in tasks:
388
+ self._compiled.add(labels, vectors, copy=copy, threads=threads)
389
+ pbar.update(len(labels))
332
390
 
333
- self._compiled.add(labels, vectors, copy=copy, threads=threads)
391
+ pbar.close()
392
+
393
+ else:
394
+ self._compiled.add(labels, vectors, copy=copy, threads=threads)
334
395
 
335
396
  return labels
336
397
 
337
398
  def search(
338
- self, vectors, k: int = 10, *, threads: int = 0, exact: bool = False
399
+ self,
400
+ vectors,
401
+ k: int = 10,
402
+ *,
403
+ threads: int = 0,
404
+ exact: bool = False,
405
+ log: Union[str, bool] = False,
406
+ batch_size: int = 0,
339
407
  ) -> Matches:
340
408
  """Performs approximate nearest neighbors search for one or more queries.
341
409
 
@@ -427,18 +495,30 @@ class Index:
427
495
  def expansion_search(self, v: int):
428
496
  self._compiled.expansion_search = v
429
497
 
430
- def save(self, path: os.PathLike):
498
+ def save(self, path: Optional[os.PathLike] = None):
499
+ path = path if path else self.path
500
+ if path is None:
501
+ raise Exception("Define `path` argument")
431
502
  self._compiled.save(path)
432
503
 
433
- def load(self, path: os.PathLike):
504
+ def load(self, path: Optional[os.PathLike] = None):
505
+ path = path if path else self.path
506
+ if path is None:
507
+ raise Exception("Define `path` argument")
434
508
  self._compiled.load(path)
435
509
 
436
- def view(self, path: os.PathLike):
510
+ def view(self, path: Optional[os.PathLike] = None):
511
+ path = path if path else self.path
512
+ if path is None:
513
+ raise Exception("Define `path` argument")
437
514
  self._compiled.view(path)
438
515
 
439
516
  def clear(self):
440
517
  self._compiled.clear()
441
518
 
519
+ def close(self):
520
+ self._compiled.close()
521
+
442
522
  @property
443
523
  def labels(self) -> np.ndarray:
444
524
  return self._compiled.labels
@@ -446,12 +526,15 @@ class Index:
446
526
  def get_vectors(
447
527
  self, labels: np.ndarray, dtype: ScalarKind = ScalarKind.F32
448
528
  ) -> np.ndarray:
449
- dtype = _normalize_dtype(dtype)
450
- return np.vstack([self._compiled.__getitem__(l, dtype) for l in labels])
529
+ dtype = _normalize_dtype(dtype, self._metric_kind)
530
+ get_dtype = _to_numpy_compatible_dtype(dtype)
531
+ vectors = np.vstack([self._compiled.__getitem__(l, get_dtype) for l in labels])
532
+ view_dtype = _to_numpy_dtype(dtype)
533
+ return vectors.view(view_dtype)
451
534
 
452
535
  @property
453
536
  def vectors(self) -> np.ndarray:
454
- return self.get_vectors(self.labels)
537
+ return self.get_vectors(self.labels, self.dtype)
455
538
 
456
539
  def __delitem__(self, label: int):
457
540
  raise NotImplementedError()
@@ -460,4 +543,8 @@ class Index:
460
543
  return self._compiled.__contains__(label)
461
544
 
462
545
  def __getitem__(self, label: int) -> np.ndarray:
463
- return self._compiled.__getitem__(label)
546
+ dtype = self.dtype
547
+ get_dtype = _to_numpy_compatible_dtype(dtype)
548
+ vector = self._compiled.__getitem__(label, get_dtype)
549
+ view_dtype = _to_numpy_dtype(dtype)
550
+ return None if vector is None else vector.view(view_dtype)
@@ -31,7 +31,7 @@ def load_matrix(
31
31
  dtype = np.int32
32
32
  scalar_size = 4
33
33
  elif filename.endswith(".bbin"):
34
- dtype = np.int8
34
+ dtype = np.uint8
35
35
  scalar_size = 1
36
36
  else:
37
37
  raise Exception("Unknown file type")
@@ -69,7 +69,7 @@ def save_matrix(vectors: np.ndarray, filename: str):
69
69
  elif filename.endswith(".ibin"):
70
70
  dtype = np.int32
71
71
  elif filename.endswith(".bbin"):
72
- dtype = np.int8
72
+ dtype = np.uint8
73
73
  else:
74
74
  raise Exception("Unknown file type")
75
75
 
package/setup.py CHANGED
@@ -5,97 +5,94 @@ from setuptools import setup
5
5
  from pybind11.setup_helpers import Pybind11Extension
6
6
 
7
7
 
8
- compile_args = ['-std=c++11', '-O3']
8
+ compile_args = []
9
9
  link_args = []
10
10
  macros_args = [
11
- ('USEARCH_USE_NATIVE_F16', '0'),
12
- ('USEARCH_USE_SIMSIMD', '0'),
11
+ ("USEARCH_USE_NATIVE_F16", "0"),
12
+ ("USEARCH_USE_SIMSIMD", "0"),
13
13
  ]
14
14
 
15
- if sys.platform == 'linux':
16
- compile_args.append('-fopenmp')
17
- link_args.append('-lgomp')
18
- compile_args.append('-Wno-unknown-pragmas')
19
- macros_args.append(('USEARCH_USE_OPENMP', '1'))
15
+ if sys.platform == "linux":
16
+ compile_args.append("-std=c++11")
17
+ compile_args.append("-O3")
18
+ compile_args.append("-Wno-unknown-pragmas")
20
19
 
21
- if sys.platform == 'darwin':
20
+ macros_args.append(("USEARCH_USE_OPENMP", "1"))
21
+ compile_args.append("-fopenmp")
22
+ link_args.append("-lgomp")
23
+
24
+ if sys.platform == "darwin":
22
25
  # MacOS 10.15 or higher is needed for `aligned_alloc` support.
23
26
  # https://github.com/unum-cloud/usearch/actions/runs/4975434891/jobs/8902603392
24
- compile_args.append('-mmacosx-version-min=10.15')
25
- compile_args.append('-Wno-unknown-pragmas')
27
+ compile_args.append("-mmacosx-version-min=10.15")
28
+ compile_args.append("-std=c++11")
29
+ compile_args.append("-O3")
30
+ compile_args.append("-Wno-unknown-pragmas")
26
31
 
27
- if sys.platform == 'win32':
28
- compile_args = ['/std:c++14', '/O2']
32
+ # Linking OpenMP requires additional preparion in CIBuildWheel
33
+ # macros_args.append(("USEARCH_USE_OPENMP", "1"))
34
+ # compile_args.append("-Xpreprocessor -fopenmp")
35
+ # link_args.append("-Xpreprocessor -lomp")
29
36
 
37
+ if sys.platform == "win32":
38
+ compile_args.append("/std:c++14")
39
+ compile_args.append("/O2")
30
40
 
31
41
  ext_modules = [
32
42
  Pybind11Extension(
33
- 'usearch.compiled',
34
- ['python/lib.cpp'],
43
+ "usearch.compiled",
44
+ ["python/lib.cpp"],
35
45
  extra_compile_args=compile_args,
36
46
  extra_link_args=link_args,
37
47
  define_macros=macros_args,
38
48
  ),
39
49
  ]
40
50
 
41
- __version__ = open('VERSION', 'r').read().strip()
42
- __lib_name__ = 'usearch'
51
+ __version__ = open("VERSION", "r").read().strip()
52
+ __lib_name__ = "usearch"
43
53
 
44
54
 
45
55
  this_directory = os.path.abspath(os.path.dirname(__file__))
46
- with open(os.path.join(this_directory, 'README.md')) as f:
56
+ with open(os.path.join(this_directory, "README.md")) as f:
47
57
  long_description = f.read()
48
58
 
49
59
 
50
60
  setup(
51
-
52
61
  name=__lib_name__,
53
62
  version=__version__,
54
- packages=['usearch'],
55
- package_dir={'usearch': 'python/usearch'},
56
-
57
- description='Smaller & Faster Single-File Vector Search Engine from Unum',
63
+ packages=["usearch"],
64
+ package_dir={"usearch": "python/usearch"},
65
+ description="Smaller & Faster Single-File Vector Search Engine from Unum",
58
66
  long_description=long_description,
59
- long_description_content_type='text/markdown',
60
- license='Apache-2.0',
61
-
67
+ long_description_content_type="text/markdown",
68
+ license="Apache-2.0",
62
69
  classifiers=[
63
- 'Development Status :: 5 - Production/Stable',
64
-
65
- 'Natural Language :: English',
66
- 'Intended Audience :: Developers',
67
- 'Intended Audience :: Information Technology',
68
- 'License :: OSI Approved :: Apache Software License',
69
-
70
- 'Programming Language :: C++',
71
- 'Programming Language :: Python :: 3 :: Only',
72
- 'Programming Language :: Python :: Implementation :: CPython',
73
- 'Programming Language :: Java',
74
- 'Programming Language :: JavaScript',
75
- 'Programming Language :: Objective C',
76
- 'Programming Language :: Rust',
77
- 'Programming Language :: Other',
78
-
79
- 'Operating System :: MacOS',
80
- 'Operating System :: Unix',
81
- 'Operating System :: Microsoft :: Windows',
82
-
83
- 'Topic :: System :: Clustering',
84
- 'Topic :: Database :: Database Engines/Servers',
85
- 'Topic :: Scientific/Engineering :: Artificial Intelligence',
70
+ "Development Status :: 5 - Production/Stable",
71
+ "Natural Language :: English",
72
+ "Intended Audience :: Developers",
73
+ "Intended Audience :: Information Technology",
74
+ "License :: OSI Approved :: Apache Software License",
75
+ "Programming Language :: C++",
76
+ "Programming Language :: Python :: 3 :: Only",
77
+ "Programming Language :: Python :: Implementation :: CPython",
78
+ "Programming Language :: Java",
79
+ "Programming Language :: JavaScript",
80
+ "Programming Language :: Objective C",
81
+ "Programming Language :: Rust",
82
+ "Programming Language :: Other",
83
+ "Operating System :: MacOS",
84
+ "Operating System :: Unix",
85
+ "Operating System :: Microsoft :: Windows",
86
+ "Topic :: System :: Clustering",
87
+ "Topic :: Database :: Database Engines/Servers",
88
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
86
89
  ],
87
-
88
- include_dirs=[
89
- 'include',
90
- 'fp16/include', 'robin-map/include', 'simsimd/include'],
91
-
90
+ include_dirs=["include", "fp16/include", "robin-map/include", "simsimd/include"],
92
91
  ext_modules=ext_modules,
93
92
  install_requires=[
94
- 'numpy',
95
- 'pandas',
93
+ "numpy",
94
+ "pandas",
95
+ "tqdm",
96
96
  'ucall; python_version >= "3.9"',
97
97
  ],
98
- extras_require={
99
- 'jit': ['numba', 'cppyy'],
100
- },
101
98
  )
package/swift/Test.swift CHANGED
@@ -6,16 +6,20 @@
6
6
  //
7
7
 
8
8
  import Foundation
9
+ import XCTest
9
10
  import USearch
10
11
 
11
12
  @available(iOS 13, macOS 10.15, tvOS 13.0, watchOS 6.0, *)
12
- func testUnit() throws {
13
- let index = USearchIndex.make(metric:USearchMetric.l2sq, dimensions: 4, connectivity: 8, quantization:USearchScalar.F32)
14
- let vectorA: [Float32] = [0.3, 0.5, 1.2, 1.4]
15
- let vectorB: [Float32] = [0.4, 0.2, 1.2, 1.1]
16
- index.add(label: 42, vector: vectorA[...])
17
- index.add(label: 43, vector: vectorB[...])
18
-
19
- let results = index.search(vector: vectorA[...], count: 10)
20
- assert(results.0[0] == 42)
13
+ class Test: XCTestCase {
14
+ func testUnit() throws {
15
+ let index = USearchIndex.make(metric:USearchMetric.l2sq, dimensions: 4, connectivity: 8, quantization:USearchScalar.F32)
16
+ let vectorA: [Float32] = [0.3, 0.5, 1.2, 1.4]
17
+ let vectorB: [Float32] = [0.4, 0.2, 1.2, 1.1]
18
+ index.reserve(2)
19
+ index.add(label: 42, vector: vectorA[...])
20
+ index.add(label: 43, vector: vectorB[...])
21
+
22
+ let results = index.search(vector: vectorA[...], count: 10)
23
+ assert(results.0[0] == 42)
24
+ }
21
25
  }