small demo fixes

This commit is contained in:
Alex Garcia 2024-05-12 00:15:50 -07:00
parent f532692f7b
commit c446cc3963
3 changed files with 8 additions and 14 deletions

View file

@ -10,25 +10,16 @@ def serialize_f32(vector: List[float]) -> bytes:
return struct.pack("%sf" % len(vector), *vector)
def serialize_int8(vector: List[int]) -> bytes:
"""serializes a list of int8 into a compact "raw bytes" format"""
return struct.pack("%sb" % len(vector), *vector)
db = sqlite3.connect(":memory:")
db.enable_load_extension(True)
sqlite_vec.load(db)
db.enable_load_extension(False)
sqlite_version, vec_version = db.execute(
"select sqlite_version(), vec_version()"
).fetchone()
print(f"sqlite_version={sqlite_version}, vec_version={vec_version}")
db.execute("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
items = [
(1, [0.1, 0.1, 0.1, 0.1]),
(2, [0.2, 0.2, 0.2, 0.2]),
@ -36,6 +27,10 @@ items = [
(4, [0.4, 0.4, 0.4, 0.4]),
(5, [0.5, 0.5, 0.5, 0.5]),
]
query = [0.3, 0.3, 0.3, 0.3]
db.execute("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
with db:
for item in items:
db.execute(
@ -43,7 +38,6 @@ with db:
[item[0], serialize_f32(item[1])],
)
query = [0.3, 0.3, 0.3, 0.3]
rows = db.execute(
"""
SELECT