bangla-retriever

A general-purpose Bangla (Bengali) passage retriever / reranker (bi-encoder) fine-tuned from csebuetnlp/banglabert, using the cosine/dot-product dual-encoder architecture from sayedshaun/wsd. Given a query and a shortlist of candidate passages (from an upstream retrieval step such as BM25 or a dense retriever), it scores each passage and ranks them.

The query and each passage are encoded separately with the same encoder; the score is the dot product of their [CLS] vectors. Merged and self-contained: loads with plain AutoModel, no custom class needed.

Training

• Data: the commercial-ok config of SayedShaun/bangla-retrieval-collection (MIRACL bn, Mr.TyDi bn, MIRACL-nlpai, SWIM-IR bn). 126,706 training examples.
• Example format: 1 query vs. 8 candidates (1 positive + 7 negatives), softmax cross-entropy. MIRACL / Mr.TyDi use their mined hard negatives (≤4 examples per query–positive pair); pair-only SWIM-IR uses random in-subset passages as negatives.
• Settings: max length 512, batch size 1, lr 2e-5, AdamW (wd 0.01), linear warmup 10%, bf16, gradient checkpointing, single 4 GB laptop GPU.
• Early stopping: evaluated every 5,000 steps on a held-out 2% of queries (2,665 examples, split by query); stopped after no improvement ≥0.002 for 3 evaluations.

Checkpoint: step 55000 (of a 70,000-step run; early-stopped) — 0.9122 accuracy on the held-out split (correct passage ranked first out of 8). This is a training-time validation signal on the same data distribution, not a benchmark score. It has not yet been evaluated on mteb/MIRACLReranking or other external benchmarks.

Limitations

• Validation negatives are mostly random (SWIM-IR makes up ~80% of the split), so the held-out accuracy overstates performance against hard, topically close distractors.
• Trained on Wikipedia-style passages; not evaluated on other domains.
• Scores are raw dot products, not calibrated probabilities; don't threshold them without validating on your own data.

Usage

import torch
from transformers import AutoModel, AutoTokenizer

REPO = "SayedShaun/bangla-retriever"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(REPO)
model = AutoModel.from_pretrained(REPO).to(DEVICE).eval()


@torch.inference_mode()
def encode(texts: list[str]) -> torch.Tensor:
    batch = tokenizer(
        texts, return_tensors="pt", padding=True,
        truncation=True, max_length=512,
    ).to(DEVICE)
    return model(**batch).last_hidden_state[:, 0, :].float()  # [CLS]


def rerank(query: str, passages: list[str]) -> list[tuple[str, float]]:
    q, p = encode([query])[0], encode(passages)
    scores = (p @ q).tolist()
    return sorted(zip(passages, scores), key=lambda x: x[1], reverse=True)

License

License for the weights not yet chosen by the author. Training data licenses: MIRACL / Mr.TyDi (Apache-2.0, Wikipedia passages CC-BY-SA), SWIM-IR (CC-BY-SA-4.0); see the dataset card.

Citations

Base encoder:

@inproceedings{bhattacharjee-etal-2022-banglabert,
    title = "{B}angla{BERT}: Language Model Pretraining and Benchmarks for Low-Resource Language Understanding Evaluation in {B}angla",
    author = "Bhattacharjee, Abhik and Hasan, Tahmid and Ahmad, Wasi and Mubasshir, Kazi Samin and Islam, Md Saiful and Iqbal, Anindya and Rahman, M. Sohel and Shahriyar, Rifat",
    booktitle = "Findings of the Association for Computational Linguistics: NAACL 2022",
    year = "2022", address = "Seattle, United States",
    publisher = "Association for Computational Linguistics", pages = "1318--1327",
}

Training pipeline / architecture:

@misc{wsd,
    author = {Md Abu Sayed Shaun},
    title = {Dual Architecture pipeline for Word Sense Disambiguation (WSD)},
    year = {2025},
    howpublished = {\url{https://github.com/sayedshaun/wsd}},
}