cluster_documents.py

# cluster_documents.py
import json
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans

def load_index():
    with open("index.json") as f:
        data = json.load(f)
    index = data["index"]
    docs = data["documents"]
    return index, docs

def reconstruct_docs(index, docs):
    doc_texts = {}

    for term, posting in index.items():
        for doc_id, freq in posting.items():
            doc_texts.setdefault(doc_id, [])
            doc_texts[doc_id].extend([term] * freq)

    # Only keep documents that actually have tokens
    valid_docs = []
    valid_texts = []

    for d in docs:
        if d in doc_texts and len(doc_texts[d]) > 0:
            valid_docs.append(d)
            valid_texts.append(" ".join(doc_texts[d]))

    print(f"Using {len(valid_docs)} documents for clustering")

    return valid_texts, valid_docs

def cluster_texts(texts, doc_ids, k):
    vectorizer = TfidfVectorizer(max_features=5000)
    X = vectorizer.fit_transform(texts)

    model = KMeans(n_clusters=k, random_state=42)
    labels = model.fit_predict(X)

    terms = vectorizer.get_feature_names_out()
    centers = model.cluster_centers_

    clusters = {}
    for cluster_idx in range(k):
        weights = centers[cluster_idx]
        top_idx = weights.argsort()[::-1][:50]
        clusters[cluster_idx] = [terms[i] for i in top_idx]

    return clusters, labels

def main():
    index, docs = load_index()
    texts, doc_ids = reconstruct_docs(index, docs)

    for k in [2, 10, 20]:
        clusters, labels = cluster_texts(texts, doc_ids, k)
        print(f"\n======== K = {k} ========")
        for c, terms in clusters.items():
            print(f"\nCluster {c}:")
            print(", ".join(terms))

if __name__ == "__main__":
    main()