RankLLM

To fine-tune a model for RankLLM, use the accelerate launch train_rankllm.py command. You can choose between three training objectives via the --objective flag:

• generation: Traditional language modelling objective.
• ranking: Learning-to-rank objective.
• combined: A combination of both objectives.

Required and common arguments:

• --model_name_or_path: Path to the model to fine-tune.
• --train_dataset_path: Path to the training dataset.
• --num_train_epochs: Number of training epochs.
• --seed: Random seed.
• --per_device_train_batch_size: Batch size per device.
• --gradient_accumulation_steps: Number of gradient accumulation steps.
• --num_warmup_steps: Number of warmup steps.
• --gradient_checkpointing: Enable gradient checkpointing.
• --output_dir: Directory to save the output.
• --noisy_embedding_alpha: Alpha value for noisy embeddings.
• --objective: The training objective (generation, ranking, or combined).

DS_SKIP_CUDA_CHECK=1 NCCL_IB_DISABLE=1 NCCL_P2P_DISABLE=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True accelerate launch train_rankllm.py \
    --model_name_or_path <path-to-model> \
    --train_dataset_path <path-to-train-dataset> \
    --num_train_epochs <num-epochs> \
    --seed <seed> \
    --per_device_train_batch_size <batch-size> \
    --gradient_accumulation_steps <gradient-accumulation-steps> \
    --num_warmup_steps <num-warmup-steps> \
    --gradient_checkpointing \
    --output_dir <output-dir> \
    --noisy_embedding_alpha <noisy-embedding-alpha> \
    --objective <objective>

RankLLM can be used as a post-processor in Llama Index to rerank nodes retrieved from a vector index.

Installation:

pip install llama-index-core llama-index-embeddings-huggingface llama-index-postprocessor-rank-llm rank_llm transformers requests

Usage Pattern:

1. Build a standard Llama Index VectorStoreIndex using HuggingFaceEmbedding.
2. Retrieve nodes using a VectorIndexRetriever.
3. Pass the retrieved nodes to RankLLMRerank.postprocess_nodes().
4. Memory Management: If using rank_zephyr, call del reranker and torch.cuda.empty_cache() after post-processing to free the ~16GB of VRAM required.

from llama_index.core import VectorStoreIndex, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.postprocessor.rankllm_rerank import RankLLMRerank
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core import QueryBundle
import torch

# Setup Index
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
index = VectorStoreIndex.from_documents(documents)

# Retrieval and Reranking
query_bundle = QueryBundle("Which date did Paul Gauguin arrive in Arles?")
retriever = VectorIndexRetriever(index=index, similarity_top_k=50)
retrieved_nodes = retriever.retrieve(query_bundle)

# Apply RankLLM
reranker = RankLLMRerank(model="rank_zephyr", top_n=3, window_size=15)
reranked_nodes = reranker.postprocess_nodes(retrieved_nodes, query_bundle)

# Cleanup
del reranker
torch.cuda.empty_cache()

To use RankLLM within a LangChain workflow, install the necessary dependencies and use the RankLLMRerank compressor with a ContextualCompressionRetriever. This allows you to take a base retriever (like FAISS) and apply RankLLM reranking to the retrieved documents.

Installation:

pip install langchain-community faiss-gpu torch transformers sentence-transformers huggingface-hub rank_llm

Usage Pattern:

1. Set up a standard LangChain retriever (e.g., using FAISS and HuggingFaceEmbeddings).
2. Initialize RankLLMRerank with your desired top_n and model_path.
3. Wrap the base retriever in a ContextualCompressionRetriever using the RankLLM compressor.
4. Note: If using rank_zephyr, it consumes approximately 16GB of GPU VRAM. It is recommended to del compressor and call torch.cuda.empty_cache() after use to free memory.

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.retrievers import ContextualCompressionRetriever
from rank_llm import RankLLMRerank
import torch

# 1. Setup base retriever
documents = TextLoader("state_of_the_union.txt").load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
texts = text_splitter.split_documents(documents)
embedding = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en", model_kwargs={'device': 'cuda'})
retriever = FAISS.from_documents(texts, embedding).as_retriever(search_kwargs={"k": 20})

# 2. Setup Reranker
torch.cuda.empty_cache()
compressor = RankLLMRerank(top_n=3, model_path="rank_zephyr")
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor, base_retriever=retriever
)

# 3. Use
query = "What was done to Russia?"
compressed_docs = compression_retriever.invoke(query)

# 4. Cleanup
del compressor

The rerankers library provides a streamlined way to use RankLLM by specifying model_type="rankllm".

Installation:

pip install "rerankers[rankllm]"

Usage: Initialize the Reranker class with the model name and the specific model_type.

Configuration Arguments for Reranker (model_type="rankllm"):

• model: str (default: "rank_zephyr")
• window_size: int (default: 20)
• context_size: int (default: 4096)
• prompt_template_path: str
• num_few_shot_examples: int (default: 0)
• few_shot_file: Optional[str] (default: None)
• num_gpus: int (default: 1)
• variable_passages: bool (default: False)
• use_logits: bool (default: False)
• use_alpha: bool (default: False)
• stride: int (default: 10)
• use_azure_openai: bool (default: False)

from rerankers import Reranker

ranker = Reranker('rank_zephyr', model_type="rankllm")
results = ranker.rank(
    query="I love you", 
    docs=["I hate you", "I really like you"], 
    doc_ids=[0, 1]
)
print(results)