Add steps for reproducibility

build_ann.py

@@ -0,0 +1,39 @@
+from annoy import AnnoyIndex
+from safetensors import safe_open
+from tqdm import trange
+
+safetensors_path = "definitions.safetensors"
+
+with safe_open(safetensors_path, framework="numpy") as f:
+    vectors = f.get_tensor("vectors")
+
+num_vectors, vector_dim = vectors.shape
+
+print(f"Loaded {num_vectors} vectors of dimension {vector_dim}")
+
+index = AnnoyIndex(vector_dim, "angular")
+
+for i in trange(num_vectors):
+    index.add_item(i, vectors[i])
+
+num_trees = 25
+index.build(num_trees)
+
+index.save("definitions.ann")
+
+query_vector = vectors[0]
+num_neighbors = 5
+
+nearest_neighbors = index.get_nns_by_vector(query_vector, num_neighbors)
+print(f"Indices of {num_neighbors} nearest neighbors:", nearest_neighbors)
+
+neighbors_with_distances = index.get_nns_by_vector(
+    query_vector, num_neighbors, include_distances=True
+)
+print("Neighbors with distances:", neighbors_with_distances)
+
+import duckdb
+
+conn = duckdb.connect("sonajaht.db")
+query = "SELECT word_id, value FROM definitions WHERE entry_id in (SELECT unnest(?))"
+print(conn.execute(query, [nearest_neighbors]).df())

create_db.py

@@ -0,0 +1,24 @@
+import re
+
+import duckdb
+from datasets import load_dataset
+
+definitions_ds = load_dataset("adorkin/sonajaht", "definitions")
+words_ds = load_dataset("adorkin/sonajaht", "words")
+
+definitions = definitions_ds["definitions"].to_pandas()
+definitions.value = definitions.value.str.replace(
+    re.compile(r"<[^>]*>"), "", regex=True
+).apply(lambda el: " ".join(el.split()))
+
+definitions = duckdb.query(
+    "SELECT * FROM definitions WHERE lang = 'est' AND LENGTH(value) > 5"
+).df()
+
+definitions.reset_index(inplace=True, names="entry_id")
+
+words = words_ds["words"].to_pandas()
+
+conn = duckdb.connect("sonajaht.db")
+conn.execute("CREATE TABLE definitions AS SELECT * FROM definitions")
+conn.execute("CREATE TABLE words AS SELECT * FROM words")

vectorize.py

@@ -0,0 +1,34 @@
+import duckdb
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from safetensors.numpy import save_file
+from tqdm import tqdm
+
+conn = duckdb.connect("sonajaht.db")
+
+model = SentenceTransformer("sentence-transformers/LaBSE")
+
+
+query = "SELECT value FROM definitions"
+result = conn.execute(query)
+
+vectors = []
+batch_size = 64
+
+p_bar = tqdm()
+while True:
+    chunk = result.fetchmany(batch_size)
+    if not chunk:
+        break
+    values = [row[0] for row in chunk]
+    vectors.append(
+        model.encode(
+            values, show_progress_bar=False, batch_size=batch_size, device="mps"
+        )
+    )
+    p_bar.update(batch_size)
+
+vectors = np.concatenate(vectors)
+save_file(dict(vectors=vectors), "definitions.safetensors")
+
+conn.close()