FreeChunker is a training-free embedding optimization method that dynamically chunks text to improve retrieval performance. This repository contains the FreeChunker model initialized with jinaai/jina-embeddings-v2-small-en embeddings.

Features

Dynamic Chunking: Automatically groups sentences into semantically coherent chunks.
Optimized for RAG: Improves retrieval augmented generation by providing better context segments.
Backbone: Built on top of jinaai/jina-embeddings-v2-small-en sentence embeddings.

Requirements

pip install torch transformers sentence-transformers numpy

Usage

Abstract Usage

python

from transformers import AutoModel
import torch

# 1. Load Model (UnifiedEncoder)
model = AutoModel.from_pretrained("XiaSheng/FreeChunk-jina", trust_remote_code=True)

# 2. Build Vector Store from Text
text = "Your text..."
model.build_vector_store(text)

# 3. Query with Post-Aggregation (Default)
query = "Your query..."
results = model.query(query, top_k=1, aggregation_mode='post')

print(f"Query: {query}")
print(f"Result: {results}")

Manual Pipeline

If you prefer to use the components separately:

Split and Encode: Use Sentenceizer (wrapping jinaai/jina-embeddings-v2-small-en) to get sentence embeddings.
FreeChunker: Pass embeddings to FreeChunkerModel.
Process: Use the output shift_matrix to group sentences.

python

from sentenizer import Sentenceizer
from modeling_freechunker import FreeChunkerModel
import torch

# 1. Setup Sentenceizer with Backbone
sentenceizer = Sentenceizer(model_name="jinaai/jina-embeddings-v2-small-en")

# 2. Load FreeChunker Model
model = FreeChunkerModel.from_pretrained(".", trust_remote_code=True)
model.eval()

# 3. Process Text
text = "Your text..."
sentences, embeddings = sentenceizer.split_and_encode(text)

# 4. Forward pass through FreeChunker
inputs_embeds = torch.tensor(embeddings).unsqueeze(0) # Batch size 1
with torch.no_grad():
    outputs = model(inputs_embeds=inputs_embeds)
    
# outputs['embedding'] contains refined embeddings
# outputs['shift_matrix'] contains chunking information

Files

model.safetensors: The FreeChunker model weights.
encoder.py: High-level interface (UnifiedEncoder) for end-to-end usage.
sentenizer.py: Helper for text splitting and backbone embedding.
aggregator.py: Helper for aggregating retrieved results.
configuration_freechunker.py & modeling_freechunker.py: Model definition.

Citation

If you use this model in your research, please cite:

bibtex

@article{zhang2025freechunker, 
   title={FreeChunker: A Cross-Granularity Chunking Framework}, 
   author={Zhang, Wenxuan and Jiang, Yuan-Hao and Wu, Yonghe}, 
   journal={arXiv preprint arXiv:2510.20356}, 
   year={2025} 
 }