Binary Quantization from Scratch

Setup: Install Dependencies, Imports & Download Embeddings

!pip install matplotlib tqdm pandas numpy datasets --quiet --upgrade

import numpy as np
import pandas as pd
from datasets import load_dataset
from tqdm import tqdm

👨🏾‍💻 Code Walkthrough

Here’s an explanation of the code structure provided:

1. Loading Data: OpenAI embeddings are loaded from a parquet files (we can load upto 1M embedding) and concatenated into one array.
2. Binary Conversion: A new array with the same shape is initialized with zeros, and the positive values in the original vectors are set to 1.
3. Accuracy Function: The accuracy function compares original vectors with binary vectors for a given index, limit, and oversampling rate. The comparison is done using dot products and logical XOR, sorting the results, and measuring the intersection.
4. Testing: The accuracy is tested for different oversampling rates (1, 2, 4), revealing a correctness of ~0.96 for an oversampling of 4.

💿 Loading Data

# Download from Huggingface Hub
ds = load_dataset(
    "Qdrant/dbpedia-entities-openai3-text-embedding-3-large-3072-100K", split="train"
)
openai_vectors = np.array(ds["text-embedding-3-large-3072-embedding"])
del ds

openai_bin = np.zeros_like(openai_vectors, dtype=np.int8)
openai_bin[openai_vectors > 0] = 1

n_dim = openai_vectors.shape[1]
n_dim

3072

🎯 Accuracy Function

We will use the accuracy function to compare the original vectors with the binary vectors for a given index, limit, and oversampling rate. The comparison is done using dot products and logical XOR, sorting the results, and measuring the intersection.

def accuracy(idx, limit: int, oversampling: int):
    scores = np.dot(openai_vectors, openai_vectors[idx])
    dot_results = np.argsort(scores)[-limit:][::-1]
    bin_scores = n_dim - np.logical_xor(openai_bin, openai_bin[idx]).sum(axis=1)
    bin_results = np.argsort(bin_scores)[-(limit * oversampling) :][::-1]
    return len(set(dot_results).intersection(set(bin_results))) / limit

📊 Results

number_of_samples = 10
limits = [3, 10]
sampling_rate = [1, 2, 3, 5]
results = []
def mean_accuracy(number_of_samples, limit, sampling_rate):
    return np.mean(
        [accuracy(i, limit=limit, oversampling=sampling_rate) for i in range(number_of_samples)]
    )
for i in tqdm(sampling_rate):
    for j in tqdm(limits):
        result = {
            "sampling_rate": i,
            "limit": j,
            "mean_acc": mean_accuracy(number_of_samples, j, i),
        }
        print(result)
        results.append(result)

  0%|          | 0/4 [00:00<?, ?it/s]
  0%|          | 0/2 [00:00<?, ?it/s]
 50%|█████     | 1/2 [00:02<00:02,  2.05s/it]

{'sampling_rate': 1, 'limit': 3, 'mean_acc': 0.9}

100%|██████████| 2/2 [00:04<00:00,  2.02s/it]
 25%|██▌       | 1/4 [00:04<00:12,  4.05s/it]

{'sampling_rate': 1, 'limit': 10, 'mean_acc': 0.8300000000000001}

  0%|          | 0/2 [00:00<?, ?it/s]
 50%|█████     | 1/2 [00:01<00:01,  1.72s/it]

{'sampling_rate': 2, 'limit': 3, 'mean_acc': 1.0}

100%|██████████| 2/2 [00:03<00:00,  1.76s/it]
 50%|█████     | 2/4 [00:07<00:07,  3.75s/it]

{'sampling_rate': 2, 'limit': 10, 'mean_acc': 0.9700000000000001}

  0%|          | 0/2 [00:00<?, ?it/s]
 50%|█████     | 1/2 [00:01<00:01,  1.72s/it]

{'sampling_rate': 3, 'limit': 3, 'mean_acc': 1.0}

100%|██████████| 2/2 [00:03<00:00,  1.69s/it]
 75%|███████▌  | 3/4 [00:10<00:03,  3.58s/it]

{'sampling_rate': 3, 'limit': 10, 'mean_acc': 0.9800000000000001}

  0%|          | 0/2 [00:00<?, ?it/s]
 50%|█████     | 1/2 [00:01<00:01,  1.68s/it]

{'sampling_rate': 5, 'limit': 3, 'mean_acc': 1.0}

100%|██████████| 2/2 [00:03<00:00,  1.65s/it]
100%|██████████| 4/4 [00:14<00:00,  3.57s/it]

{'sampling_rate': 5, 'limit': 10, 'mean_acc': 0.99}

㆓ Binary Conversion

Here, we will use 0 as the threshold for the binary conversion. All values greater than 0 will be set to 1, and others will remain 0. This is a simple and effective way to convert continuous values into binary values for OpenAI embeddings.

results = pd.DataFrame(results)
results