Function calculate_server_similarities

tests/test_multiple_datasets.py:166–252

(parquet_path: str, bf_csv_path: str)

def calculate_server_similarities(parquet_path: str, bf_csv_path: str):
    df = pl.read_parquet(parquet_path, low_memory=True)
    constants = {
        "COLLECTION_NAME": os.path.splitext(os.path.basename(parquet_path))[0],
        "DIMENSIONS": max(len(emb) for emb in df["emb"]),
        "MAX_VAL": 1.0,
        "MIN_VAL": -1.0,
        "BATCH_SIZE": 500,
        "BATCH_COUNT": None,
    }

    vectors = []
    df = df[["id", "emb"]]
    for v in df.iter_rows():
        vector = {"id": v[0], "values": v[1]}
        vectors.append(vector)
    constants["TOKEN"] = login()
    create_db(
        constants["COLLECTION_NAME"], "Embeddings from dataset", constants["DIMENSIONS"]
    )

    def upsert_with_retry(start_idx, retries=20):
        for attempt in range(retries):
            try:
                upsert_in_transaction(
                    constants["COLLECTION_NAME"],
                    txn_id,
                    vectors[start_idx : start_idx + batch_size],
                )
                return
            except Exception as e:
                print(
                    f"Upsert attempt {attempt + 1} failed for batch starting at index {start_idx}: {e}"
                )
                time.sleep(random.uniform(1, 3))
        print(
            f"Failed to upsert batch starting at index {start_idx} after {retries} attempts"
        )

    transaction = create_transaction(constants["COLLECTION_NAME"])
    txn_id = transaction["transaction_id"]

    batch_size = constants["BATCH_SIZE"]
    with ThreadPoolExecutor() as executor:
        futures = [
            executor.submit(upsert_with_retry, i)
            for i in range(0, len(vectors), batch_size)
        ]
        for future in futures:
            future.result()

    commit_transaction(constants["COLLECTION_NAME"], txn_id)
    print(colored("Vectors upserted successfully!", "green"))

    def search_dataset_vectors(query_id, query_emb):
        result = ann_vector(query_id, constants["COLLECTION_NAME"], query_emb)
        return result