clusterkw
Version:
A package for clustering keywords using OpenAI embeddings
210 lines (209 loc) • 8.6 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.KeywordClusterer = void 0;
const openai_1 = __importDefault(require("openai"));
const simple_clustering_1 = require("./algorithms/simple-clustering");
const kmeans_clustering_1 = require("./algorithms/kmeans-clustering");
const hierarchical_clustering_1 = require("./algorithms/hierarchical-clustering");
const direct_clustering_1 = require("./algorithms/direct-clustering");
/**
* KeywordClusterer class for clustering keywords using OpenAI embeddings
*/
class KeywordClusterer {
/**
* Creates a new KeywordClusterer instance
* @param options Configuration options
*/
constructor(options) {
if (!options.apiKey) {
throw new Error('OpenAI API key is required');
}
this.openai = new openai_1.default({
apiKey: options.apiKey
});
this.embeddingModel = options.embeddingModel || 'text-embedding-3-small';
this.completionModel = options.completionModel || 'gpt-4o-mini-2024-07-18';
this.minClusterSize = options.minClusterSize || 2;
this.distanceThreshold = options.distanceThreshold || 0.3;
this.algorithm = options.algorithm || 'simple';
this.k = options.k;
this.maxIterations = options.maxIterations;
this.linkage = options.linkage || 'average';
this.context = options.context;
}
/**
* Clusters the provided keywords
* @param keywords Array of keywords to cluster
* @returns Array of clusters with names and descriptions
*/
async clusterKeywords(keywords) {
if (!keywords || keywords.length === 0) {
return [];
}
// For direct clustering, we don't need embeddings
if (this.algorithm === 'direct') {
return await (0, direct_clustering_1.directClustering)(keywords, this.openai, {
model: this.completionModel,
minClusterSize: this.minClusterSize,
maxClusters: this.k || 10,
context: this.context
});
}
// For other algorithms, proceed with embeddings
const embeddings = await this.getEmbeddings(keywords);
// Calculate distances between all embeddings
const distances = this.calculateDistances(embeddings);
// Generate clusters based on the selected algorithm
const clusters = this.generateClusters(keywords, distances, embeddings);
// Generate names and descriptions for clusters
const namedClusters = await this.nameAndDescribeClusters(clusters);
return namedClusters;
}
/**
* Gets embeddings for the provided texts
* @param texts Array of texts to get embeddings for
* @returns Array of embeddings
*/
async getEmbeddings(texts) {
const embeddings = [];
// Process in batches of 100 to avoid API limits
for (let i = 0; i < texts.length; i += 100) {
const batch = texts.slice(i, i + 100);
const response = await this.openai.embeddings.create({
model: this.embeddingModel,
input: batch
});
const batchEmbeddings = response.data.map(item => item.embedding);
embeddings.push(...batchEmbeddings);
}
return embeddings;
}
/**
* Calculates cosine distances between all embeddings
* @param embeddings Array of embeddings
* @returns Matrix of distances
*/
calculateDistances(embeddings) {
const n = embeddings.length;
const distances = Array(n).fill(0).map(() => Array(n).fill(0));
for (let i = 0; i < n; i++) {
for (let j = i; j < n; j++) {
if (i === j) {
distances[i][j] = 0;
}
else {
const distance = 1 - this.cosineSimilarity(embeddings[i], embeddings[j]);
distances[i][j] = distance;
distances[j][i] = distance; // Distance matrix is symmetric
}
}
}
return distances;
}
/**
* Calculates cosine similarity between two vectors
* @param a First vector
* @param b Second vector
* @returns Cosine similarity
*/
cosineSimilarity(a, b) {
if (a.length !== b.length) {
throw new Error('Vectors must have the same length');
}
let dotProduct = 0;
let normA = 0;
let normB = 0;
for (let i = 0; i < a.length; i++) {
dotProduct += a[i] * b[i];
normA += a[i] * a[i];
normB += b[i] * b[i];
}
if (normA === 0 || normB === 0) {
return 0;
}
return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
}
/**
* Generates clusters based on the selected algorithm
* @param keywords Array of keywords
* @param distances Matrix of distances
* @returns Array of clusters
*/
generateClusters(keywords, distances, embeddings) {
switch (this.algorithm) {
case 'kmeans':
if (!embeddings) {
throw new Error('Embeddings are required for k-means clustering');
}
return (0, kmeans_clustering_1.kmeansClustering)(keywords, embeddings, {
k: this.k,
maxIterations: this.maxIterations,
minClusterSize: this.minClusterSize
});
case 'hierarchical':
return (0, hierarchical_clustering_1.hierarchicalClustering)(keywords, distances, {
minClusterSize: this.minClusterSize,
distanceThreshold: this.distanceThreshold,
linkage: this.linkage
});
case 'simple':
default:
return (0, simple_clustering_1.simpleClustering)(keywords, distances, {
minClusterSize: this.minClusterSize,
distanceThreshold: this.distanceThreshold
});
}
}
/**
* Generates names and descriptions for clusters
* @param clusters Array of clusters
* @returns Array of clusters with names and descriptions
*/
async nameAndDescribeClusters(clusters) {
const namedClusters = [];
for (const cluster of clusters) {
const items = cluster.items.join(', ');
// Prepare the prompt with context if available
const contextPrompt = this.context
? `I have a cluster of keywords related to "${this.context}": ${items}. Please provide a short, descriptive name for this cluster and a brief description of what unifies these keywords in the context of ${this.context}. Format your response as JSON with "name" and "description" fields.`
: `I have a cluster of keywords: ${items}. Please provide a short, descriptive name for this cluster and a brief description of what unifies these keywords. Format your response as JSON with "name" and "description" fields.`;
const response = await this.openai.chat.completions.create({
model: this.completionModel,
messages: [
{
role: 'system',
content: 'You are a helpful assistant that names and describes clusters of keywords.'
},
{
role: 'user',
content: contextPrompt
}
],
response_format: { type: 'json_object' }
});
try {
const content = response.choices[0].message.content;
if (content) {
const result = JSON.parse(content);
namedClusters.push({
items: cluster.items,
name: result.name,
description: result.description
});
}
else {
namedClusters.push(cluster);
}
}
catch (error) {
// If parsing fails, just use the original cluster
namedClusters.push(cluster);
}
}
return namedClusters;
}
}
exports.KeywordClusterer = KeywordClusterer;