Reranking

জিনিসটা কী, কেন দরকার

একটা চাকরির জন্য company 2,000 টা CV পায়। কেউ 2,000 টা CV মন দিয়ে পড়ে না। একজন recruiter কয়েকটা keyword দেখে দ্রুত চোখ বুলিয়ে 50 টা রাখে; তারপর hiring manager ওই 50 টা ঠিকঠাক পড়ে সেরা 5 টাকে সাজায়। আগে দ্রুত আর চওড়া, তারপর ধীর আর যত্নের।

Retrieval ও একইভাবে কাজ করে। Vector search (lesson 4) লাখ লাখ chunk এর জন্য যথেষ্ট দ্রুত, কারণ query আর প্রতিটা document আলাদা আলাদা embed হয়, কিন্তু এজন্যই সে একটু ঝাপসাও: দুইটা সারাংশ মেলায়, আসল শব্দগুলো কখনো পাশাপাশি না। একটা reranker (একটা cross-encoder) query আর একটা candidate chunk একসাথে পড়ে আর বিচার করে ওই chunk ওই প্রশ্নের কতটা ভালো answer। লাখ লাখের জন্য বেশি ধীর, আর প্রথম 50 টা নতুন করে সাজানোর জন্য একদম ঠিক। RAG এ সেরা 3 থেকে 5 টা chunk কে প্রথমে রাখা, "প্রথম 20 এর কোথাও" এর বদলে, মান বাড়ানোর সবচেয়ে সস্তা উপায়গুলোর একটা।

ভেতরে কীভাবে চলে

Bi-encoder (ধাপ 1)। Query → vector, document → vector, dot product দিয়ে তুলনা। Document এর vector একবারই, offline এ হিসাব হয়। প্রতিটা query র খরচ একটা embedding আর একটা index খোঁজা: লাখ লাখ chunk এ কয়েক millisecond। খরচটা হলো: model কখনো query আর document একসাথে দেখে না, তাই ধরতে পারে না যে query তে "bKash" শব্দটা "refund" এর চেয়ে বেশি জরুরি, বা একটা document bKash এর কথা বলছে ঠিকই কিন্তু অন্য প্রসঙ্গে।

Cross-encoder (ধাপ 2)। Query আর একটা document একটা input হিসেবে transformer এ যায়, [query] [SEP] [document], আর model একটাই relevance score দেয়। Query র প্রতিটা শব্দ document এর প্রতিটা শব্দে attend করতে পারে (Part 4), তাই হুবহু মিল, না-বোধক কথা আর প্রসঙ্গ দেখে। খরচ: প্রতিটা (query, document) জোড়ায় একবার পুরো model চালানো, আর কিছুই আগে থেকে হিসাব করে রাখা যায় না।

দুই ধাপের retrieval দুইটা মেলায়: ধাপ 1 প্রথম 50 থেকে 100 টা আনে (recall বাড়াও: সঠিক chunk টা কোথাও না কোথাও থাকতে হবে), ধাপ 2 সেগুলো নতুন করে সাজিয়ে প্রথম কয়েকটা রাখে (precision বাড়াও: সেরাগুলো আগে)। Rerank করা score শুধু একটা query র ভেতরে ক্রম ঠিক করতে অর্থবহ, absolute confidence হিসেবে না।

Latency বাজেট। GPU তে 50 টা ছোট chunk এ cross-encoder কয়েক দশ থেকে কয়েকশো millisecond নেয়। কতগুলো candidate rerank করবে সেটা ঠিক করো কতটা সময় দিতে পারবে তা দেখে, উল্টোটা না।

আগে অনুমান করো

Embedding এর মতো সব document এর cross-encoder score আগে থেকে হিসাব করে রাখা যায় না কেন?

Score নির্ভর করে query আর document দুইটার উপর একসাথে, আর query আগে থেকে জানা নেইCross-encoder জমা রাখার পক্ষে বেশি বড়যায়, শুধু খুব একটা করা হয় নাপ্রতিবার model চালালে score বদলে যায়

ছবিতে দেখো

Play

Code এ দুই ধাপের retrieval

আটটা ছোট policy chunk, দুইটা প্রশ্ন। ধাপ 1 local multilingual-e5-small bi-encoder দিয়ে প্রথম 5 টা রাখে; ধাপ 2 সেই 5 টাকে একটা ছোট বহুভাষিক cross-encoder (mmarco-mMiniLMv2-L12-H384-v1) দিয়ে rerank করে, যেটা প্রতিটা query-chunk জোড়া একসাথে পড়ে।

"""Two-stage retrieval: a fast bi-encoder pulls candidates, a cross-encoder reorders the top few."""
import numpy as np
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from embed import embed

