Metadata-Version: 2.3
Name: chunkey-bert
Version: 1.0.0
Summary: Modification of the KeyBERT method to extract keywords and keyphrases using chunks. This provides better results, especialy when handling long documents.
Keywords: machine learning
Author: Yaniv Shulman
Author-email: yaniv@shulman.info
Requires-Python: >=3.10,<3.13
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Dist: keybert (>=0.9.0,<0.10.0)
Project-URL: Homepage, https://github.com/yaniv-shulman/chunkey-bert
Project-URL: Repository, https://github.com/yaniv-shulman/chunkey-bert
Description-Content-Type: text/markdown

[![Downloads](https://static.pepy.tech/badge/chunkey-bert)](https://pepy.tech/project/chunkey-bert)
![Tests](https://github.com/yaniv-shulman/chunkey-bert/actions/workflows/linting_and_tests.yml/badge.svg?branch=main)
[![Pyversions](https://img.shields.io/pypi/pyversions/chunkey-bert.svg?style=flat-square)](https://pypi.python.org/pypi/chunkey-bert)

# ChunkeyBERT - Unsupervised Keyword Extraction from Long Documents #
## Overview ##
ChunkeyBert is a minimal and easy-to-use keyword extraction technique that leverages embeddings for unsupervised
keyphrase extraction from text documents. ChunkeyBert is a modification of the
[KeyBERT method](https://towardsdatascience.com/keyword-extraction-with-bert-724efca412ea) to handle documents with
arbitrary length with better results. ChunkeyBERT works by chunking the documents and uses KeyBERT to extract candidate
keywords/keyphrases from all chunks followed by a configurable merge stage to produce the final keywords for the entire
document. ChunkeyBert can use any document chunking method as long as it can be wrapped in a simple function, however it
can also work without a chunker and process the entire document as a single chunk. ChunkeyBert works with any
configuration of KeyBERT and can handle batches of documents.

## Installation ##
Install from [PyPI](https://pypi.org/project/chunkey-bert/) using pip (preferred method):
```bash
pip install chunkey-bert
```

## Use Cases ##
ChunkeyBERT is most useful when you want lightweight, deterministic keyword extraction for long documents without relying
on an LLM. Typical use cases include:

- tagging and indexing large document collections
- search, clustering, and downstream retrieval features
- privacy-sensitive or offline processing pipelines
- batch extraction where LLM latency or cost would be too high
- generating candidate phrases for later reranking or normalization by an LLM

In LLM-based systems, ChunkeyBERT is often best used as a cheap first-stage candidate extraction layer rather than as a
replacement for generative semantic analysis.

## Details ##
### How does ChunkeyBERT differs to KeyBERT? ###
ChunkeyBERT differs from KeyBERT primarily in its approach to handling long documents for keyword extraction. While 
KeyBERT directly applies keyword extraction techniques to the entire document, ChunkeyBERT introduces an additional step
of chunking the document into smaller, manageable pieces before applying KeyBERT's keyword extraction methods. This 
modification aims to improve the performance and relevance of the extracted keywords, especially for longer documents 
where directly applying KeyBERT might not yield optimal results due to the complexity and size of the document. Here are
the key differences:

**Document chunking**: ChunkeyBERT uses a chunking method to divide a long document into smaller chunks. This is done through the chunker 
parameter in the extract_keywords method. The chunker can be any callable that takes a string (the document) and returns
a list of strings (the chunks). This allows ChunkeyBERT to process each chunk independently, making it more effective at
handling long documents. A chunker could be as simple as 
```python
chunker: Callable[[str], List[str]] = lambda text: [t for t in text.split("\n\n") if len(t) > 25]
```
or can wrap more complicated logic such as a Langchain chunker for example.

**Handling of chunks**: After chunking, ChunkeyBERT applies KeyBERT's keyword extraction to each chunk separately.

**Keyword scoring and selection**: ChunkeyBERT introduces additional logic to score and select keywords based on their
occurrence across different chunks and their similarity.

**Merge strategies**: ChunkeyBERT supports multiple strategies for merging chunk-level keywords into document-level
results. The default `"similarity"` strategy ranks keywords by semantic centrality, while `"count"` ranks them by
repetition across chunks. Custom merge callables are also supported.

### Flexibility in keyword extraction ###
ChunkeyBERT offers flexibility in keyword extraction in a number of ways. It can work with any configuration of KeyBERT
and exposes a superset of KeyBERT's extract_keywords() API, which allow fine-tuning of the keyword extraction process 
based on the characteristics of the chunks and the overall document. It can also work with any chunking method including
semantic chunking, chunk filtering and even sampling from the document to finetune the process. ChunkeyBERT can be
configured to consider the multiplicity of keywords across chunks to account for repetitions, and the merge strategy
can be selected explicitly depending on the behavior you want.

### Batching and GPU support ###
ChunkeyBERT works with document batches and attempts to process these batches in parallel on the GPU if possible.

### Compatible with KeyBERT return values ###
ChunkeyBERT returns results in a format similar to KeyBERT but can also optionally return the embeddings for each of the
keywords extracted.

## Usage ##

The following steps describe a basic example on how use ChunkeyBert for keyword extraction:

**Install ChunkeyBert**: First, ensure that ChunkeyBert is installed in your environment. You can install it using pip as
shown below:

```bash
pip install chunkey-bert
```
**Import required libraries**: Import the necessary libraries including ChunkeyBert, KeyBERT, and any other dependencies you
might need for your specific use case.

```python
from keybert import KeyBERT
from sentence_transformers import SentenceTransformer
from chunkey_bert.model import ChunkeyBert
```

**Initialize KeyBERT**: this could be done for example using a Sentence Transformer model that is used to generate embeddings
for the text. _Note that the quality of extracted keywords depends greatly on how KeyBERT is configured_, so it is 
required to understand how to use KeyBERT effectively.

```python
sentence_model = SentenceTransformer(model_name_or_path="all-MiniLM-L6-v2")
keybert = KeyBERT(model=sentence_model)
```

**Define a chunker function (optional)**: If you want to chunk your text into smaller parts (which is the main feature of
ChunkeyBert), define a chunker function. This function takes a string and returns a list of strings (chunks). If you
don't provide a chunker, ChunkeyBert will process the entire document as a single chunk but will still apply a different
keywords selection method to KeyBERT. Here is an example of a very simple chunker:

```python
chunker = lambda text: [t for t in text.split("\n\n") if len(t) > 25]  # Example chunker that splits text into paragraphs
```

**Create a ChunkeyBert instance**: Initialize ChunkeyBert with the KeyBERT instance you created earlier.

```python
chunkey_bert = ChunkeyBert(keybert=keybert)
```
**Extract keywords**: Use the extract_keywords method of ChunkeyBert to extract keywords from your document. You can specify
the number of keywords, whether to use the chunker, the merge strategy, and other parameters related to keyword extraction and to
KeyBERT.extract_keywords.

```python
text = "Your long document text goes here..."

keywords = chunkey_bert.extract_keywords(
    docs=text,
    num_keywords=10,
    chunker=chunker,  # Pass your chunker here. If None, the entire document is treated as a single chunk.
    merge_strategy="similarity",  # Built-in strategies: "similarity" (default) or "count".
    top_n=3,  # Number of keywords to extract from each chunk
    nr_candidates=20,  # Number of candidate keywords/keyphrases to consider from each chunk
)

print(keywords)
```
To rank keywords by repetition across chunks instead of semantic centrality:

```python
keywords = chunkey_bert.extract_keywords(
    docs=text,
    num_keywords=10,
    chunker=chunker,
    merge_strategy="count",
    use_count_weights=True,  # Required by the built-in "count" strategy.
    top_n=3,
)
```

Custom merge logic can also be provided as a callable:

```python
def my_merge_strategy(
    embeddings_doc: np.ndarray,
    counts_doc: Optional[np.ndarray],
    top_k: Optional[int],
) -> tuple[np.ndarray, np.ndarray]:
    ...
```

See a more advanced example in this notebook: https://nbviewer.org/github/yaniv-shulman/chunkey-bert/tree/main/src/experiments/

## Experimental results ##
Example notebooks are available at https://nbviewer.org/github/yaniv-shulman/chunkey-bert/tree/main/src/experiments/.

## Benchmarks ##
The repository now includes a reproducible long-document benchmark runner:

```bash
source configure.sh
python src/experiments/run_keyword_benchmarks.py
```

After sourcing `configure.sh`, benchmark outputs are written to `${REPO_DIR}/out/benchmarks` and `${OUT_DIR}` is
exported automatically. The runner writes:

- per-profile JSON results with per-document metrics
- a markdown summary report
- simple SVG comparison charts

The current built-in benchmark profiles focus on long scientific papers from `midas/krapivin`, which is a better fit
for ChunkeyBert than short web snippets because the method is explicitly designed for long documents.

### Current benchmark snapshot ###
On the `100` longest documents from the `midas/krapivin` test split:

- average document length: `14,937.68` tokens
- exact `F1@10`: `0.0117` for KeyBERT vs `0.0756` for ChunkeyBert
- stemmed `F1@10`: `0.0182` for KeyBERT vs `0.0896` for ChunkeyBert
- exact `hit@10` improvement: `+0.39`
- bootstrap `95%` CI for exact `F1@10` delta: `[0.0487, 0.0802]`
- bootstrap `95%` CI for stemmed `F1@10` delta: `[0.0554, 0.0876]`

This is enough to say, with reasonable confidence, that ChunkeyBert is materially stronger than plain KeyBERT on very
long scientific documents under present-keyphrase matching. It is not evidence that ChunkeyBert is universally better
across all keyword extraction settings or short-document benchmarks.

For more detail, see:

- [benchmark analysis notebook](src/experiments/benchmark_results_analysis.ipynb)
- [qualitative comparison notebook](src/experiments/compare_with_keybert_wikipedia_example.ipynb)

The generated markdown report and SVG charts are written to `out/benchmarks/` when the benchmark runner is executed.

## Contribution and feedback ##
Contributions and feedback are most welcome. Please see
[CONTRIBUTING.md](https://github.com/yaniv-shulman/chunkey-bert/tree/main/CONTRIBUTING.md) for further details.

