Skip to content

Commit

Permalink
Modify Python integration test to store data files in ignored directory
Browse files Browse the repository at this point in the history
  • Loading branch information
Bryan Gillespie committed Jan 10, 2025
1 parent e3eb94c commit 9bce4fd
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 5 deletions.
1 change: 1 addition & 0 deletions iris-mpc-py/examples-py/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
data
16 changes: 11 additions & 5 deletions iris-mpc-py/examples-py/test_integration.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,11 @@
import os
from iris_mpc_py import PyIrisCode, PyPlaintextStore, PyGraphStore, PyHnswSearcher

vector_filename = "./data/vector.ndjson"
graph_filename = "./data/graph1.dat"
if not os.path.exists("./data/"):
os.makedirs("./data/")

print("Generating 100k uniform random iris codes...")
vector_init = PyPlaintextStore()
iris0 = PyIrisCode.uniform_random()
Expand All @@ -9,13 +15,13 @@

# write vector store to file
print("Writing vector store to file...")
vector_init.write_to_ndjson("vector.ndjson")
vector_init.write_to_ndjson(vector_filename)

print("Generating HNSW graphs for 10k imported iris codes...")
hnsw = PyHnswSearcher(32, 64, 32)
vector1 = PyPlaintextStore()
graph1 = PyGraphStore()
hnsw.fill_from_ndjson_file("vector.ndjson", vector1, graph1, 10000)
hnsw.fill_from_ndjson_file(vector_filename, vector1, graph1, 10000)

print("Imported length:", vector1.len())

Expand All @@ -27,11 +33,11 @@

# write graph store to file
print("Writing graph store to file...")
graph1.write_to_bin("graph1.dat")
graph1.write_to_bin(graph_filename)

# read HNSW graphs from disk
print("Reading vector and graph stores from file...")
vector2 = PyPlaintextStore.read_from_ndjson("vector.ndjson", 10000)
graph2 = PyGraphStore.read_from_bin("graph1.dat")
vector2 = PyPlaintextStore.read_from_ndjson(vector_filename, 10000)
graph2 = PyGraphStore.read_from_bin(graph_filename)

print("Search for random query iris code:", hnsw.search(query, vector2, graph2))

0 comments on commit 9bce4fd

Please sign in to comment.