CE = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
ce_tok = AutoTokenizer.from_pretrained(CE)
ce = AutoModelForSequenceClassification.from_pretrained(CE).eval()

def cross_scores(query, docs):
    """Read query and document TOGETHER and output one relevance score per pair."""
    enc = ce_tok([query] * len(docs), docs, padding=True, truncation=True, max_length=256, return_tensors="pt")
    with torch.no_grad():
        return ce(**enc).logits[:, 0].numpy()

DOCS = [
    "Refunds for cancelled orders go back to the original payment method",
    "bKash refunds arrive within 3 working days; card refunds take up to 7",
    "Refund requests for missing items must be made within 2 hours",
    "You can pay with bKash, Nagad, card or cash on delivery",
    "bKash outage: payments may fail between 2 and 4 am tonight",
    "Cash on delivery orders are refunded as app credit within 24 hours",
    "To add money to your bKash wallet, visit an agent",
    "Late deliveries over 45 minutes get a 100 taka voucher",
]
D = embed(DOCS)

for q in ["how many days for my bKash money to come back after cancelling",
          "I paid cash and cancelled, how do I get it back"]:
    first = np.argsort(-(D @ embed([q], kind="query")[0]))[:5]       # stage 1: bi-encoder, top 5
    cs = cross_scores(q, [DOCS[i] for i in first])                   # stage 2: cross-encoder on those 5
    print(f"\nquery: {q}")
    print("  stage 1, bi-encoder top 5:")
    for r, i in enumerate(first):
        print(f"    {r + 1}. {DOCS[i]}")
    print("  stage 2, after cross-encoder reranking:")
    for r, j in enumerate(np.argsort(-cs)):
        print(f"    {r + 1}. {cs[j]:6.2f}  {DOCS[first[j]]}")

# Why not use the cross-encoder for everything? It must run once per (query, document) pair.
many = DOCS * 25                                                     # pretend we have 200 documents
cross_scores("bKash refund days", many)
print(f"\ncross-encoder on 200 documents: {len(many)} model runs; on the top 5: 5 runs ({len(many) // 5}x fewer)")

Output:

query: how many days for my bKash money to come back after cancelling
  stage 1, bi-encoder top 5:
    1. bKash refunds arrive within 3 working days; card refunds take up to 7
    2. You can pay with bKash, Nagad, card or cash on delivery
    3. Cash on delivery orders are refunded as app credit within 24 hours
    4. Refunds for cancelled orders go back to the original payment method
    5. bKash outage: payments may fail between 2 and 4 am tonight
  stage 2, after cross-encoder reranking:
    1.   1.73  bKash refunds arrive within 3 working days; card refunds take up to 7
    2.  -1.52  bKash outage: payments may fail between 2 and 4 am tonight
    3.  -4.09  Cash on delivery orders are refunded as app credit within 24 hours
    4.  -7.17  Refunds for cancelled orders go back to the original payment method
    5.  -8.50  You can pay with bKash, Nagad, card or cash on delivery

query: I paid cash and cancelled, how do I get it back
  stage 1, bi-encoder top 5:
    1. Refunds for cancelled orders go back to the original payment method
    2. You can pay with bKash, Nagad, card or cash on delivery
    3. bKash refunds arrive within 3 working days; card refunds take up to 7
    4. To add money to your bKash wallet, visit an agent
    5. Cash on delivery orders are refunded as app credit within 24 hours
  stage 2, after cross-encoder reranking:
    1.  -3.42  Refunds for cancelled orders go back to the original payment method
    2.  -6.09  Cash on delivery orders are refunded as app credit within 24 hours
    3.  -6.22  You can pay with bKash, Nagad, card or cash on delivery
    4.  -6.42  bKash refunds arrive within 3 working days; card refunds take up to 7
    5.  -8.54  To add money to your bKash wallet, visit an agent

cross-encoder on 200 documents: 200 model runs; on the top 5: 5 runs (40x fewer)

কী দেখলাম:

