update source code and pipeline
This commit is contained in:
649
cluster/gmm_extensive.py
Normal file
649
cluster/gmm_extensive.py
Normal file
@@ -0,0 +1,649 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Extensive Gaussian Mixture Model clustering with grid search for optimal parameters
|
||||
Includes BIC and AIC metrics for model selection
|
||||
"""
|
||||
|
||||
import json
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from sklearn.mixture import GaussianMixture
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
from sklearn.metrics import silhouette_score, calinski_harabasz_score, davies_bouldin_score
|
||||
from sklearn.decomposition import PCA
|
||||
import datetime
|
||||
import csv
|
||||
import argparse
|
||||
import warnings
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
class GMMExtensiveClustering:
|
||||
def __init__(self, embeddings_path):
|
||||
self.embeddings_path = embeddings_path
|
||||
self.embeddings = None
|
||||
self.file_paths = None
|
||||
self.load_embeddings()
|
||||
|
||||
def load_embeddings(self):
|
||||
"""Load embeddings from JSON file"""
|
||||
print(f"Loading embeddings from {self.embeddings_path}...")
|
||||
with open(self.embeddings_path, 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
self.file_paths = []
|
||||
embeddings_list = []
|
||||
|
||||
for item in data:
|
||||
self.file_paths.append(item['filepath'])
|
||||
embeddings_list.append(item['embedding'])
|
||||
|
||||
self.embeddings = np.array(embeddings_list, dtype=np.float32)
|
||||
print(f"Loaded {len(self.file_paths)} samples with embedding dimension {self.embeddings.shape[1]}")
|
||||
|
||||
# Standardize embeddings for better clustering
|
||||
self.scaler = StandardScaler()
|
||||
self.embeddings_scaled = self.scaler.fit_transform(self.embeddings)
|
||||
|
||||
def run_gmm_grid_search(self):
|
||||
"""Run GMM with optimized grid search for faster execution"""
|
||||
print("\n" + "="*70)
|
||||
print("RUNNING GAUSSIAN MIXTURE MODEL CLUSTERING WITH OPTIMIZED GRID SEARCH")
|
||||
print("="*70)
|
||||
|
||||
# Optimized GMM parameter candidates for faster execution
|
||||
|
||||
# Smart n_components range with larger steps
|
||||
max_components = min(50, len(self.embeddings_scaled) // 20) # Reduced max and increased divisor
|
||||
n_components_candidates = []
|
||||
|
||||
# Progressive step sizes: smaller steps for low numbers, larger for high
|
||||
for n in range(2, max_components + 1):
|
||||
if n <= 5:
|
||||
n_components_candidates.append(n) # 2, 3, 4, 5
|
||||
elif n <= 10:
|
||||
if n % 2 == 0: # 6, 8, 10
|
||||
n_components_candidates.append(n)
|
||||
else:
|
||||
if n % 3 == 2: # 11, 14, 17, 20
|
||||
n_components_candidates.append(n)
|
||||
|
||||
# Reduced covariance types - focus on most important ones
|
||||
covariance_types = [
|
||||
# 'full', 'diag',
|
||||
'tied', 'spherical'
|
||||
] # Removed 'tied' and 'spherical' as they're less common
|
||||
|
||||
# Simplified regularization - focus on key values
|
||||
reg_covar_candidates = [1e-5, 1e-4, 1e-3] # Removed extreme values
|
||||
|
||||
# Reduced n_init - 1 is often sufficient for good initialization methods
|
||||
n_init_candidates = [1, 5] # Removed 10 to save time
|
||||
|
||||
# Focus on best initialization methods
|
||||
init_params_candidates = ['kmeans', 'k-means++'] # Removed 'random' and 'random_from_data'
|
||||
|
||||
# Simplified max_iter - most problems converge quickly
|
||||
max_iter_candidates = [100, 300] # Removed 500, added 300 as middle ground
|
||||
|
||||
print(f"Optimized parameter combinations:")
|
||||
print(f" - n_components: {len(n_components_candidates)} values {n_components_candidates}")
|
||||
print(f" - covariance_types: {len(covariance_types)} options {covariance_types}")
|
||||
print(f" - reg_covar: {len(reg_covar_candidates)} values {reg_covar_candidates}")
|
||||
print(f" - n_init: {len(n_init_candidates)} values {n_init_candidates}")
|
||||
print(f" - init_params: {len(init_params_candidates)} options {init_params_candidates}")
|
||||
print(f" - max_iter: {len(max_iter_candidates)} values {max_iter_candidates}")
|
||||
|
||||
total_combinations = (len(n_components_candidates) * len(covariance_types) *
|
||||
len(reg_covar_candidates) * len(n_init_candidates) *
|
||||
len(init_params_candidates) * len(max_iter_candidates))
|
||||
print(f"Total combinations: {total_combinations} (optimized for speed)")
|
||||
|
||||
# Estimate time
|
||||
estimated_time_per_combination = 0.5 # seconds (conservative estimate)
|
||||
estimated_total_time = total_combinations * estimated_time_per_combination
|
||||
print(f"Estimated runtime: {estimated_total_time/60:.1f} minutes")
|
||||
print("This should be much faster...\n")
|
||||
|
||||
# Track all results for analysis
|
||||
all_results = []
|
||||
|
||||
# Early stopping criteria for speed optimization
|
||||
early_stopping_threshold = 0.7 # If we find a very good silhouette score, we can be less exhaustive
|
||||
good_results_found = 0
|
||||
max_good_results = 5 # Stop early if we find several very good results
|
||||
|
||||
best_bic_score = float('inf')
|
||||
best_aic_score = float('inf')
|
||||
best_silhouette_score = -1
|
||||
best_params_bic = None
|
||||
best_params_aic = None
|
||||
best_params_silhouette = None
|
||||
best_labels_bic = None
|
||||
best_labels_aic = None
|
||||
best_labels_silhouette = None
|
||||
|
||||
current_combination = 0
|
||||
|
||||
# Optimized iteration order: test simpler models first (fewer components, simpler covariance)
|
||||
for covariance_type in covariance_types: # Start with covariance type
|
||||
for n_components in n_components_candidates: # Then components
|
||||
for init_params in init_params_candidates: # Good initialization methods
|
||||
for reg_covar in reg_covar_candidates: # Regularization
|
||||
for n_init in n_init_candidates: # Number of initializations
|
||||
for max_iter in max_iter_candidates: # Iterations last
|
||||
current_combination += 1
|
||||
|
||||
# Progress indicator with time estimation
|
||||
if current_combination % 50 == 0 or current_combination == total_combinations:
|
||||
progress = (current_combination / total_combinations) * 100
|
||||
print(f"Progress: {current_combination}/{total_combinations} ({progress:.1f}%) - "
|
||||
f"Best scores so far: BIC={best_bic_score:.2f}, Silhouette={best_silhouette_score:.3f}")
|
||||
|
||||
try:
|
||||
# Early convergence check for faster models
|
||||
tol = 1e-3 if n_components <= 5 else 1e-4 # Less strict tolerance for simple models
|
||||
|
||||
# Run GMM
|
||||
gmm = GaussianMixture(
|
||||
n_components=n_components,
|
||||
covariance_type=covariance_type,
|
||||
reg_covar=reg_covar,
|
||||
n_init=n_init,
|
||||
init_params=init_params,
|
||||
max_iter=max_iter,
|
||||
tol=tol, # Added tolerance for faster convergence
|
||||
random_state=42
|
||||
)
|
||||
|
||||
# Fit and predict
|
||||
gmm.fit(self.embeddings_scaled)
|
||||
labels = gmm.predict(self.embeddings_scaled)
|
||||
|
||||
# Quick validation - skip if model didn't converge properly
|
||||
if not gmm.converged_ and max_iter <= 100:
|
||||
continue # Skip non-converged simple models
|
||||
|
||||
# Calculate metrics
|
||||
bic_score = gmm.bic(self.embeddings_scaled)
|
||||
aic_score = gmm.aic(self.embeddings_scaled)
|
||||
log_likelihood = gmm.score(self.embeddings_scaled)
|
||||
|
||||
# Only calculate clustering metrics if we have multiple clusters
|
||||
if len(set(labels)) > 1:
|
||||
silhouette = silhouette_score(self.embeddings_scaled, labels)
|
||||
calinski_harabasz = calinski_harabasz_score(self.embeddings_scaled, labels)
|
||||
davies_bouldin = davies_bouldin_score(self.embeddings_scaled, labels)
|
||||
|
||||
# Early stopping check
|
||||
if silhouette > early_stopping_threshold:
|
||||
good_results_found += 1
|
||||
print(f"🎯 Excellent result found: n_comp={n_components}, cov={covariance_type}, "
|
||||
f"silhouette={silhouette:.4f}")
|
||||
|
||||
else:
|
||||
silhouette = -1
|
||||
calinski_harabasz = 0
|
||||
davies_bouldin = float('inf')
|
||||
|
||||
# Store result for analysis
|
||||
result_info = {
|
||||
'n_components': n_components,
|
||||
'covariance_type': covariance_type,
|
||||
'reg_covar': reg_covar,
|
||||
'n_init': n_init,
|
||||
'init_params': init_params,
|
||||
'max_iter': max_iter,
|
||||
'bic_score': bic_score,
|
||||
'aic_score': aic_score,
|
||||
'log_likelihood': log_likelihood,
|
||||
'silhouette_score': silhouette,
|
||||
'calinski_harabasz_score': calinski_harabasz,
|
||||
'davies_bouldin_score': davies_bouldin,
|
||||
'converged': gmm.converged_,
|
||||
'n_iter': gmm.n_iter_,
|
||||
'unique_clusters': len(set(labels))
|
||||
}
|
||||
|
||||
all_results.append(result_info)
|
||||
|
||||
# Print promising results
|
||||
if (silhouette > 0.3 and bic_score < np.percentile([r['bic_score'] for r in all_results], 25)):
|
||||
print(f"n_components={n_components}, cov={covariance_type}, init={init_params}: "
|
||||
f"BIC={bic_score:.2f}, AIC={aic_score:.2f}, silhouette={silhouette:.4f}")
|
||||
|
||||
# Track best results for different criteria
|
||||
if bic_score < best_bic_score:
|
||||
best_bic_score = bic_score
|
||||
best_params_bic = {
|
||||
'n_components': n_components,
|
||||
'covariance_type': covariance_type,
|
||||
'reg_covar': reg_covar,
|
||||
'n_init': n_init,
|
||||
'init_params': init_params,
|
||||
'max_iter': max_iter
|
||||
}
|
||||
best_labels_bic = labels
|
||||
|
||||
if aic_score < best_aic_score:
|
||||
best_aic_score = aic_score
|
||||
best_params_aic = {
|
||||
'n_components': n_components,
|
||||
'covariance_type': covariance_type,
|
||||
'reg_covar': reg_covar,
|
||||
'n_init': n_init,
|
||||
'init_params': init_params,
|
||||
'max_iter': max_iter
|
||||
}
|
||||
best_labels_aic = labels
|
||||
|
||||
if silhouette > best_silhouette_score and len(set(labels)) > 1:
|
||||
best_silhouette_score = silhouette
|
||||
best_params_silhouette = {
|
||||
'n_components': n_components,
|
||||
'covariance_type': covariance_type,
|
||||
'reg_covar': reg_covar,
|
||||
'n_init': n_init,
|
||||
'init_params': init_params,
|
||||
'max_iter': max_iter
|
||||
}
|
||||
best_labels_silhouette = labels
|
||||
|
||||
# Early stopping check
|
||||
if good_results_found >= 5 and silhouette > 0.6:
|
||||
print(f"🛑 Early stopping triggered: Found {good_results_found} excellent results. "
|
||||
f"Stopping at {current_combination}/{total_combinations} combinations.")
|
||||
break
|
||||
|
||||
except Exception:
|
||||
# Skip problematic parameter combinations
|
||||
continue
|
||||
|
||||
# Break from nested loops if early stopping triggered
|
||||
if good_results_found >= 5 and best_silhouette_score > 0.6:
|
||||
break
|
||||
if good_results_found >= 5 and best_silhouette_score > 0.6:
|
||||
break
|
||||
if good_results_found >= 5 and best_silhouette_score > 0.6:
|
||||
break
|
||||
if good_results_found >= 5 and best_silhouette_score > 0.6:
|
||||
break
|
||||
if good_results_found >= 5 and best_silhouette_score > 0.6:
|
||||
break
|
||||
|
||||
# Analysis of results
|
||||
print("\n" + "="*70)
|
||||
print("GAUSSIAN MIXTURE MODEL GRID SEARCH ANALYSIS")
|
||||
print("="*70)
|
||||
|
||||
if all_results:
|
||||
import pandas as pd
|
||||
df_results = pd.DataFrame(all_results)
|
||||
|
||||
print(f"Total parameter combinations tested: {len(df_results)}")
|
||||
|
||||
# Filter results with valid clustering (more than 1 cluster)
|
||||
valid_results = df_results[df_results['unique_clusters'] > 1]
|
||||
print(f"Combinations with valid clustering: {len(valid_results)}")
|
||||
|
||||
if len(valid_results) > 0:
|
||||
# Best scores analysis
|
||||
print(f"\nModel Selection Metrics:")
|
||||
print(f"Best BIC score: {df_results['bic_score'].min():.2f}")
|
||||
print(f"Best AIC score: {df_results['aic_score'].min():.2f}")
|
||||
print(f"Best Log-Likelihood: {df_results['log_likelihood'].max():.2f}")
|
||||
|
||||
print(f"\nClustering Quality Metrics:")
|
||||
print(f"Best silhouette score: {valid_results['silhouette_score'].max():.4f}")
|
||||
print(f"Mean silhouette score: {valid_results['silhouette_score'].mean():.4f}")
|
||||
print(f"Best Calinski-Harabasz score: {valid_results['calinski_harabasz_score'].max():.2f}")
|
||||
print(f"Best Davies-Bouldin score: {valid_results['davies_bouldin_score'].min():.4f}")
|
||||
|
||||
# Top results by different criteria
|
||||
print(f"\nTop 5 results by BIC (lower is better):")
|
||||
top_bic = df_results.nsmallest(5, 'bic_score')
|
||||
for idx, row in top_bic.iterrows():
|
||||
print(f" n_comp={row['n_components']}, cov={row['covariance_type']}: "
|
||||
f"BIC={row['bic_score']:.2f}, AIC={row['aic_score']:.2f}")
|
||||
|
||||
print(f"\nTop 5 results by AIC (lower is better):")
|
||||
top_aic = df_results.nsmallest(5, 'aic_score')
|
||||
for idx, row in top_aic.iterrows():
|
||||
print(f" n_comp={row['n_components']}, cov={row['covariance_type']}: "
|
||||
f"BIC={row['bic_score']:.2f}, AIC={row['aic_score']:.2f}")
|
||||
|
||||
if len(valid_results) > 0:
|
||||
print(f"\nTop 5 results by Silhouette Score:")
|
||||
top_silhouette = valid_results.nlargest(5, 'silhouette_score')
|
||||
for idx, row in top_silhouette.iterrows():
|
||||
print(f" n_comp={row['n_components']}, cov={row['covariance_type']}: "
|
||||
f"silhouette={row['silhouette_score']:.4f}")
|
||||
|
||||
# Component count analysis
|
||||
component_performance = df_results.groupby('n_components').agg({
|
||||
'bic_score': 'min',
|
||||
'aic_score': 'min',
|
||||
'silhouette_score': 'max'
|
||||
}).reset_index()
|
||||
|
||||
print(f"\nComponent count analysis (top 10 by BIC):")
|
||||
top_components = component_performance.nsmallest(10, 'bic_score')
|
||||
for idx, row in top_components.iterrows():
|
||||
print(f" {row['n_components']} components: "
|
||||
f"BIC={row['bic_score']:.2f}, AIC={row['aic_score']:.2f}, "
|
||||
f"silhouette={row['silhouette_score']:.4f}")
|
||||
|
||||
print(f"\n📁 SAVING DETAILED RESULTS...")
|
||||
print("="*30)
|
||||
|
||||
# Save detailed grid search results
|
||||
self.save_gmm_grid_search_results(all_results,
|
||||
best_params_bic, best_bic_score,
|
||||
best_params_aic, best_aic_score,
|
||||
best_params_silhouette, best_silhouette_score)
|
||||
|
||||
# Return best results based on BIC (primary), AIC (secondary), Silhouette (tertiary)
|
||||
results = {
|
||||
'bic': (best_labels_bic, best_params_bic, best_bic_score),
|
||||
'aic': (best_labels_aic, best_params_aic, best_aic_score),
|
||||
'silhouette': (best_labels_silhouette, best_params_silhouette, best_silhouette_score)
|
||||
}
|
||||
|
||||
# Print best results
|
||||
if best_labels_bic is not None:
|
||||
print(f"\nBest GMM result by BIC:")
|
||||
print(f"Parameters: {best_params_bic}")
|
||||
print(f"BIC score: {best_bic_score:.2f}")
|
||||
|
||||
if best_labels_aic is not None:
|
||||
print(f"\nBest GMM result by AIC:")
|
||||
print(f"Parameters: {best_params_aic}")
|
||||
print(f"AIC score: {best_aic_score:.2f}")
|
||||
|
||||
if best_labels_silhouette is not None:
|
||||
print(f"\nBest GMM result by Silhouette:")
|
||||
print(f"Parameters: {best_params_silhouette}")
|
||||
print(f"Silhouette score: {best_silhouette_score:.4f}")
|
||||
|
||||
return results
|
||||
|
||||
def save_gmm_grid_search_results(self, all_results,
|
||||
best_params_bic, best_bic_score,
|
||||
best_params_aic, best_aic_score,
|
||||
best_params_silhouette, best_silhouette_score):
|
||||
"""Save detailed GMM grid search results to JSON file"""
|
||||
|
||||
# Prepare comprehensive results data
|
||||
grid_search_data = {
|
||||
"experiment_info": {
|
||||
"timestamp": datetime.datetime.now().isoformat(),
|
||||
"dataset_path": self.embeddings_path,
|
||||
"total_samples": len(self.file_paths),
|
||||
"embedding_dimension": self.embeddings.shape[1],
|
||||
"total_combinations_tested": len(all_results),
|
||||
"method": "Gaussian Mixture Model"
|
||||
},
|
||||
"best_results": {
|
||||
"by_bic": {
|
||||
"parameters": best_params_bic,
|
||||
"bic_score": best_bic_score if best_bic_score != float('inf') else None
|
||||
},
|
||||
"by_aic": {
|
||||
"parameters": best_params_aic,
|
||||
"aic_score": best_aic_score if best_aic_score != float('inf') else None
|
||||
},
|
||||
"by_silhouette": {
|
||||
"parameters": best_params_silhouette,
|
||||
"silhouette_score": best_silhouette_score if best_silhouette_score > -1 else None
|
||||
}
|
||||
},
|
||||
"all_trials": []
|
||||
}
|
||||
|
||||
# Add all trial results
|
||||
for i, result in enumerate(all_results):
|
||||
trial_data = {
|
||||
"trial_id": i + 1,
|
||||
"parameters": {
|
||||
"n_components": result['n_components'],
|
||||
"covariance_type": result['covariance_type'],
|
||||
"reg_covar": result['reg_covar'],
|
||||
"n_init": result['n_init'],
|
||||
"init_params": result['init_params'],
|
||||
"max_iter": result['max_iter']
|
||||
},
|
||||
"results": {
|
||||
"bic_score": result['bic_score'],
|
||||
"aic_score": result['aic_score'],
|
||||
"log_likelihood": result['log_likelihood'],
|
||||
"silhouette_score": result['silhouette_score'],
|
||||
"calinski_harabasz_score": result['calinski_harabasz_score'],
|
||||
"davies_bouldin_score": result['davies_bouldin_score'],
|
||||
"converged": result['converged'],
|
||||
"n_iter": result['n_iter'],
|
||||
"unique_clusters": result['unique_clusters']
|
||||
}
|
||||
}
|
||||
grid_search_data["all_trials"].append(trial_data)
|
||||
|
||||
# Calculate summary statistics
|
||||
if all_results:
|
||||
bic_scores = [r['bic_score'] for r in all_results]
|
||||
aic_scores = [r['aic_score'] for r in all_results]
|
||||
log_likelihoods = [r['log_likelihood'] for r in all_results]
|
||||
|
||||
valid_silhouette = [r['silhouette_score'] for r in all_results if r['silhouette_score'] > -1]
|
||||
|
||||
grid_search_data["summary_statistics"] = {
|
||||
"total_trials": len(all_results),
|
||||
"valid_clustering_trials": len(valid_silhouette),
|
||||
"bic_score": {
|
||||
"best": min(bic_scores),
|
||||
"worst": max(bic_scores),
|
||||
"mean": sum(bic_scores) / len(bic_scores),
|
||||
"median": sorted(bic_scores)[len(bic_scores)//2]
|
||||
},
|
||||
"aic_score": {
|
||||
"best": min(aic_scores),
|
||||
"worst": max(aic_scores),
|
||||
"mean": sum(aic_scores) / len(aic_scores),
|
||||
"median": sorted(aic_scores)[len(aic_scores)//2]
|
||||
},
|
||||
"log_likelihood": {
|
||||
"best": max(log_likelihoods),
|
||||
"worst": min(log_likelihoods),
|
||||
"mean": sum(log_likelihoods) / len(log_likelihoods)
|
||||
}
|
||||
}
|
||||
|
||||
if valid_silhouette:
|
||||
grid_search_data["summary_statistics"]["silhouette_score"] = {
|
||||
"best": max(valid_silhouette),
|
||||
"worst": min(valid_silhouette),
|
||||
"mean": sum(valid_silhouette) / len(valid_silhouette),
|
||||
"median": sorted(valid_silhouette)[len(valid_silhouette)//2]
|
||||
}
|
||||
|
||||
# Top 10 results by different criteria
|
||||
sorted_by_bic = sorted(all_results, key=lambda x: x['bic_score'])
|
||||
sorted_by_aic = sorted(all_results, key=lambda x: x['aic_score'])
|
||||
valid_results = [r for r in all_results if r['silhouette_score'] > -1]
|
||||
sorted_by_silhouette = sorted(valid_results, key=lambda x: x['silhouette_score'], reverse=True)
|
||||
|
||||
grid_search_data["top_10_results"] = {
|
||||
"by_bic": [],
|
||||
"by_aic": [],
|
||||
"by_silhouette": []
|
||||
}
|
||||
|
||||
for i, result in enumerate(sorted_by_bic[:10]):
|
||||
grid_search_data["top_10_results"]["by_bic"].append({
|
||||
"rank": i + 1,
|
||||
"parameters": {
|
||||
"n_components": result['n_components'],
|
||||
"covariance_type": result['covariance_type'],
|
||||
"init_params": result['init_params']
|
||||
},
|
||||
"bic_score": result['bic_score'],
|
||||
"aic_score": result['aic_score']
|
||||
})
|
||||
|
||||
for i, result in enumerate(sorted_by_aic[:10]):
|
||||
grid_search_data["top_10_results"]["by_aic"].append({
|
||||
"rank": i + 1,
|
||||
"parameters": {
|
||||
"n_components": result['n_components'],
|
||||
"covariance_type": result['covariance_type'],
|
||||
"init_params": result['init_params']
|
||||
},
|
||||
"bic_score": result['bic_score'],
|
||||
"aic_score": result['aic_score']
|
||||
})
|
||||
|
||||
for i, result in enumerate(sorted_by_silhouette[:10]):
|
||||
grid_search_data["top_10_results"]["by_silhouette"].append({
|
||||
"rank": i + 1,
|
||||
"parameters": {
|
||||
"n_components": result['n_components'],
|
||||
"covariance_type": result['covariance_type'],
|
||||
"init_params": result['init_params']
|
||||
},
|
||||
"silhouette_score": result['silhouette_score']
|
||||
})
|
||||
|
||||
# Save to file with timestamp
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"gmm_grid_search_detailed_{timestamp}.json"
|
||||
# print()
|
||||
|
||||
# with open(filename, 'w') as f:
|
||||
# json.dump(grid_search_data, f, indent=4, ensure_ascii=False)
|
||||
|
||||
print(f"Detailed grid search results saved to: {filename}")
|
||||
|
||||
# Also save a CSV summary for easy analysis
|
||||
csv_filename = f"gmm_grid_search_summary_{timestamp}.csv"
|
||||
self.save_grid_search_csv(all_results, csv_filename)
|
||||
print(f"Grid search summary CSV saved to: {csv_filename}")
|
||||
|
||||
def save_grid_search_csv(self, all_results, filename):
|
||||
"""Save grid search results as CSV for easy analysis"""
|
||||
|
||||
with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
|
||||
fieldnames = ['trial_id', 'n_components', 'covariance_type', 'reg_covar',
|
||||
'n_init', 'init_params', 'max_iter', 'bic_score', 'aic_score',
|
||||
'log_likelihood', 'silhouette_score', 'calinski_harabasz_score',
|
||||
'davies_bouldin_score', 'converged', 'n_iter', 'unique_clusters']
|
||||
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
|
||||
|
||||
writer.writeheader()
|
||||
for i, result in enumerate(all_results):
|
||||
writer.writerow({
|
||||
'trial_id': i + 1,
|
||||
'n_components': result['n_components'],
|
||||
'covariance_type': result['covariance_type'],
|
||||
'reg_covar': result['reg_covar'],
|
||||
'n_init': result['n_init'],
|
||||
'init_params': result['init_params'],
|
||||
'max_iter': result['max_iter'],
|
||||
'bic_score': result['bic_score'],
|
||||
'aic_score': result['aic_score'],
|
||||
'log_likelihood': result['log_likelihood'],
|
||||
'silhouette_score': result['silhouette_score'],
|
||||
'calinski_harabasz_score': result['calinski_harabasz_score'],
|
||||
'davies_bouldin_score': result['davies_bouldin_score'],
|
||||
'converged': result['converged'],
|
||||
'n_iter': result['n_iter'],
|
||||
'unique_clusters': result['unique_clusters']
|
||||
})
|
||||
|
||||
def visualize_results(self, results):
|
||||
"""Visualize clustering results using PCA"""
|
||||
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
|
||||
|
||||
# Reduce dimensions for visualization
|
||||
pca = PCA(n_components=2, random_state=42)
|
||||
embeddings_2d = pca.fit_transform(self.embeddings_scaled)
|
||||
|
||||
methods = ['bic', 'aic', 'silhouette']
|
||||
titles = ['Best by BIC', 'Best by AIC', 'Best by Silhouette']
|
||||
|
||||
for idx, (method, title) in enumerate(zip(methods, titles)):
|
||||
labels, params, score = results[method]
|
||||
|
||||
if labels is not None:
|
||||
unique_labels = set(labels)
|
||||
colors = plt.cm.Set3(np.linspace(0, 1, len(unique_labels)))
|
||||
|
||||
for label, color in zip(unique_labels, colors):
|
||||
mask = labels == label
|
||||
axes[idx].scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1],
|
||||
c=[color], s=50, alpha=0.7, label=f'Cluster {label}')
|
||||
|
||||
axes[idx].set_title(f'{title}\nn_components={params["n_components"]}, '
|
||||
f'cov={params["covariance_type"]}')
|
||||
else:
|
||||
axes[idx].text(0.5, 0.5, 'No valid clustering', ha='center', va='center',
|
||||
transform=axes[idx].transAxes, fontsize=12)
|
||||
axes[idx].set_title(f'{title}\n(Failed)')
|
||||
|
||||
axes[idx].set_xlabel('PCA Component 1')
|
||||
axes[idx].set_ylabel('PCA Component 2')
|
||||
axes[idx].grid(True, alpha=0.3)
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig('gmm_clustering_results.png', dpi=300, bbox_inches='tight')
|
||||
plt.show()
|
||||
|
||||
print(f"Visualization saved as 'gmm_clustering_results.png'")
|
||||
|
||||
def save_clustering_results(self, results):
|
||||
"""Save final clustering results to JSON files"""
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
for method in ['bic', 'aic', 'silhouette']:
|
||||
labels, params, score = results[method]
|
||||
|
||||
if labels is not None:
|
||||
clustering_results = []
|
||||
for filepath, label in zip(self.file_paths, labels):
|
||||
clustering_results.append({
|
||||
"filepath": filepath,
|
||||
"cluster": int(label)
|
||||
})
|
||||
|
||||
filename = f"gmm_final_results_{method}_{timestamp}.json"
|
||||
|
||||
with open(filename, 'w') as f:
|
||||
json.dump({
|
||||
"method": f"GMM (best by {method.upper()})",
|
||||
"parameters": params,
|
||||
"n_components": params['n_components'],
|
||||
"n_samples": len(labels),
|
||||
f"{method}_score": score,
|
||||
"results": clustering_results
|
||||
}, f, indent=4)
|
||||
|
||||
print(f"Final clustering results ({method}) saved to: {filename}")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Run extensive Gaussian Mixture Model clustering on document embeddings")
|
||||
parser.add_argument("--embeddings_path", required=True, help="Path to embeddings JSON file")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Initialize clustering
|
||||
clustering = GMMExtensiveClustering(args.embeddings_path)
|
||||
|
||||
# Run extensive grid search
|
||||
results = clustering.run_gmm_grid_search()
|
||||
|
||||
if any(labels is not None for labels, _, _ in results.values()):
|
||||
# Visualize and save results
|
||||
clustering.visualize_results(results)
|
||||
clustering.save_clustering_results(results)
|
||||
print("\nGMM extensive clustering completed successfully!")
|
||||
else:
|
||||
print("\nGMM extensive clustering did not find suitable clusters.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user