pgvector 0.4.1__tar.gz → 0.4.2__tar.gz
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.
- {pgvector-0.4.1/pgvector.egg-info → pgvector-0.4.2}/PKG-INFO +37 -7
- pgvector-0.4.1/PKG-INFO → pgvector-0.4.2/README.md +35 -18
- pgvector-0.4.2/pgvector/django/extensions.py +11 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/sqlalchemy/bit.py +12 -0
- pgvector-0.4.1/README.md → pgvector-0.4.2/pgvector.egg-info/PKG-INFO +48 -5
- {pgvector-0.4.1 → pgvector-0.4.2}/pyproject.toml +2 -2
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_sqlalchemy.py +21 -2
- pgvector-0.4.1/pgvector/django/extensions.py +0 -6
- {pgvector-0.4.1 → pgvector-0.4.2}/LICENSE.txt +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/asyncpg/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/asyncpg/register.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/bit.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/django/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/django/bit.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/django/functions.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/django/halfvec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/django/indexes.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/django/sparsevec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/django/vector.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/halfvec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/peewee/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/peewee/bit.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/peewee/halfvec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/peewee/sparsevec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/peewee/vector.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/pg8000/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/pg8000/register.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg/bit.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg/halfvec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg/register.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg/sparsevec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg/vector.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg2/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg2/halfvec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg2/register.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg2/sparsevec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/psycopg2/vector.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/sparsevec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/sqlalchemy/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/sqlalchemy/functions.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/sqlalchemy/halfvec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/sqlalchemy/sparsevec.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/sqlalchemy/vector.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/utils/__init__.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector/vector.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector.egg-info/SOURCES.txt +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector.egg-info/dependency_links.txt +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector.egg-info/requires.txt +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/pgvector.egg-info/top_level.txt +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/setup.cfg +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_asyncpg.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_bit.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_django.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_half_vector.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_peewee.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_pg8000.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_psycopg.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_psycopg2.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_sparse_vector.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_sqlmodel.py +0 -0
- {pgvector-0.4.1 → pgvector-0.4.2}/tests/test_vector.py +0 -0
|
@@ -1,9 +1,9 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: pgvector
|
|
3
|
-
Version: 0.4.
|
|
3
|
+
Version: 0.4.2
|
|
4
4
|
Summary: pgvector support for Python
|
|
5
5
|
Author-email: Andrew Kane <andrew@ankane.org>
|
|
6
|
-
License: MIT
|
|
6
|
+
License-Expression: MIT
|
|
7
7
|
Project-URL: Homepage, https://github.com/pgvector/pgvector-python
|
|
8
8
|
Requires-Python: >=3.9
|
|
9
9
|
Description-Content-Type: text/markdown
|
|
@@ -272,7 +272,6 @@ index = Index(
|
|
|
272
272
|
'my_index',
|
|
273
273
|
func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
|
|
274
274
|
postgresql_using='hnsw',
|
|
275
|
-
postgresql_with={'m': 16, 'ef_construction': 64},
|
|
276
275
|
postgresql_ops={'embedding': 'halfvec_l2_ops'}
|
|
277
276
|
)
|
|
278
277
|
```
|
|
@@ -284,6 +283,37 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])
|
|
|
284
283
|
session.scalars(select(Item).order_by(order).limit(5))
|
|
285
284
|
```
|
|
286
285
|
|
|
286
|
+
#### Binary Quantization
|
|
287
|
+
|
|
288
|
+
Use expression indexing for binary quantization
|
|
289
|
+
|
|
290
|
+
```python
|
|
291
|
+
from pgvector.sqlalchemy import BIT
|
|
292
|
+
from sqlalchemy.sql import func
|
|
293
|
+
|
|
294
|
+
index = Index(
|
|
295
|
+
'my_index',
|
|
296
|
+
func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
|
|
297
|
+
postgresql_using='hnsw',
|
|
298
|
+
postgresql_ops={'embedding': 'bit_hamming_ops'}
|
|
299
|
+
)
|
|
300
|
+
```
|
|
301
|
+
|
|
302
|
+
Get the nearest neighbors by Hamming distance
|
|
303
|
+
|
|
304
|
+
```python
|
|
305
|
+
order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
|
|
306
|
+
session.scalars(select(Item).order_by(order).limit(5))
|
|
307
|
+
```
|
|
308
|
+
|
|
309
|
+
Re-rank by the original vectors for better recall
|
|
310
|
+
|
|
311
|
+
```python
|
|
312
|
+
order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
|
|
313
|
+
subquery = session.query(Item).order_by(order).limit(20).subquery()
|
|
314
|
+
session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5))
|
|
315
|
+
```
|
|
316
|
+
|
|
287
317
|
#### Arrays
|
|
288
318
|
|
|
289
319
|
Add an array column
|
|
@@ -422,7 +452,7 @@ Enable the extension
|
|
|
422
452
|
conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
423
453
|
```
|
|
424
454
|
|
|
425
|
-
Register the
|
|
455
|
+
Register the types with your connection
|
|
426
456
|
|
|
427
457
|
```python
|
|
428
458
|
from pgvector.psycopg import register_vector
|
|
@@ -485,7 +515,7 @@ cur = conn.cursor()
|
|
|
485
515
|
cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
486
516
|
```
|
|
487
517
|
|
|
488
|
-
Register the
|
|
518
|
+
Register the types with your connection or cursor
|
|
489
519
|
|
|
490
520
|
```python
|
|
491
521
|
from pgvector.psycopg2 import register_vector
|
|
@@ -531,7 +561,7 @@ Enable the extension
|
|
|
531
561
|
await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
532
562
|
```
|
|
533
563
|
|
|
534
|
-
Register the
|
|
564
|
+
Register the types with your connection
|
|
535
565
|
|
|
536
566
|
```python
|
|
537
567
|
from pgvector.asyncpg import register_vector
|
|
@@ -585,7 +615,7 @@ Enable the extension
|
|
|
585
615
|
conn.run('CREATE EXTENSION IF NOT EXISTS vector')
|
|
586
616
|
```
|
|
587
617
|
|
|
588
|
-
Register the
|
|
618
|
+
Register the types with your connection
|
|
589
619
|
|
|
590
620
|
```python
|
|
591
621
|
from pgvector.pg8000 import register_vector
|
|
@@ -1,16 +1,3 @@
|
|
|
1
|
-
Metadata-Version: 2.4
|
|
2
|
-
Name: pgvector
|
|
3
|
-
Version: 0.4.1
|
|
4
|
-
Summary: pgvector support for Python
|
|
5
|
-
Author-email: Andrew Kane <andrew@ankane.org>
|
|
6
|
-
License: MIT
|
|
7
|
-
Project-URL: Homepage, https://github.com/pgvector/pgvector-python
|
|
8
|
-
Requires-Python: >=3.9
|
|
9
|
-
Description-Content-Type: text/markdown
|
|
10
|
-
License-File: LICENSE.txt
|
|
11
|
-
Requires-Dist: numpy
|
|
12
|
-
Dynamic: license-file
|
|
13
|
-
|
|
14
1
|
# pgvector-python
|
|
15
2
|
|
|
16
3
|
[pgvector](https://github.com/pgvector/pgvector) support for Python
|
|
@@ -272,7 +259,6 @@ index = Index(
|
|
|
272
259
|
'my_index',
|
|
273
260
|
func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
|
|
274
261
|
postgresql_using='hnsw',
|
|
275
|
-
postgresql_with={'m': 16, 'ef_construction': 64},
|
|
276
262
|
postgresql_ops={'embedding': 'halfvec_l2_ops'}
|
|
277
263
|
)
|
|
278
264
|
```
|
|
@@ -284,6 +270,37 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])
|
|
|
284
270
|
session.scalars(select(Item).order_by(order).limit(5))
|
|
285
271
|
```
|
|
286
272
|
|
|
273
|
+
#### Binary Quantization
|
|
274
|
+
|
|
275
|
+
Use expression indexing for binary quantization
|
|
276
|
+
|
|
277
|
+
```python
|
|
278
|
+
from pgvector.sqlalchemy import BIT
|
|
279
|
+
from sqlalchemy.sql import func
|
|
280
|
+
|
|
281
|
+
index = Index(
|
|
282
|
+
'my_index',
|
|
283
|
+
func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
|
|
284
|
+
postgresql_using='hnsw',
|
|
285
|
+
postgresql_ops={'embedding': 'bit_hamming_ops'}
|
|
286
|
+
)
|
|
287
|
+
```
|
|
288
|
+
|
|
289
|
+
Get the nearest neighbors by Hamming distance
|
|
290
|
+
|
|
291
|
+
```python
|
|
292
|
+
order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
|
|
293
|
+
session.scalars(select(Item).order_by(order).limit(5))
|
|
294
|
+
```
|
|
295
|
+
|
|
296
|
+
Re-rank by the original vectors for better recall
|
|
297
|
+
|
|
298
|
+
```python
|
|
299
|
+
order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
|
|
300
|
+
subquery = session.query(Item).order_by(order).limit(20).subquery()
|
|
301
|
+
session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5))
|
|
302
|
+
```
|
|
303
|
+
|
|
287
304
|
#### Arrays
|
|
288
305
|
|
|
289
306
|
Add an array column
|
|
@@ -422,7 +439,7 @@ Enable the extension
|
|
|
422
439
|
conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
423
440
|
```
|
|
424
441
|
|
|
425
|
-
Register the
|
|
442
|
+
Register the types with your connection
|
|
426
443
|
|
|
427
444
|
```python
|
|
428
445
|
from pgvector.psycopg import register_vector
|
|
@@ -485,7 +502,7 @@ cur = conn.cursor()
|
|
|
485
502
|
cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
486
503
|
```
|
|
487
504
|
|
|
488
|
-
Register the
|
|
505
|
+
Register the types with your connection or cursor
|
|
489
506
|
|
|
490
507
|
```python
|
|
491
508
|
from pgvector.psycopg2 import register_vector
|
|
@@ -531,7 +548,7 @@ Enable the extension
|
|
|
531
548
|
await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
532
549
|
```
|
|
533
550
|
|
|
534
|
-
Register the
|
|
551
|
+
Register the types with your connection
|
|
535
552
|
|
|
536
553
|
```python
|
|
537
554
|
from pgvector.asyncpg import register_vector
|
|
@@ -585,7 +602,7 @@ Enable the extension
|
|
|
585
602
|
conn.run('CREATE EXTENSION IF NOT EXISTS vector')
|
|
586
603
|
```
|
|
587
604
|
|
|
588
|
-
Register the
|
|
605
|
+
Register the types with your connection
|
|
589
606
|
|
|
590
607
|
```python
|
|
591
608
|
from pgvector.pg8000 import register_vector
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
from django import VERSION
|
|
2
|
+
from django.contrib.postgres.operations import CreateExtension
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class VectorExtension(CreateExtension):
|
|
6
|
+
if VERSION[0] >= 6:
|
|
7
|
+
def __init__(self, hints=None):
|
|
8
|
+
super().__init__('vector', hints=hints)
|
|
9
|
+
else:
|
|
10
|
+
def __init__(self):
|
|
11
|
+
self.name = 'vector'
|
|
@@ -14,6 +14,18 @@ class BIT(UserDefinedType):
|
|
|
14
14
|
return 'BIT'
|
|
15
15
|
return 'BIT(%d)' % self.length
|
|
16
16
|
|
|
17
|
+
def bind_processor(self, dialect):
|
|
18
|
+
if dialect.__class__.__name__ == 'PGDialect_asyncpg':
|
|
19
|
+
import asyncpg
|
|
20
|
+
|
|
21
|
+
def process(value):
|
|
22
|
+
if isinstance(value, str):
|
|
23
|
+
return asyncpg.BitString(value)
|
|
24
|
+
return value
|
|
25
|
+
return process
|
|
26
|
+
else:
|
|
27
|
+
return super().bind_processor(dialect)
|
|
28
|
+
|
|
17
29
|
class comparator_factory(UserDefinedType.Comparator):
|
|
18
30
|
def hamming_distance(self, other):
|
|
19
31
|
return self.op('<~>', return_type=Float)(other)
|
|
@@ -1,3 +1,16 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pgvector
|
|
3
|
+
Version: 0.4.2
|
|
4
|
+
Summary: pgvector support for Python
|
|
5
|
+
Author-email: Andrew Kane <andrew@ankane.org>
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/pgvector/pgvector-python
|
|
8
|
+
Requires-Python: >=3.9
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
License-File: LICENSE.txt
|
|
11
|
+
Requires-Dist: numpy
|
|
12
|
+
Dynamic: license-file
|
|
13
|
+
|
|
1
14
|
# pgvector-python
|
|
2
15
|
|
|
3
16
|
[pgvector](https://github.com/pgvector/pgvector) support for Python
|
|
@@ -259,7 +272,6 @@ index = Index(
|
|
|
259
272
|
'my_index',
|
|
260
273
|
func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
|
|
261
274
|
postgresql_using='hnsw',
|
|
262
|
-
postgresql_with={'m': 16, 'ef_construction': 64},
|
|
263
275
|
postgresql_ops={'embedding': 'halfvec_l2_ops'}
|
|
264
276
|
)
|
|
265
277
|
```
|
|
@@ -271,6 +283,37 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])
|
|
|
271
283
|
session.scalars(select(Item).order_by(order).limit(5))
|
|
272
284
|
```
|
|
273
285
|
|
|
286
|
+
#### Binary Quantization
|
|
287
|
+
|
|
288
|
+
Use expression indexing for binary quantization
|
|
289
|
+
|
|
290
|
+
```python
|
|
291
|
+
from pgvector.sqlalchemy import BIT
|
|
292
|
+
from sqlalchemy.sql import func
|
|
293
|
+
|
|
294
|
+
index = Index(
|
|
295
|
+
'my_index',
|
|
296
|
+
func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
|
|
297
|
+
postgresql_using='hnsw',
|
|
298
|
+
postgresql_ops={'embedding': 'bit_hamming_ops'}
|
|
299
|
+
)
|
|
300
|
+
```
|
|
301
|
+
|
|
302
|
+
Get the nearest neighbors by Hamming distance
|
|
303
|
+
|
|
304
|
+
```python
|
|
305
|
+
order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
|
|
306
|
+
session.scalars(select(Item).order_by(order).limit(5))
|
|
307
|
+
```
|
|
308
|
+
|
|
309
|
+
Re-rank by the original vectors for better recall
|
|
310
|
+
|
|
311
|
+
```python
|
|
312
|
+
order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
|
|
313
|
+
subquery = session.query(Item).order_by(order).limit(20).subquery()
|
|
314
|
+
session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5))
|
|
315
|
+
```
|
|
316
|
+
|
|
274
317
|
#### Arrays
|
|
275
318
|
|
|
276
319
|
Add an array column
|
|
@@ -409,7 +452,7 @@ Enable the extension
|
|
|
409
452
|
conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
410
453
|
```
|
|
411
454
|
|
|
412
|
-
Register the
|
|
455
|
+
Register the types with your connection
|
|
413
456
|
|
|
414
457
|
```python
|
|
415
458
|
from pgvector.psycopg import register_vector
|
|
@@ -472,7 +515,7 @@ cur = conn.cursor()
|
|
|
472
515
|
cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
473
516
|
```
|
|
474
517
|
|
|
475
|
-
Register the
|
|
518
|
+
Register the types with your connection or cursor
|
|
476
519
|
|
|
477
520
|
```python
|
|
478
521
|
from pgvector.psycopg2 import register_vector
|
|
@@ -518,7 +561,7 @@ Enable the extension
|
|
|
518
561
|
await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
519
562
|
```
|
|
520
563
|
|
|
521
|
-
Register the
|
|
564
|
+
Register the types with your connection
|
|
522
565
|
|
|
523
566
|
```python
|
|
524
567
|
from pgvector.asyncpg import register_vector
|
|
@@ -572,7 +615,7 @@ Enable the extension
|
|
|
572
615
|
conn.run('CREATE EXTENSION IF NOT EXISTS vector')
|
|
573
616
|
```
|
|
574
617
|
|
|
575
|
-
Register the
|
|
618
|
+
Register the types with your connection
|
|
576
619
|
|
|
577
620
|
```python
|
|
578
621
|
from pgvector.pg8000 import register_vector
|
|
@@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "pgvector"
|
|
7
|
-
version = "0.4.
|
|
7
|
+
version = "0.4.2"
|
|
8
8
|
description = "pgvector support for Python"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
authors = [
|
|
11
11
|
{name = "Andrew Kane", email = "andrew@ankane.org"}
|
|
12
12
|
]
|
|
13
|
-
license =
|
|
13
|
+
license = "MIT"
|
|
14
14
|
requires-python = ">= 3.9"
|
|
15
15
|
dependencies = [
|
|
16
16
|
"numpy"
|
|
@@ -103,7 +103,6 @@ half_precision_index = Index(
|
|
|
103
103
|
'sqlalchemy_orm_half_precision_index',
|
|
104
104
|
func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
|
|
105
105
|
postgresql_using='hnsw',
|
|
106
|
-
postgresql_with={'m': 16, 'ef_construction': 64},
|
|
107
106
|
postgresql_ops={'embedding': 'halfvec_l2_ops'}
|
|
108
107
|
)
|
|
109
108
|
half_precision_index.create(setup_engine)
|
|
@@ -112,7 +111,6 @@ binary_quantize_index = Index(
|
|
|
112
111
|
'sqlalchemy_orm_binary_quantize_index',
|
|
113
112
|
func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
|
|
114
113
|
postgresql_using='hnsw',
|
|
115
|
-
postgresql_with={'m': 16, 'ef_construction': 64},
|
|
116
114
|
postgresql_ops={'embedding': 'bit_hamming_ops'}
|
|
117
115
|
)
|
|
118
116
|
binary_quantize_index.create(setup_engine)
|
|
@@ -528,6 +526,22 @@ class TestSqlalchemy:
|
|
|
528
526
|
items = session.query(Item).order_by(distance).all()
|
|
529
527
|
assert [v.id for v in items] == [2, 3, 1]
|
|
530
528
|
|
|
529
|
+
def test_binary_quantize_reranking(self, engine):
|
|
530
|
+
# recreate index (could also vacuum table)
|
|
531
|
+
binary_quantize_index.drop(setup_engine)
|
|
532
|
+
binary_quantize_index.create(setup_engine)
|
|
533
|
+
|
|
534
|
+
with Session(engine) as session:
|
|
535
|
+
session.add(Item(id=1, embedding=[-1, -2, -3]))
|
|
536
|
+
session.add(Item(id=2, embedding=[1, -2, 3]))
|
|
537
|
+
session.add(Item(id=3, embedding=[1, 2, 3]))
|
|
538
|
+
session.commit()
|
|
539
|
+
|
|
540
|
+
distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
|
|
541
|
+
subquery = session.query(Item).order_by(distance).limit(20).subquery()
|
|
542
|
+
items = session.query(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5).all()
|
|
543
|
+
assert [v.id for v in items] == [2, 3, 1]
|
|
544
|
+
|
|
531
545
|
|
|
532
546
|
@pytest.mark.parametrize('engine', array_engines)
|
|
533
547
|
class TestSqlalchemyArray:
|
|
@@ -596,6 +610,11 @@ class TestSqlalchemyAsync:
|
|
|
596
610
|
item = await session.get(Item, 1)
|
|
597
611
|
assert item.binary_embedding == embedding
|
|
598
612
|
|
|
613
|
+
if engine == asyncpg_engine:
|
|
614
|
+
session.add(Item(id=2, binary_embedding='101'))
|
|
615
|
+
item = await session.get(Item, 2)
|
|
616
|
+
assert item.binary_embedding == embedding
|
|
617
|
+
|
|
599
618
|
await engine.dispose()
|
|
600
619
|
|
|
601
620
|
@pytest.mark.asyncio
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|