• প্রশ্ন 2 এ বোঝা যায় reranking কেন দরকার। "I paid cash and cancelled" এর দরকার cash-on-delivery chunk, যেটা bi-encoder রেখেছিল 5 নম্বরে, "To add money to your bKash wallet" এর নিচে। Cross-encoder সেটাকে 2 নম্বরে আনে, সাধারণ refund এর নিয়মের ঠিক পরে। LLM এ শুধু প্রথম 3 টা গেলে, এটাই ঠিক আর ভুল answer এর পার্থক্য।
• প্রশ্ন 1: সঠিক chunk ("bKash refunds arrive within 3 working days") আগে থেকেই 1 নম্বরে ছিল আর 1 নম্বরেই থাকে, এবার পরিষ্কার ব্যবধানে (1.73 বনাম -1.52)। Reranker "You can pay with bKash..." কে 2 থেকে শেষে ঠেলে দিয়েছে: এতে bKash আছে কিন্তু প্রশ্নের answer নেই।
• Reranker ও ভুল করে: প্রশ্ন 1 এ "bKash outage" উঠে 2 নম্বরে এসেছে। bKash আর সময়ের শব্দ মেলে কিন্তু অপ্রাসঙ্গিক। Reranker বেশি ধারালো, নিখুঁত না; তোমার data তে মাপো (Part 8)।
• Score শুধু ক্রমের জন্য। প্রশ্ন 2 এর সেরা score -3.42, প্রশ্ন 1 এর দ্বিতীয়টার চেয়েও কম। নিজের data তে মিলিয়ে না নিয়ে cross-encoder score কে absolute "এটা প্রাসঙ্গিক কিনা" threshold হিসেবে ব্যবহার কোরো না।
• খরচ: 200 টা document এ প্রতি query তে 200 বার cross-encoder চালাতে হয়; শুধু প্রথম 5 টা rerank করলে 5 বার, 40 গুণ কম। ধাপ 1 ই ধাপ 2 কে সাধ্যের মধ্যে আনে।

Production এ যা মাথায় রাখবে

• 20 থেকে 100 টা candidate rerank করো, LLM কে 3 থেকে 8 টা পাঠাও। দুইটা সংখ্যা tune করতে ধাপ 1 এর recall@50 আর reranking এর পরের precision@5 মাপো (Part 8)।
• GPU বা hosted reranker ব্যবহার করো ছোট load এর বেশি হলেই; CPU তে cross-encoder দ্রুত ধীর হয়ে যায়। একটা query র সব জোড়া একটা call এ batch করো।
• Document এর দৈর্ঘ্য মাথায় রাখো। Cross-encoder কেটে ফেলে (প্রায়ই 512 token এ)। পুরো document না, chunk rerank করো।
• বহুভাষিক হওয়া জরুরি। তোমার user রা যে ভাষায় লেখে সেগুলোতে train করা reranker বাছো; শুধু English এর reranker ভালো বাংলা chunk কে নিচে ঠেলে দিতে পারে।
• পারলে cache করো: জনপ্রিয় query (আর তাদের rerank করা result) বারবার আসে। অল্প সময়ের একটা cache দুই ধাপই বাঁচায়।
• LLM ও rerank করতে পারে ("এই 20 টা অংশ সাজাও"), ভালো যুক্তি দিয়ে কিন্তু বেশি খরচ আর latency তে; সাধারণত আলাদা একটা cross-encoder ই ভালো সমঝোতা।

এক কথায়

Bi-encoder query আর document আলাদা করে embed করে, তাই লাখ লাখে দ্রুত কিন্তু ঝাপসাভাবে খুঁজতে পারে; cross-encoder প্রতিটা query-document জোড়া একসাথে পড়ে আর অনেক নির্ভুলভাবে সাজায় কিন্তু বড় হতে পারে না; প্রথমটা দিয়ে 50 টা আনো, দ্বিতীয়টা দিয়ে rerank করো, আর শুধু সেরা কয়েকটা LLM কে পাঠাও।

নিজে করে দেখো

"my card refund is late" প্রশ্নটা যোগ করো আর দুইটা ধাপই দেখো। তারপর ধাপ 1 সরিয়ে 8 টা chunk এই cross-encoder চালাও। ক্রম কি বদলায়? প্রতি জোড়ায় 5 ms করে 10 লাখ chunk এ, ধাপ 1 ছাড়া একটা query তে কত সময় লাগবে?

Answer

Card এর প্রশ্নে reranking এর পরে "bKash refunds ... card refunds take up to 7" উপরে আসা উচিত, কারণ শুধু এই chunk এই card refund এর সময়ের কথা আছে; দেখো ধাপ 1 এর প্রথম 5 এ এটা আগে থেকেই ছিল কিনা। মাত্র 8 টা chunk এ, সবগুলো rerank করলে আর প্রথম 5 টা rerank করলে একই সেরা result আসে, যদি ধাপ 1 এর প্রথম 5 এ সঠিকগুলো থাকে। দশ লাখ chunk এ ধাপ 1 ছাড়া: 1,000,000 × 5 ms = 5,000 second, প্রতি query তে প্রায় 83 মিনিট। দুই ধাপের পুরো কারণ এটাই।