luluw/bge-reranker-v2-m3-eng-nep-16k-trimmed

BGE-Reranker-v2-M3 (English + Nepali) - Trimmed Vocabulary (16k)

This is a vocabulary-trimmed version of BAAI/bge-reranker-v2-m3.

It is a cross-encoder reranker: it scores a (query, passage) pair jointly and outputs a single relevance logit. It does not produce embeddings for a vector index — use it to re-score/re-order a candidate list already retrieved by a first-stage retriever (e.g. BM25 or a dense embedding model).

• Original vocab: 250,002 tokens (100+ languages)
• Trimmed vocab: 16,384 tokens focused on English + Nepali
• Classification head: unchanged (num_labels=1), it operates on pooled hidden states and does not depend on vocab size.
• No fine-tuning was performed - the original trained embedding vectors for kept tokens were copied as-is.

How it was created

1. Token frequency was counted on real English and Nepali text from the lbourdois/fineweb-2-trimming dataset.
2. Special tokens + the first 1,000 original IDs were always kept.
3. The remaining budget was filled with the highest-frequency English and Nepali tokens (weighted 50/50).
4. The input embedding matrix was rebuilt by copying the original rows for every kept token (get_input_embeddings() / set_input_embeddings()), the classification head was left untouched.
5. An old_id -> new_id mapping is provided so the original XLM-R tokenizer can still be used for subword splitting.

Usage

Because the vocabulary was remapped, you must use the provided mapping and pair-encode queries with passages (<s> query </s></s> passage </s>). Do not feed raw original tokenizer IDs into the trimmed model.

import json
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from huggingface_hub import hf_hub_download

repo_id = "luluw/bge-reranker-v2-m3-eng-nep-16k-trimmed"

model = AutoModelForSequenceClassification.from_pretrained(repo_id)
model.eval()

tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder="original_tokenizer")
with open(hf_hub_download(repo_id, "vocab_mapping.json"), "r", encoding="utf-8") as f:
    vocab_data = json.load(f)

old_to_new = {int(k): v for k, v in vocab_data["old_to_new"].items()}
new_unk_id = old_to_new[tokenizer.unk_token_id]
new_pad_id = old_to_new[tokenizer.pad_token_id]

def encode_pair_trimmed(query, passage, max_length=1024):
    old_ids = tokenizer.encode(query, passage, add_special_tokens=True, truncation=True, max_length=max_length)
    return [old_to_new.get(i, new_unk_id) for i in old_ids]

def rerank(pairs, max_length=1024, batch_size=16, apply_sigmoid=True):
    scores = []
    with torch.no_grad():
        for start in range(0, len(pairs), batch_size):
            batch = pairs[start:start + batch_size]
            id_lists = [encode_pair_trimmed(q, p, max_length=max_length) for q, p in batch]
            max_len = max(len(ids) for ids in id_lists)
            input_ids = torch.full((len(id_lists), max_len), new_pad_id, dtype=torch.long)
            attn_mask = torch.zeros((len(id_lists), max_len), dtype=torch.long)
            for i, ids in enumerate(id_lists):
                input_ids[i, :len(ids)] = torch.tensor(ids, dtype=torch.long)
                attn_mask[i, :len(ids)] = 1
            logits = model(input_ids=input_ids, attention_mask=attn_mask).logits.view(-1)
            batch_scores = torch.sigmoid(logits) if apply_sigmoid else logits
            scores.extend(batch_scores.tolist())
    return scores

pairs = [
    ("Where is the Eiffel Tower?", "The Eiffel Tower is located in Paris, France."),
    ("Where is the Eiffel Tower?", "Python is a popular programming language."),
]
print(rerank(pairs))  # relevant pair should score higher

Notes

• Only the reranker's dense backbone + input embeddings are affected; the classification head is untouched.
• Very rare domain-specific tokens may map to <unk>.
• A short fine-tune can still improve quality, but the model is usable out of the box.

Citation

@misc{bge_m3,
  title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
  author={Chen, Jianlv and Xiao, Shitao and Zhang, Peitian and Luo, Kun and Lian, Defu and Liu, Zheng},
  year={2024},
  eprint={2402.03216},
  archivePrefix={arXiv},
  primaryClass={cs.CL}
}