XiaSheng/FreeChunk-bge-m3

FreeChunker-BGE-M3

FreeChunker is a training-free embedding optimization method that dynamically chunks text to improve retrieval performance. This repository contains the FreeChunker model initialized with BAAI/bge-m3 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 BAAI/bge-m3 sentence embeddings.

Requirements

pip install torch transformers sentence-transformers numpy

Usage

Abstract Usage

from transformers import AutoModel
import torch

# 1. Load Model (UnifiedEncoder)
model = AutoModel.from_pretrained("XiaSheng/FreeChunk-bge-m3", 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:

1. Split and Encode: Use Sentenceizer (wrapping BAAI/bge-m3) to get sentence embeddings.
2. FreeChunker: Pass embeddings to FreeChunkerModel.
3. Process: Use the output shift_matrix to group sentences.

from sentenizer import Sentenceizer
from modeling_freechunker import FreeChunkerModel
import torch

# 1. Setup Sentenceizer with Backbone
sentenceizer = Sentenceizer(model_name="BAAI/bge-m3")

# 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:

@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} 
 }