cluster_myCollection.py

import json
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans

INDEX_FILE = "index.json"
MY_COLLECTION_FILE = "my_collection.json"

def load_index():
    with open(INDEX_FILE, "r", encoding="utf-8") as f:
        return json.load(f)

def load_my_collection():
    with open(MY_COLLECTION_FILE, "r", encoding="utf-8") as f:
        return set(json.load(f))

def reconstruct_my_docs(index, my_docs):
    doc_texts = {doc: [] for doc in my_docs}

    for term, posting in index["index"].items():
        for doc_id in posting:
            if doc_id in my_docs:
                doc_texts[doc_id].append(term)

    texts = [" ".join(doc_texts[d]) for d in doc_texts]
    doc_ids = list(doc_texts.keys())
    return texts, doc_ids

def run_clustering(texts, k, df_mode="local"):
    vectorizer = TfidfVectorizer(
        stop_words="english",
        max_features=5000,
        min_df=2
    )

    X = vectorizer.fit_transform(texts)

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

    terms = vectorizer.get_feature_names_out()
    order_centroids = model.cluster_centers_.argsort()[:, ::-1]

    clusters = {}

    for i in range(k):
        top_terms = [terms[ind] for ind in order_centroids[i, :50]]
        clusters[i] = top_terms

    return clusters

def main():
    index = load_index()
    my_docs = load_my_collection()

    texts, doc_ids = reconstruct_my_docs(index, my_docs)
    print(f"Using {len(texts)} documents for clustering")

    for k in [2, 10, 20]:
        print(f"\n======== K = {k} ========\n")
        clusters = run_clustering(texts, k)

        for cid, terms in clusters.items():
            print(f"Cluster {cid}:")
            print(", ".join(terms))
            print()

if __name__ == "__main__":
    main()