Text clustering groups similar documents without requiring a predefined list of categories. It is useful when topics, such as workplace email subjects, keep changing. section 4.4
The workflow turns documents into semantic embedding vectors, then applies a classical clustering algorithm to group nearby vectors. section 4.4.1
Figure 4.3 Clustering emails. We first calculate embedding vectors for all emails. Then we cluster those vectors to assign emails with similar content to the same cluster.
See the two-stage process: create email embeddings, then cluster them by content similarity. section 4.4.1
Listing 4.3 Clustering text documents using language models
import argparse
import openai
import pandas as pd
import time
from sklearn.cluster import KMeans
client = openai.OpenAI()
def get_embedding(text): #1
""" Calculate embedding vector for input text.
Args:
text: calculate embedding for this text.
Returns:
Vector representation of input text.
"""
for nr_retries in range(1, 4):
try:
response = client.embeddings.create(
model='text-embedding-ada-002',
input=text)
return response.data[0].embedding
except:
time.sleep(nr_retries * 2)
raise Exception('Cannot query OpenAI model!')
def get_kmeans(embeddings, k): #2
""" Cluster embedding vectors using K-means.
Args:
embeddings: embedding vectors.
k: number of result clusters.
Returns:
cluster IDs in embedding order.
"""
kmeans = KMeans(n_clusters=k, init='k-means++')
kmeans.fit(embeddings)
return kmeans.labels_
if __name__ == '__main__': #3
parser = argparse.ArgumentParser()
parser.add_argument('file_path', type=str, help='Path to input file')
parser.add_argument('nr_clusters', type=int, help='Number of clusters')
args = parser.parse_args()
df = pd.read_csv(args.file_path)
embeddings = df['text'].apply(get_embedding)
df['clusterid'] = get_kmeans(list(embeddings), args.nr_clusters)
df.to_csv('result.csv')
The Python implementation reads text from a CSV, creates embeddings, clusters them with K-means, and writes cluster IDs to the results. section 4.4.4
Separate poems from emails
Run the clustering script on a mixed dataset with two target clusters. The output adds a cluster ID to each row; the expected result places emails and poems in separate groups. section 4.4.5
python listing3.py textmix.csv 2
Example command for clustering the mixed text dataset into two groups. section 4.4.5
