Metadata-Version: 2.4
Name: k-llmmeans
Version: 0.3.0
Summary: k-LLMmeans clustering algorithm
Author: Joël Weber
Author-email: Joël Weber <joel@joelweber.nl>
License-Expression: MIT
Classifier: License :: OSI Approved :: MIT License
Requires-Dist: litellm>=1.0.0
Requires-Dist: python-dotenv>=1.2.2
Requires-Dist: scikit-learn>=1.7.2
Requires-Dist: tqdm>=4.67.1
Requires-Python: >=3.10
Description-Content-Type: text/markdown

# k-llmmeans

Scikit-learn compatible implementation of **k-LLMmeans** for text clustering with summary-based centroids.

This package adapts the original research code into an estimator API you can use with familiar `fit`, `predict`, and `fit_predict` workflows.

- Original implementation: [jairoadiazr/k-LLMmeans](https://github.com/jairoadiazr/k-LLMmeans)
- Paper: [Summaries as Centroids for Interpretable and Scalable Text Clustering (arXiv:2502.09667)](https://arxiv.org/abs/2502.09667)

## What This Package Provides

- `kLLMmeans` estimator implementing `BaseEstimator` + `ClusterMixin`
- scikit-learn style methods:
  - `fit(X)`
  - `predict(X)`
  - `fit_predict(X)`
- configurable document embedding function (`embedding_fn`)
- configurable cluster summarization function (`summarizer_fn`) or LiteLLM-backed LLM summarization
- optional per-cluster sampling before summarization to keep prompts bounded on large clusters
- optional precomputed embedding support for faster iterative experimentation

## Installation

```bash
pip install k-llmmeans
```

Or from source:

```bash
pip install -e .
```

## Quick Start

```python
from k_llmmeans import kLLMmeans

docs = [
    "How to optimize SQL queries for large tables?",
    "What is the best way to tune a random forest model?",
    "PostgreSQL index strategy for analytics workloads",
    "Cross-validation tips for imbalanced classification",
]

model = kLLMmeans(
    n_clusters=2,
    llm="openai/gpt-4o-mini",
    max_llm_iter=5,
    random_state=0,
)

labels = model.fit_predict(docs)
print(labels)
print(model.summaries_)  # human-readable cluster summaries
```

## LiteLLM Configuration

Pass a model string (any [LiteLLM-supported model](https://docs.litellm.ai/docs/providers)) or a dict of kwargs forwarded to `litellm.completion`:

```python
model = kLLMmeans(
    n_clusters=2,
    llm={
        "model": "openai/gpt-4o-mini",
        "temperature": 0.2,
    },
)
```

You can also set `LITELLM_MODEL` or `OPENAI_MODEL` in the environment instead of passing `llm`.

## Using Custom Embeddings and Summarization

You can fully control both the embedding and summarization steps:

```python
from sentence_transformers import SentenceTransformer
from k_llmmeans import kLLMmeans

encoder = SentenceTransformer("all-MiniLM-L6-v2")

def embedding_fn(texts: list[str]):
    return encoder.encode(texts)

def summarizer_fn(cluster_texts: list[str]) -> str:
    # Replace with your own deterministic or LLM summarizer
    return " | ".join(cluster_texts[:2])

model = kLLMmeans(
    n_clusters=3,
    embedding_fn=embedding_fn,
    summarizer_fn=summarizer_fn,
)

model.fit(["text a", "text b", "text c", "text d"])
```

## Sampling Cluster Documents for Summaries

The paper's few-shot k-LLMmeans variant summarizes at most `m` representative documents from each cluster instead of sending the full cluster to the LLM. Use `cluster_sample_size` for this prompt budget:

```python
model = kLLMmeans(
    n_clusters=3,
    llm="openai/gpt-4o-mini",
    cluster_sample_size=10,
    cluster_sampling_strategy="kmeans++",
)
```

Set `cluster_sample_size=None` to summarize every document assigned to each cluster. Supported sampling strategies are `kmeans++`, `random`, `centroid`, and `edge`. The selected document indices for each LLM iteration are stored in `sampled_indices_evolution_`.

## Debugging Runs

Use `show_progress=True` to display progress bars for LLM iterations and cluster summarization. Use `verbose=True` to print timing, input character counts, sampled document counts, summary lengths, token usage when available, and convergence details:

```python
model = kLLMmeans(
    n_clusters=3,
    llm="openai/gpt-4o-mini",
    cluster_sample_size=10,
    show_progress=True,
    verbose=True,
)
```

After fitting, inspect `summaries_evolution_` to see how cluster summaries changed, `centroids_evolution_` to inspect summary-derived centroid embeddings, and `sampled_indices_evolution_` to trace which input texts were used for each cluster summary. `summary_workers` can parallelize cluster summarization when your LLM provider and rate limits allow concurrent requests.

## API Notes

- Input `X` should be `list[str]`.
- The estimator stores standard fitted attributes such as:
  - `labels_`
  - `cluster_centers_`
  - `n_iter_`
- Additional clustering interpretability attributes:
  - `summaries_`
  - `summary_embeddings_`
  - `summaries_evolution_`
  - `centroids_evolution_`
  - `sampled_indices_evolution_`

## Citation

If you use this package in research or production work, please cite the original paper:

```bibtex
@article{diazrodriguez2025summaries,
  title={Summaries as Centroids for Interpretable and Scalable Text Clustering},
  author={Diaz-Rodriguez, Jairo},
  journal={arXiv preprint arXiv:2502.09667},
  year={2025}
}
```

Paper URL: [https://arxiv.org/abs/2502.09667](https://arxiv.org/abs/2502.09667)

## Acknowledgment

This package is a scikit-learn compatible adaptation of the original project:
[https://github.com/jairoadiazr/k-LLMmeans](https://github.com/jairoadiazr/k-LLMmeans)
