Embeddings Quantization & Dequantization Cookbook

!pip install torch numpy mistralai

text = "Quantization reduces the precision of stored vector values."
embedding = get_text_embedding([text])

# Inspect what kind of floats the API actually returned.
# A Python `float` is a 64-bit (double) value;
print("type of one value:", type(embedding[0]).__name__)        # 'float'
print("value kind:", float(embedding[0]).__class__.__name__)    # 'float' (C double)

# We downcast to a float32 tensor for the rest of the cookbook
x = torch.tensor(embedding, dtype=torch.float32)

# Keep the dimension count — the bit-packed formats need it to trim padding bits
orig_dim = x.shape[0]
print("dims:", orig_dim, "| range:", x.min().item(), x.max().item())

q_int8 = quantize_int8(x)
print(q_int8)  # signed bytes in [-127, 127], same length as the embedding

q_uint8 = quantize_uint8(x)
print(q_uint8)  # unsigned bytes in [0, 255], same length as the embedding

q_int4, scale_int4 = quantize_int4(x)
print(q_int4)  # uint8 bytes, ceil(orig_dim/2) elements (packed nibbles, padded)

q_uint4, scale_uint4 = quantize_uint4(x)
print(q_uint4)  # uint8 bytes, ceil(orig_dim/2) elements (packed nibbles, padded)

q_bin = quantize_binary(x)
q_ubin = quantize_ubinary(x)
print(q_bin)   # signed bytes, ~orig_dim/8 elements (sign bits, padded)
print(q_ubin)  # unsigned bytes, ~orig_dim/8 elements (sign bits, padded)

# Signed 8-bit -> float
x_int8 = dequantize(q_int8, quant="int8")

# Unsigned 8-bit -> float
x_uint8 = dequantize(q_uint8, quant="uint8")

# Signed 4-bit -> float
x_int4 = dequantize(q_int4, quant="int4", orig_dim=orig_dim, scale=scale_int4)

# Unsigned 4-bit -> float
x_uint4 = dequantize(q_uint4, quant="uint4", orig_dim=orig_dim, scale=scale_uint4)

# Packed sign bits -> float (orig_dim trims the padding bits)
x_bin  = dequantize(q_bin,  quant="binary",  orig_dim=orig_dim)
x_ubin = dequantize(q_ubin, quant="ubinary", orig_dim=orig_dim)

print(x_int8.shape, x_int4.shape, x_bin.shape)  # all restore the original dimension count

import torch.nn.functional as F

def cosine(a, b):
    return F.cosine_similarity(a.unsqueeze(0).float(), b.unsqueeze(0).float()).item()

print("int8   :", cosine(x, dequantize(q_int8, "int8")))
print("uint8  :", cosine(x, dequantize(q_uint8, "uint8")))
print("int4   :", cosine(x, dequantize(q_int4, "int4", orig_dim, scale_int4)))
print("uint4  :", cosine(x, dequantize(q_uint4, "uint4", orig_dim, scale_uint4)))
print("binary :", cosine(x, dequantize(q_bin, "binary", orig_dim)))
print("ubinary:", cosine(x, dequantize(q_ubin, "ubinary", orig_dim)))

import pandas as pd

# Download the dataset
URL = "https://raw.githubusercontent.com/mistralai/cookbook/main/data/Symptom2Disease.csv"
df = pd.read_csv(URL, index_col=0)

# Remove full-row duplicates, keeping the first occurrence
before = len(df)
df = df.drop_duplicates()
after = len(df)

print(f"Rows before: {before}")
print(f"Rows after:  {after}")
print(f"Removed:     {before - after} duplicate rows")


def get_embeddings_by_chunks(data, chunk_size):
    chunks = [data[x : x + chunk_size] for x in range(0, len(data), chunk_size)]
    embeddings_response = [
        client.embeddings.create(model=model, inputs=c) for c in chunks
    ]
    return [d.embedding for e in embeddings_response for d in e.data]


df["embeddings"] = get_embeddings_by_chunks(df["text"].tolist(), 50)

# Original embedding dimension — the bit-packed formats need it to trim padding
orig_dim = len(df["embeddings"].iloc[0])
print("dims:", orig_dim, "| examples:", len(df))

# Create a train / test split
from sklearn.model_selection import train_test_split

train_x, test_x, train_y, test_y = train_test_split(
    df["embeddings"], df["label"], test_size=0.5, random_state=42
)
# Normalize features
from sklearn.preprocessing import StandardScaler

scaler = StandardScaler()
train_x = scaler.fit_transform(train_x.to_list())
test_x = scaler.transform(test_x.to_list())

# Train a classifier and compute the test accuracy
from sklearn.linear_model import LogisticRegression

# For a real problem, C should be properly cross validated and the confusion matrix analyzed
clf = LogisticRegression(random_state=0, C=1.0, max_iter=500).fit(
    train_x, train_y.to_list()
)

# you can also try the sag algorithm:
# clf = LogisticRegression(random_state=0, C=1.0, max_iter=1000, solver='sag').fit(train_x, train_y)

mistral_classifier_precision = 100 * np.mean(clf.predict(test_x) == test_y.to_list())

print(f"Precision: {mistral_classifier_precision:.2f}%")

# @title Train LogisticRegression models for each quantization format
formats = ["float", "int8", "uint8", "int4", "uint4", "binary", "ubinary"]
results = {}

for quant in formats:
    # Quantize -> dequantize every embedding back to floats to simulate loss of information
    X = quantize_dequantize(df["embeddings"].tolist(), quant, orig_dim)
    y = df["label"].tolist()

    # Same split as the baseline (same length + same random_state => same indices)
    train_x, test_x, train_y, test_y = train_test_split(
        X, y, test_size=0.5, random_state=42
    )

    # Same normalization
    scaler = StandardScaler()
    train_x = scaler.fit_transform(train_x)
    test_x = scaler.transform(test_x)

    # Same classifier & hyper-parameters
    clf = LogisticRegression(random_state=0, C=1.0, max_iter=500).fit(
        train_x, train_y
    )
    precision = 100 * np.mean(clf.predict(test_x) == test_y)
    results[quant] = precision
    print(f"{quant:8s}: {precision:.2f}%")

# @title Printing all Results
# size lookup
size_of = {}
for quant in formats:
    if quant == "float":
        size_of[quant] = "4 (float32)"
    elif quant in ("int8", "uint8"):
        size_of[quant] = "1"
    elif quant in ("int4", "uint4"):
        size_of[quant] = "0.5"
    else:  # binary, ubinary
        size_of[quant] = "1/8"

# Build rows
headers = ["quant", "bytes/value", "precision"]
rows = [[quant, size_of[quant], f"{results[quant]:.2f}%"] for quant in formats]

# Compute column widths
widths = [len(h) for h in headers]
for r in rows:
    for i, c in enumerate(r):
        widths[i] = max(widths[i], len(c))

# Print aligned table
print("  ".join(headers[i].ljust(widths[i]) for i in range(len(headers))))
print("  ".join("-" * widths[i] for i in range(len(headers))))
for r in rows:
    print("  ".join(r[i].ljust(widths[i]) for i in range(len(r))))