-
Notifications
You must be signed in to change notification settings - Fork 0
/
partition.py
118 lines (93 loc) · 4.13 KB
/
partition.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import os
from sklearn.cluster import KMeans
import numpy as np
import json
import matplotlib.pyplot as plt
import seaborn as sns
import time
from itertools import groupby
desired_k_clusters = 500
# Get the directory path of the current script.
script_dir = os.path.dirname(os.path.abspath(__file__))
combined_embeddings_file = os.path.join(script_dir, "../data/embeddings/combined.json")
output_dir_partitions = os.path.join(script_dir, "../data/clusters")
if not os.path.exists(combined_embeddings_file):
print(
f"\t✋ Combined data file ({combined_embeddings_file}) deosn't exist. Run create-embeddings.py first."
)
exit(1)
print("Loading combined embeddings...")
with open(combined_embeddings_file, "r") as f:
combined_data = json.load(f)
# Partition the embeddings.
print(f"\n🪣 Performing k-means clustering, desired clusters {desired_k_clusters}:")
os.environ["TOKENIZERS_PARALLELISM"] = "true"
start_time = time.time()
# Convert the vectors from the combined data to a DataFrame.
vectors = np.array([entry["vector"] for entry in combined_data])
# Perform k-means quantization.
kmeans = KMeans(n_clusters=desired_k_clusters, random_state=0).fit(vectors)
# Predict the cluster assignments and get the centroid vector for each cluster.
cluster_assignments = kmeans.predict(vectors)
centroid_vectors = kmeans.cluster_centers_
# Update the combined data with cluster assignment.
for i in range(len(combined_data)):
combined_data[i]["cluster"] = str(cluster_assignments[i])
os.makedirs(output_dir_partitions, exist_ok=True)
# Store centroids.
# The format is compact: [ f64[], ... ]
# The array index is the centroid ID.
# f64[] is the centroid vector, corresponding to the arithmetic mean of embeddings in the cluster..
centroids = [vector.tolist() for _, vector in enumerate(centroid_vectors)]
with open(f"{output_dir_partitions}/centroids.json", "w") as f:
json.dump(centroids, f)
# Store combined data with cluster assignments.
# The format is compact: [ [ (cluster0) [embedding_id, f64[]], ... ], [ (cluster1) [], ... ], ...]
# The outermost array index is the cluster ID.
# The innermost arrays contain the embedding ID (corresponding to an id in the dataset)
# and the embedding vector.
# Sort the combined data by cluster.
sorted_objects = sorted(combined_data, key=lambda obj: int(obj["cluster"]))
# Group the sorted objects by cluster.
grouped_objects = groupby(sorted_objects, key=lambda obj: int(obj["cluster"]))
# Transform the grouped objects into the compact combined data format.
compact_combined_data = [
[[obj["id"], obj["vector"]] for obj in group] for _, group in grouped_objects
]
with open(
f"{output_dir_partitions}/combined.json",
"w",
) as f:
json.dump(compact_combined_data, f)
print(
f"\n\t🤙 Partitioned {len(combined_data)} datapoints in {desired_k_clusters} clusters: {time.time() - start_time:.2f} seconds"
)
# Glean the shape of the data.
# Plot the distribution of vectors per cluster.
unique, counts = np.unique(cluster_assignments, return_counts=True)
cluster_distribution = dict(zip(unique, counts))
min_count = min(cluster_distribution.values())
max_count = max(cluster_distribution.values())
print(f"\n\tPer cluster: {min_count} to {max_count} vectors")
# Calculate mean and standard deviation.
mean_count = np.mean(counts)
std_dev_count = np.std(counts)
# Define a threshold for "even" distribution.
threshold = 0.1 * mean_count
# Check if the distribution is even.
if all(abs(count - mean_count) <= threshold for count in counts):
print("\tDistribution is roughly even.")
else:
print(f"\tMean count: {mean_count:.2f}, standard deviation: {std_dev_count:.2f}")
# Pretty print the distribution as a histogram.
plot_file_name = f"{output_dir_partitions}/distribution-k{desired_k_clusters}.png"
plt.figure(figsize=(15, 6))
sns.barplot(x=list(cluster_distribution.keys()), y=list(cluster_distribution.values()))
plt.xticks([])
plt.xlabel(f"{desired_k_clusters} clusters")
plt.ylabel("Number of vectors")
plt.title("Distribution of vectors per cluster")
plt.savefig(plot_file_name, bbox_inches="tight")
plt.close()
print(f"\n\t📊 Cluster distribution plot: {plot_file_name}")
print("\tDone! 🎉")