From pipelines to prompts

NLP in the LLM era

Author

Marie-Hélène Burle

Content from the webinar slides for easier browsing.

4 eras of NLP

Era 1: Symbolic era

1950s–1970s

Complex, hard-coded rules written by linguists.

Era 3: Dense retrieval

2000s–2010s

Machine learning & early deep learning models.

Vector embeddings, e.g. Word2Vec (2013): mathematical representation of words based on their semantic meanings & contextual relationships.

Recurrent Neural Networks (RNNs) such as Long Short-Term Memory (LSTM) networks for speech recognition & translation.

Era 4: LLMs

2017–present

Transformer-based large language models.

Self-attention mechanism, massive scaling, generative AI.

BERT, GPT, Gemini, Claude, Kimi, GLM, Llama, etc.

The question

If LLMs can do everything, are task-specific methods dead?

Let’s find out with one test case.

Benchmark: BEIR/SciFact

The SciFact [1] dataset is a benchmark for scientific claim verification (2020), part of the BEIR [2] standard suite of information retrieval benchmarks:

  • 5,183 scientific abstracts,
  • 1,409 scientific claims,
  • Expert annotators’ ground truth.

Get the data

curl -L -o scifact.zip \
  https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/scifact.zip

Alternative methods to get the data:

from datasets import load_dataset
dataset = load_dataset("allenai/scifact")

Structure of the data

scifact.zip unzips to a scifact dir with:

corpus.jsonl: a corpus of research abstracts
➔ _id, title, text (abstract of scientific paper).

queries.jsonl: scientific claims to be verified against the corpus
➔ id, text (claim), label (contradict / support).

qrels dir with train.tsv & test.tsv: relevance judgments for training & testing
➔ corpus-id, query-id, score.

Methods

Lexical ranking algorithm
Sparse keyword search

Dense retrieval
Semantic vectors via embeddings

Reciprocal Rank Fusion
Hybrid keyword + embeddings

LLMs
Reranking the shortlist

Metric: nDCG@10

Normalized Discounted Cumulative Gain at rank 10 is a standard ranking quality metric in information retrieval, search engines & recommendation systems.

Measures how well a system places the most relevant items within the first 10 results.

def ndcg_at_k(ranked, qrels, k=10):
    dcg = sum(qrels.get(d, 0) / np.log2(r + 2)
              for r, d in enumerate(ranked[:k]))
    ideal = sorted(qrels.values(), reverse=True)[:k]
    idcg = sum(g / np.log2(r + 2) for r, g in enumerate(ideal))
    return dcg / idcg if idcg else 0.0

Scores in context

Across the 18 datasets of BEIR, BM25 averages an nDCG@10 of 0.48 [2].

SciFact is one of the most keyword-friendly datasets of the suite: scientific claims are full of rare exact terms.

Rule of thumb on this benchmark: 0.70+ is strong & 0.80+ is exceptional.

Hardware: one small GPU

For this test, the hardware used was a laptop with 8 cores, 16 threads, 30 GiB of RAM, and a small 6 GiB GPU.

Load packages

import json
import re

import numpy as np
# https://github.com/ollama/ollama-python
import ollama
import torch
import torch.nn.functional as F
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
# https://github.com/dorianbrown/rank_bm25
from rank_bm25 import BM25Okapi
from torch.utils.data import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    Trainer,
    TrainingArguments,
)

Load corpus

corpus = {}  # doc id -> {"title": ..., "text": ...}

with open("scifact" / "corpus.jsonl", encoding="utf-8") as f:
    for line in f:
        doc = json.loads(line)
        corpus[doc["_id"]] = {
            "title": doc.get("title") or "",
            "text": doc["text"],
        }

Load queries

queries = {}  # query id -> text

with open("scifact" / "queries.jsonl", encoding="utf-8") as f:
    for line in f:
        q = json.loads(line)
        queries[q["_id"]] = q["text"]

Load relevance judgments

qrels = {}  # query id -> {doc id: relevance}

with open("scifact" / "qrels" / "test.tsv", encoding="utf-8") as f:
    next(f)
    for line in f:
        qid, did, score = line.split("\t")
        qrels.setdefault(qid, {})[did.strip()] = int(score)

Lexical ranking algorithm

BM25

A bag-of-words retrieval function ranks documents by how relevant they are to a search query. It uses literal keyword matches (1994).

Improved version of term frequency–inverse document frequency (TF-IDF) with length normalization.

No training, CPU-based, still widely used (e.g. Elasticsearch).

Fails on vocabulary mismatch & paraphrase.

Example: “heart attack” does not match “myocardial infarction”.

Code

def tokenize(text):
    """Crude but sufficient: lowercase alphanumeric tokens."""
    return re.findall(r"[a-z0-9]+", text.lower())

# Build the index
bm25 = BM25Okapi([tokenize(full_text[d]) for d in docs])

# Score every document
scores = bm25.get_scores(tokenize(query))

# Sort
ranking = [docs[i] for i in np.argsort(-scores)]

Result

BM25  nDCG@10 0.652  (15.0 ms/query)

0.652 with zero training, one second of indexing & 15 ms per query.

A very serious baseline on this dataset.

Dense retrieval

The idea

A trained embedding model maps texts to vectors based on meaning:

  • Index: embed every abstract once,
  • Search: embed the query, rank by cosine similarity,
  • Solves vocabulary mismatch: “heart attack” & “myocardial infarction” have close vectors.

Dense retrieval does not scan text; it compares coordinates. That is why it needs an index of embeddings.

3 model sizes

Model Params Year Disk
all-minilm 23M 2020 45 MB
nomic-embed-text 137M 2024 274 MB
mxbai-embed-large 335M 2024 669 MB


Pulled by ollama:

ollama pull all-minilm nomic-embed-text mxbai-embed-large

Served locally:

ollama serve

Embedding the corpus

def embed(texts, model, batch=64):
    out = []
    for i in range(0, len(texts), batch):
        out.extend(ollama.embed(model=model,
                                input=texts[i:i+batch])["embeddings"])
    emb = np.asarray(out, dtype=np.float32)
    # L2-normalization
    emb /= np.linalg.norm(emb, axis=1, keepdims=True)  # unit vectors
    return emb

doc_emb = embed([full_text[d] for d in docs], model)

L2-normalization makes every later dot product a cosine similarity.

On the laptop GPU, indexing 5,183 abstracts took 150 s with the largest model.

Searching

q = np.asarray(
    ollama.embed(model=model, input=[query])["embeddings"][0],
    dtype=np.float32,
)
# L2-normalization
q /= np.linalg.norm(q)

ranking = [docs[i] for i in np.argsort(-(doc_emb @ q))]

One matrix-vector product against the whole corpus: 20 milliseconds.

Each model gets its recommended query format
Abstracts are truncated at each model’s own context limit.

Results

all-minilm         nDCG@10 0.648  (20.1 ms/query)
nomic-embed-text   nDCG@10 0.703  (18.1 ms/query)
mxbai-embed-large  nDCG@10 0.741  (22.3 ms/query)

The 2020 miniature model ties BM25. The modern models do better.

Bigger embedding model → better rankings at the same speed.

Reciprocal Rank Fusion (RRF)

The idea

Fuse two rankings by rank position, no score calibration needed [3].

RRF_K = 60  # classic constant (Cormack et al. 2009)

def rrf(rankings):
    fused = {}
    for qid in eval_qids:
        scores = {}
        for ranking in rankings.values():
            for rank, d in enumerate(ranking[qid]):
                scores[d] = scores.get(d, 0.0) + 1.0 / (RRF_K + rank)
        fused[qid] = sorted(scores, key=scores.get, reverse=True)
    return fused

Results

RRF (nomic + BM25):   0.702 -> 0.706   # helped a little
RRF (mxbai + BM25):   0.741 -> 0.731   # hurt

Whether or not a hybrid method helps depends:

  • With a mediocre embedder, fusion buys back what each method misses,
  • Once one signal is much stronger, fusing drags it down.

Hybrid search is a patch for weak retrieval.

LLMs

Commercial LLMs

It is cheap and easy to use a large, last-generation LLM running on the platform of a company (e.g. GPT, Gemini, Claude).

These models are now excellent at most tasks, including NLP ones.

There are many situations, however, in which they are not suitable:

  • Reproducibility (companies can tweak, change, or take down models as they wish),
  • Sensitive data,
  • Concerns about bias in the models,
  • Training on specific data.

Open-weight models

Open-weight models offer a great alternative, but very large ones (e.g. Kimi, GLM, Llama) require massive amounts of compute to run locally, fine-tune, or train.

Thankfully, small open-weight LLMs that run on a laptop can be powerful after fine-tuning.

Some open-source tools make this easy:

  • Ollama to run them locally,
  • unsloth to train, fine-tune & serve them,
  • SGLang to serve them in production.

We will use two very small versions of Llama.

The constraint

A generative model cannot scan 5,183 abstracts: it has no index & reads tokens one prompt at a time.

The standard pattern is retrieve-then-rerank:

  • Retrieve a shortlist: top-10 BM25 or top-10 dense (≤ 20 documents),
  • Ask an LLM to reorder the shortlist,
  • One LLM call per query: extremely high compute demand.

Naive prompt

resp = ollama.chat(
    model=model,
    messages=[
        {"role": "system", "content": "You are a search reranker."},
        {"role": "user", "content":
            f"Query: {query}\n\nDocuments:\n{listing}\n\n"
            "Rank the documents by relevance to the query. "
            "Reply with the document numbers, most relevant first."},
    ],
    format=RERANK_FORMAT,   # JSON schema: {"ranking": [integers]}
    options={"temperature": 0, "num_predict": 200, "num_ctx": 16384},
)

Full abstracts (median ~360 tokens), so the prompt needs a 16k-token context window.

Issues

Several replies per run could not be parsed:

query: 'The combination of H3K4me3 and H3K79me2 is found in
        quiescent hair follicle stem cells.'
response: '{ "ranking": [1, 3, 11, 12, 13, 14, 15, 16, 17, 18,
            19, 4, 6, 8, 9, 10, 5, 7, 2, 5, 5, 5, 5, 5, 5,
error: Expecting value: line 32 column 3 (char 168)

The model counts & loops over numbers until the generation cap cuts the JSON mid-list: 9 of 50 replies unparseable.

The fallback hid the failures

The reply is text: code has to turn it into a ranking.

try:
    nums = json.loads(resp["message"]["content"])["ranking"]
    ranked = [shortlist[n - 1] for n in nums
              if 1 <= n <= len(shortlist)]
except (ValueError, KeyError, IndexError):
    ranked = shortlist   # fallback: keep the retriever order

Every unparseable reply silently keeps the retriever’s ranking:
the measured score was partly the retriever’s work, not the LLM’s.

Better prompt

State the number of documents, the valid range & the reply length.

resp = ollama.chat(
    model=model,
    messages=[
        {"role": "system", "content": "You are a search reranker."},
        {"role": "user", "content":
            f"Query: {query}\n\nDocuments:\n{listing}\n\n"
            f"There are {len(shortlist)} documents, numbered "
            f"1 to {len(shortlist)}. Reply with the 10 most relevant "
            "document numbers, most relevant first."},
    ],
    format=RERANK_FORMAT,   # JSON schema: {"ranking": [integers]}
    options={"temperature": 0, "num_predict": 200, "num_ctx": 16384},
)

All replies parsed.

New issues with 3B

Real reply to the better prompt, for a query whose gold document sat at position 1 of the shortlist:

shortlist: 20 docs | gold document at position 1

llama3.2 replied: { "ranking": [21, 20, 19, 18, 17, 16, 15, 14, 13, 12] }

21 does not even exist: the model counts down from N
The better prompt fixed the format, but not the judgment.

Results

llama3.2 3B   listwise rerank   nDCG@10 0.295  (5.1 s/query)
llama3.1 8B   listwise rerank   nDCG@10 0.557  (14.4 s/query)

The 8B produces genuine orderings,
but both models still damage the 0.693 retrieval they rerank.

Pointwise vs listwise

The listwise prompt asks the model to do arithmetic over list positions — something small transformers are not good at.

A much better approach is to judge one document at a time & read the answer from the model’s own distribution [4].

Score, do not generate

One small prompt per document, rather than a giant prompt per query:

resp = ollama.chat(
    model=model,
    messages=[
        {"role": "system", "content": "You are a search reranker."},
        {"role": "user", "content":
            f"Claim: {query}\n\nAbstract: {full_text[d]}\n\n"
            "Is this abstract relevant to the claim? "
            "Answer with exactly one word: yes or no."},
    ],
    logprobs=True, top_logprobs=10,
    options={"temperature": 0, "num_predict": 1, "num_ctx": 2048},
)

The score

lp = {e["token"]: e["logprob"]
      for e in resp["logprobs"][0]["top_logprobs"]}

yes = max(lp[t] for t in ("yes", "Yes", " yes"))
no  = max(lp[t] for t in ("no", "No", " no"))

score = np.exp(yes) / (np.exp(yes) + np.exp(no))  # P(yes)

Rank the shortlist by P(yes).

P(yes) straight from the model’s distribution: no list to count, no text to parse.

Pointwise results

llama3.2 3B   pointwise  nDCG@10 0.668  (3.3 s/query)
llama3.1 8B   pointwise  nDCG@10 0.722  (7.6 s/query)

The 8B now beats the 0.693 retriever it starts from.

Fine-tuning

The LLMs we used are general-purpose chat models. While extremely large models are great at any task without fine-tuning, small ones are not. It is not surprising that they perform poorly here.

Fine-tuning can teach these models to perform much better at a narrow skill.

Parameter-Efficient Fine-Tuning (PEFT) does this by only training a tiny fraction of the parameters while freezing the rest.

Low-Rank Adaptation (LoRA) [5] is a popular PEFT technique that trains tiny adapter matrices beside a frozen model.

Let’s train the 4-bit llama3.2 3B that fits the 6 GiB GPU:

  • 400 claims from the training split (the test split was never touched): the same yes/no prompt, gold abstracts & BM25 hard negatives,
  • Loss on the answer token only,
  • One epoch: 57 minutes on the laptop.

How LoRA works

LoRA Effective weights: W + B·A scaled by α/r = 2 24.3M of 3.2B parameters trainable (0.75%) x x W W 3072 × 3072 Frozen — never trained x->W A A 16 × 3072 Trainable x->A out output add + W->add B B 3072 × 16 Trainable A->B r = 16 B->add add->out

The code

bnb = BitsAndBytesConfig(load_in_4bit=True,
                         bnb_4bit_quant_type="nf4")
model = AutoModelForCausalLM.from_pretrained(
    "unsloth/Llama-3.2-3B-Instruct",
    quantization_config=bnb, device_map={"": 0})
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"]))

Results

llama3.2 3B + LoRA   pointwise  nDCG@10 0.739  (13.3 s/query)

Same 3B model:

Generating a ranking: 0.295 → Scoring: 0.668 → Fine-tuned: 0.739.

The judgment is calibrated too: gold at P(yes) 0.76 vs 0.47 for the next abstract
The untuned model tied them at 0.82.

Verdict

Full comparison

All the LLM methods run on 50 queries; the retrievers re-scored on the same 50:

Method nDCG@10 Per query
llama3.2 3B + LoRA pointwise 0.739 13.3 s
llama3.1 8B pointwise 0.722 7.6 s
nomic-embed-text (retriever) 0.706 18 ms
mxbai-embed-large (retriever) 0.693 22 ms
llama3.2 3B pointwise 0.668 3.3 s
RRF hybrid (mxbai + BM25) 0.651 —
all-minilm (retriever) 0.622 20 ms
BM25 (retriever) 0.606 15 ms
llama3.1 8B listwise 0.557 14.4 s
llama3.2 3B listwise 0.295 5.1 s

Fine-tuned 3B is now the best method.

Best per era

Era Method nDCG@10 Per query
2 BM25 0.606 15 ms
3 mxbai-embed-large (335M) 0.693 22 ms
4 llama3.2 3B + LoRA pointwise 0.739 13.3 s

All on the same 50 queries, same benchmark, same hardware: a laptop with a 6 GiB GPU.

The LLM cost

Embedding search LLM reranking
Per query 22 ms 3.3 s to 14.4 s
Slowdown 1× 150× to 650×
Model on disk 669 MB 2.0 to 4.9 GB

Only the 8B pointwise (7.6 s) & the fine-tuned 3B (13.3 s) beat the embedder.
Fine-tuning of 3B took one hour of training on the laptop.

Takeaways

  • Embeddings remain a practical and efficient choice for semantic search, and they scale gracefully: bigger models give better rankings at essentially the same speed,
  • Hybrid fusion patches weak retrieval, but drags down strong retrieval,
  • A prompted generative model is not automatically better: asking it to generate a ranking was the worst method,
  • Used right (scoring with its probabilities instead of its prose), the 8B LLM beat the retriever it reranked on the same queries, but 345× slower,
  • One hour of LoRA on the training split turned the 3B from the worst to the best method.

Reach for the LLM when

The task is generative:

  • Summarize, answer, explain, discuss,
  • Extract claims & entities.

Retrieval and generation are different jobs. Let the embedder retrieve and the LLM do what only an LLM can do.

Or score with it when ranking quality matters more than speed.

Full code

"""Semantic search across the data-driven eras of NLP: lexical, dense, generative.

Experiment backing the webinar "From pipelines to prompts: NLP in the LLM era"
(see wb_llm_era.qmd). One task - semantic search - on one benchmark:

    BEIR/SciFact: 5,183 scientific abstracts, ~300 test queries
    (scientific claims) with gold relevance judgements.

The eras (numbered as in the webinar slides, where era 1 is the symbolic
era, not tested here as it involves no data-driven method):

2. Lexical search   BM25 (no model, exact-term matching)
3. Dense retrieval  embedding similarity, three model sizes
4. Generative       LLMs (llama3.2 3B, llama3.1 8B via ollama) reranking
                    a shortlist, two ways: the brute-force listwise prompt
                    and pointwise P(yes) scoring read from the model
                    logprobs

Plus the standard hybrid trick (reciprocal rank fusion of eras 2 and 3).

Everything is scored with nDCG@10 on the same qrels. No method is given an
unfair edge: the numbers land where they land and the webinar conclusion
follows the results, not the other way around.

Data:   curl -L -o scifact.zip
        https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/scifact.zip
Models: ollama pull all-minilm nomic-embed-text mxbai-embed-large llama3.2 llama3.1:8b

Run with `uv run python -I src/experiment.py` (SMOKE=1 for a quick validation
run on a corpus subsample).
"""

import json
import os
import re
import time
from pathlib import Path

import numpy as np
import ollama
import polars as pl
from rank_bm25 import BM25Okapi

# --- Config ------------------------------------------------------------------

SMOKE = os.environ.get("SMOKE") == "1"  # tiny corpus, few queries, fast checks
HERE = Path(__file__).resolve().parent.parent  # scripts live in src/, data one level up
DATA = HERE / "data/scifact/scifact"
SRC = HERE / "src"  # scratch artifacts: .emb_cache/, lora_reranker/
K = 10  # nDCG depth
LLM_QUERIES = 5 if SMOKE else 50  # LLM reranking is slow: subset only
RRF_K = 60  # reciprocal rank fusion constant

# Embedding models, oldest/smallest first: does model quality flip the
# BM25 vs dense ranking? (SciFact is full of rare exact terms - a BM25-friendly
# domain - so this is a fair fight, not a strawman.)
EMB_MODELS = [
    "all-minilm",  # 23M params, 2020-era miniature
    "nomic-embed-text",  # 137M params, modern
    "mxbai-embed-large",  # 335M params, modern
]
# mxbai wants a task instruction prefixed to queries (per its model card);
# nomic's ollama template already adds its own prefixes.
QUERY_PREFIX = {
    "mxbai-embed-large": "Represent this sentence for searching relevant passages: "
}
LLM_MODELS = ["llama3.2", "llama3.1:8b"]  # 3B then 8B: does scaling fix era 4?
# Reranking full abstracts (median ~360 tokens, p99 ~790) means ~8-16k token
# prompts, far past ollama's 4k default context. Weights + 16k of KV cache only
# fit the 6 GiB GPU for the 8B model with flash attention and a compressed cache.
LLM_OPTS = {
    "llama3.2": {"num_ctx": 16384},
    "llama3.1:8b": {
        "num_ctx": 16384,
        "flash_attention": True,
        "kv_cache_type": "q8_0",
    },
}

RERANK_FORMAT = {  # structured output for the reranker
    "type": "object",
    "properties": {"ranking": {"type": "array", "items": {"type": "integer"}}},
    "required": ["ranking"],
}

RESULTS = []


def tokenize(text: str) -> list[str]:
    """Crude but sufficient for BM25: lowercase alphanumeric tokens."""
    return re.findall(r"[a-z0-9]+", text.lower())


def ndcg_at_k(ranked: list[str], qrels: dict[str, int], k: int = K) -> float:
    """nDCG@k, graded gains (SciFact qrels happen to be binary)."""
    dcg = sum(qrels.get(d, 0) / np.log2(r + 2) for r, d in enumerate(ranked[:k]))
    ideal = sorted(qrels.values(), reverse=True)[:k]
    idcg = sum(g / np.log2(r + 2) for r, g in enumerate(ideal))
    return dcg / idcg if idcg else 0.0


# --- Corpus, queries, relevance judgements -------------------------------------

corpus = {}  # doc id -> {"title": ..., "text": ...}
with open(DATA / "corpus.jsonl", encoding="utf-8") as f:
    for line in f:
        doc = json.loads(line)
        corpus[doc["_id"]] = {
            "title": doc.get("title") or "",
            "text": doc["text"],
        }

if SMOKE:  # shrink the corpus, keep only queries still answerable from it
    corpus = dict(list(corpus.items())[:500])
docs = list(corpus)

queries = {}  # query id -> text
with open(DATA / "queries.jsonl", encoding="utf-8") as f:
    for line in f:
        q = json.loads(line)
        queries[q["_id"]] = q["text"]

qrels = {}  # query id -> {doc id: relevance}
with open(DATA / "qrels" / "test.tsv", encoding="utf-8") as f:
    next(f)
    for line in f:
        qid, did, score = line.split("\t")
        qrels.setdefault(qid, {})[did.strip()] = int(score)

eval_qids = [
    qid for qid in qrels if qid in queries and any(d in corpus for d in qrels[qid])
]
if SMOKE:
    eval_qids = eval_qids[:30]
print(
    f"corpus: {len(corpus)} docs | eval queries: {len(eval_qids)}"
    f" | LLM subset: {LLM_QUERIES}\n"
)

full_text = {d: f"{corpus[d]['title']}. {corpus[d]['text']}" for d in docs}


# --- Era 2: lexical search (BM25) ----------------------------------------------

# No model at all: score docs by exact (unigram) term overlap with the query.
# Fails on vocabulary mismatch - and on SciFact it is still a serious
# baseline: scientific queries are full of rare exact terms.

t0 = time.perf_counter()
bm25 = BM25Okapi([tokenize(full_text[d]) for d in docs])
bm25_index_time = time.perf_counter() - t0

t0 = time.perf_counter()
bm25_rankings = {  # query id -> doc ids, best first
    qid: [docs[i] for i in np.argsort(-bm25.get_scores(tokenize(queries[qid])))]
    for qid in eval_qids
}
bm25_search_time = time.perf_counter() - t0

bm25_ndcg = float(np.mean([ndcg_at_k(bm25_rankings[q], qrels[q]) for q in eval_qids]))
RESULTS.append(
    {
        "method": "2 BM25 (lexical)",
        "ndcg10": bm25_ndcg,
        "index_s": bm25_index_time,
        "search_ms": 1000 * bm25_search_time / len(eval_qids),
    }
)
print(
    f"era 2  BM25             nDCG@10 {bm25_ndcg:.3f}"
    f"  ({bm25_search_time * 1000 / len(eval_qids):.1f} ms/query)"
)


# --- Era 3: dense retrieval (embeddings) ---------------------------------------

CACHE = SRC / ".emb_cache"


def embed(texts: list[str], model: str, batch: int = 64) -> np.ndarray:
    """Embed texts with an ollama model, L2-normalized (disk-cached:
    the doc list is deterministic, so length is a sufficient cache key)."""
    cache = CACHE / f"{model}-{len(texts)}.npy"
    if cache.exists():
        return np.load(cache)
    CACHE.mkdir(exist_ok=True)
    out = []
    for i in range(0, len(texts), batch):
        resp = ollama.embed(model=model, input=texts[i : i + batch])
        out.extend(resp["embeddings"])
        if (i // batch) % 10 == 0:
            print(f"  {model}: {min(i + batch, len(texts))}/{len(texts)}", flush=True)
    emb = np.asarray(out, dtype=np.float32)
    emb /= np.linalg.norm(emb, axis=1, keepdims=True)
    np.save(cache, emb)
    return emb


best_dense = None  # (ndcg, model, rankings) of the era 3 winner
for model in EMB_MODELS:
    t0 = time.perf_counter()
    doc_emb = embed([full_text[d] for d in docs], model)
    index_time = time.perf_counter() - t0

    t0 = time.perf_counter()
    rankings = {}
    for qid in eval_qids:
        q = ollama.embed(
            model=model, input=[QUERY_PREFIX.get(model, "") + queries[qid]]
        )["embeddings"][0]
        q = np.asarray(q, dtype=np.float32)
        q /= np.linalg.norm(q)
        rankings[qid] = [docs[i] for i in np.argsort(-(doc_emb @ q))]
    search_time = time.perf_counter() - t0

    ndcg = float(np.mean([ndcg_at_k(rankings[q], qrels[q]) for q in eval_qids]))
    RESULTS.append(
        {
            "method": f"3 {model}",
            "ndcg10": ndcg,
            "index_s": index_time,
            "search_ms": 1000 * search_time / len(eval_qids),
        }
    )
    print(
        f"era 3  {model:<17} nDCG@10 {ndcg:.3f}"
        f"  ({index_time:.0f}s index, "
        f"{1000 * search_time / len(eval_qids):.1f} ms/query)"
    )
    if best_dense is None or ndcg > best_dense[0]:
        best_dense = (ndcg, model, rankings)

_, dense_model, dense_rankings = best_dense


# --- Hybrid: reciprocal rank fusion of eras 2 and 3 -----------------------------


def rrf(rankings: dict[str, list[str]]) -> dict[str, list[str]]:
    """Fuse rankings: score = sum over methods of 1 / (RRF_K + rank)."""
    fused = {}
    for qid in eval_qids:
        scores = {}
        for ranking in rankings.values():
            for rank, d in enumerate(ranking[qid]):
                scores[d] = scores.get(d, 0.0) + 1.0 / (RRF_K + rank)
        fused[qid] = sorted(scores, key=scores.get, reverse=True)
    return fused


rrf_rankings = rrf({"bm25": bm25_rankings, "dense": dense_rankings})
rrf_ndcg = float(np.mean([ndcg_at_k(rrf_rankings[q], qrels[q]) for q in eval_qids]))
RESULTS.append(
    {
        "method": f"3+2 hybrid RRF ({dense_model} + BM25)",
        "ndcg10": rrf_ndcg,
        "index_s": np.nan,
        "search_ms": np.nan,
    }
)
print(f"hyb   RRF fusion       nDCG@10 {rrf_ndcg:.3f}")


# --- Era 4: LLM rerankers ---------------------------------------------------------

# A generative model cannot scan the corpus; it can only reorder a shortlist.
# Standard pattern: retrieve top-k with era 2/3 methods, let the LLM rerank.
# This is the only place compute explodes: one LLM call per query.


def rerank(model: str, query: str, shortlist: list[str]) -> tuple[list[str], bool]:
    listing = "\n".join(f"{n + 1}. {full_text[d]}" for n, d in enumerate(shortlist))
    resp = ollama.chat(
        model=model,
        messages=[
            {"role": "system", "content": "You are a search reranker."},
            {
                "role": "user",
                "content": f"Query: {query}\n\nDocuments:\n{listing}\n\n"
                f"There are {len(shortlist)} documents, numbered 1 to {len(shortlist)}. "
                f"Reply with the {K} most relevant document numbers, "
                "most relevant first.",
            },
        ],
        format=RERANK_FORMAT,
        options={"temperature": 0, "num_predict": 200, **LLM_OPTS[model]},
    )
    ok = True
    try:
        nums = list(
            dict.fromkeys(  # dedupe, order-preserving
                json.loads(resp["message"]["content"])["ranking"]
            )
        )[:K]
        ranked = [shortlist[n - 1] for n in nums if 1 <= n <= len(shortlist)]
    except (json.JSONDecodeError, KeyError, IndexError, TypeError) as e:
        ok = False
        ranked = []
        with open(
            HERE / "results" / ".rerank_failures.log", "a", encoding="utf-8"
        ) as f:
            f.write(
                f"query: {query!r}\nerror: {e}\n"
                f"response: {resp['message']['content']!r}\n\n"
            )
    # pad with the retriever order (docs the LLM did not list)
    return ranked + [d for d in shortlist if d not in ranked], ok


# Unload the embedding models from the ollama server so the LLMs have
# room to load (they compete for the 6 GiB of GPU memory).
for m in EMB_MODELS:
    ollama.embed(model=m, input=["."], keep_alive=0)

llm_qids = sorted(eval_qids)[:LLM_QUERIES]
shortlists = {  # same shortlists for every LLM: only the reranker changes
    qid: list(dict.fromkeys(bm25_rankings[qid][:K] + dense_rankings[qid][:K]))
    for qid in llm_qids
}
for llm in LLM_MODELS:
    t0 = time.perf_counter()
    n_fail = 0
    llm_rankings = {}
    for qid in llm_qids:
        for attempt in range(3):
            try:
                llm_rankings[qid], ok = rerank(llm, queries[qid], shortlists[qid])
                break
            except ollama.ResponseError:  # runner hiccup: give it a moment
                print(f"  rerank retry {attempt + 1} for query {qid}", flush=True)
                time.sleep(10)
        else:  # all retries failed: fall back to the retriever order
            llm_rankings[qid], ok = shortlists[qid][:K], False
        n_fail += not ok
    llm_time = time.perf_counter() - t0
    ollama.generate(model=llm, keep_alive=0)  # free memory for the next model

    llm_ndcg = float(
        np.mean([ndcg_at_k(llm_rankings[q], qrels[q]) for q in llm_qids])
    )
    RESULTS.append(
        {
            "method": f"4 {llm} listwise rerank on {dense_model}+BM25 top-{K}",
            "ndcg10": llm_ndcg,
            "index_s": np.nan,
            "search_ms": 1000 * llm_time / len(llm_qids),
        }
    )
    print(
        f"era 4 {llm:<15} listwise rerank nDCG@10 {llm_ndcg:.3f}"
        f"  ({llm_time / len(llm_qids):.1f} s/query,"
        f" {len(llm_qids)} queries, {n_fail} unparseable)"
    )


# --- Era 4, take 2: pointwise scoring ---------------------------------------------

# The listwise prompt asks the model to do arithmetic over list positions -
# exactly what small transformers fumble (counting loops, countdown replies).
# Alternative: judge one document at a time and read the answer from the
# model's own distribution: score = P(yes) / (P(yes) + P(no)) on the first
# generated token. No list to count, no JSON to truncate, no fallback. And
# one abstract per prompt fits a small context window (no 16k KV cache).

YES = ("yes", "Yes", " yes", " Yes")
NO = ("no", "No", " no", " No")


def p_yes(resp) -> float:
    """P(yes) / (P(yes) + P(no)) from the first generated token's logprobs."""
    lp = {e["token"]: e["logprob"] for e in resp["logprobs"][0]["top_logprobs"]}
    yes = max((lp[t] for t in YES if t in lp), default=None)
    no = max((lp[t] for t in NO if t in lp), default=None)
    if yes is None and no is None:
        return 0.5  # the model dodged the question: no signal
    if no is None:
        return 1.0
    if yes is None:
        return 0.0
    ey, en = np.exp(yes), np.exp(no)
    return float(ey / (ey + en))


for llm in LLM_MODELS:
    t0 = time.perf_counter()
    llm_rankings = {}
    for qid in llm_qids:
        for attempt in range(3):
            try:
                scores = []
                for d in shortlists[qid]:
                    resp = ollama.chat(
                        model=llm,
                        messages=[
                            {
                                "role": "system",
                                "content": "You are a search reranker.",
                            },
                            {
                                "role": "user",
                                "content": f"Claim: {queries[qid]}\n\n"
                                f"Abstract: {full_text[d]}\n\n"
                                "Is this abstract relevant to the claim? "
                                "Answer with exactly one word: yes or no.",
                            },
                        ],
                        logprobs=True,
                        top_logprobs=10,
                        options={"temperature": 0, "num_predict": 1,
                                 "num_ctx": 2048},
                    )
                    scores.append(p_yes(resp))
                break
            except ollama.ResponseError:  # runner hiccup: give it a moment
                print(f"  pointwise retry {attempt + 1} for query {qid}",
                      flush=True)
                time.sleep(10)
        else:  # all retries failed: neutral scores keep the retriever order
            scores = [0.5] * len(shortlists[qid])
        # stable sort: ties keep the retriever order
        order = np.argsort(-np.asarray(scores), kind="stable")
        llm_rankings[qid] = [shortlists[qid][i] for i in order]
    llm_time = time.perf_counter() - t0
    ollama.generate(model=llm, keep_alive=0)  # free memory for the next model

    llm_ndcg = float(
        np.mean([ndcg_at_k(llm_rankings[q], qrels[q]) for q in llm_qids])
    )
    RESULTS.append(
        {
            "method": f"4 {llm} pointwise rerank on {dense_model}+BM25 top-{K}",
            "ndcg10": llm_ndcg,
            "index_s": np.nan,
            "search_ms": 1000 * llm_time / len(llm_qids),
        }
    )
    print(
        f"era 4 {llm:<15} pointwise rerank nDCG@10 {llm_ndcg:.3f}"
        f"  ({llm_time / len(llm_qids):.1f} s/query,"
        f" {len(llm_qids)} queries)"
    )


# --- Summary ---------------------------------------------------------------------

df = pl.DataFrame(RESULTS)
print("\n== BEIR/SciFact - nDCG@10 ==")
print(
    str(
        df.with_columns(
            pl.col("ndcg10").map_elements(lambda x: f"{x:.3f}", return_dtype=pl.String),
            pl.col("index_s").map_elements(
                lambda x: f"{x:.0f}", return_dtype=pl.String
            ),
            pl.col("search_ms").map_elements(
                lambda x: f"{x:.1f}", return_dtype=pl.String
            ),
        )
    )
)

df.write_json(HERE / "results" / "results.json")
"""nDCG@10 of the retrieval methods restricted to the 50-query LLM subset,
so era 4 is compared apples to apples."""

import json
import re
from pathlib import Path

import numpy as np
import ollama
from rank_bm25 import BM25Okapi

HERE = Path(__file__).resolve().parent.parent  # scripts live in src/, data one level up
DATA = HERE / "data/scifact/scifact"
SRC = HERE / "src"  # scratch artifacts: .emb_cache/, lora_reranker/
CACHE = SRC / ".emb_cache"
K, RRF_K = 10, 60
QUERY_PREFIX = {
    "mxbai-embed-large": "Represent this sentence for searching relevant passages: "
}


def ndcg_at_k(ranked, qrels, k=K):
    dcg = sum(qrels.get(d, 0) / np.log2(r + 2) for r, d in enumerate(ranked[:k]))
    ideal = sorted(qrels.values(), reverse=True)[:k]
    idcg = sum(g / np.log2(r + 2) for r, g in enumerate(ideal))
    return dcg / idcg if idcg else 0.0


corpus = {}
with open(DATA / "corpus.jsonl", encoding="utf-8") as f:
    for line in f:
        doc = json.loads(line)
        corpus[doc["_id"]] = {"title": doc.get("title") or "", "text": doc["text"]}
docs = list(corpus)
queries, qrels = {}, {}
with open(DATA / "queries.jsonl", encoding="utf-8") as f:
    for line in f:
        q = json.loads(line)
        queries[q["_id"]] = q["text"]
with open(DATA / "qrels" / "test.tsv", encoding="utf-8") as f:
    next(f)
    for line in f:
        qid, did, score = line.split("\t")
        qrels.setdefault(qid, {})[did.strip()] = int(score)

eval_qids = [q for q in qrels if q in queries and any(d in corpus for d in qrels[q])]
llm_qids = sorted(eval_qids)[:50]  # same subset as experiment.py era 4

# era-4 settings from the full run: every embedder + the reranker scores
rows = json.loads((HERE / "results" / "results.json").read_text())
emb_models = [r["method"].split()[1] for r in rows
              if "rerank" not in r["method"] and "BM25" not in r["method"]
              and "hybrid" not in r["method"]]

full_text = {d: f"{corpus[d]['title']}. {corpus[d]['text']}" for d in docs}
tok = lambda s: re.findall(r"[a-z0-9]+", s.lower())
bm25 = BM25Okapi([tok(full_text[d]) for d in docs])


def mean_ndcg(rankings):
    return float(np.mean([ndcg_at_k(rankings[q], qrels[q]) for q in llm_qids]))


bm25_r = {
    qid: [docs[i] for i in np.argsort(-bm25.get_scores(tok(queries[qid])))]
    for qid in llm_qids
}
dense_r = {}
for model in emb_models:
    E = np.load(CACHE / f"{model}-{len(docs)}.npy")  # doc order == docs
    dense_r[model] = {}
    for qid in llm_qids:
        q = np.asarray(
            ollama.embed(
                model=model,
                input=[QUERY_PREFIX.get(model, "") + queries[qid]],
            )["embeddings"][0],
            dtype=np.float32,
        )
        q /= np.linalg.norm(q)
        dense_r[model][qid] = [docs[i] for i in np.argsort(-(E @ q))]

# hybrid partner: the era-3 winner of the full run (not of this subset)
full_ndcg = {r["method"].split()[1]: r["ndcg10"] for r in rows
             if r["method"].split()[1] in emb_models}
best = max(emb_models, key=lambda m: full_ndcg[m])

rrf_r = {}
for qid in llm_qids:
    s = {}
    for ranking in (bm25_r[qid], dense_r[best][qid]):
        for rank, d in enumerate(ranking):
            s[d] = s.get(d, 0.0) + 1.0 / (RRF_K + rank)
    rrf_r[qid] = sorted(s, key=s.get, reverse=True)

print(f"on the 50-query LLM subset (nDCG@10):")
print(f"  {'BM25':<18} {mean_ndcg(bm25_r):.3f}")
for model in emb_models:
    print(f"  {model:<18} {mean_ndcg(dense_r[model]):.3f}")
print(f"  {'RRF hybrid':<18} {mean_ndcg(rrf_r):.3f}")
for r in rows:  # rerankers were scored on this very subset
    if "rerank" in r["method"]:
        print(f"  {r['method'].split(' ', 1)[1]:<42} {r['ndcg10']:.3f}"
              "  (from experiment.py)")
"""Pointwise LLM reranking: score each shortlisted document with the model's
own P(yes) instead of asking it to generate a ranking.

Companion to experiment.py era 4 (listwise reranking): same 50 queries, same
shortlists (BM25 + dense top-10 union), same qrels. Instead of one giant
"rank these documents" prompt (16k-token context, JSON contract), each
document gets its own small prompt, and the score is read from the first
generated token's logprobs:

    score = P(yes) / (P(yes) + P(no))

Nothing to count, no JSON to truncate, no fallback needed.

Run with `uv run python -I src/pointwise_rerank.py` (ollama serve first).
POINTWISE_QUERIES=3 for a quick sanity check.
"""

import json
import os
import re
import time
from pathlib import Path

import numpy as np
import ollama
from rank_bm25 import BM25Okapi

HERE = Path(__file__).resolve().parent.parent  # scripts live in src/, data one level up
DATA = HERE / "data/scifact/scifact"
SRC = HERE / "src"  # scratch artifacts: .emb_cache/, lora_reranker/
CACHE = SRC / ".emb_cache"
K = 10
N_QUERIES = int(os.environ.get("POINTWISE_QUERIES", "50"))
QUERY_PREFIX = {
    "mxbai-embed-large": "Represent this sentence for searching relevant passages: "
}
LLM_MODELS = ["llama3.2", "llama3.1:8b"]
# one abstract per prompt: a small context window is enough (and cheap)
LLM_OPTS = {"temperature": 0, "num_predict": 1, "num_ctx": 2048}
YES = ("yes", "Yes", " yes", " Yes")
NO = ("no", "No", " no", " No")


def ndcg_at_k(ranked, qrels, k=K):
    dcg = sum(qrels.get(d, 0) / np.log2(r + 2) for r, d in enumerate(ranked[:k]))
    ideal = sorted(qrels.values(), reverse=True)[:k]
    idcg = sum(g / np.log2(r + 2) for r, g in enumerate(ideal))
    return dcg / idcg if idcg else 0.0


def p_yes(resp) -> float:
    """P(yes)/(P(yes)+P(no)) from the first generated token's top logprobs."""
    lp = {e["token"]: e["logprob"]
          for e in resp["logprobs"][0]["top_logprobs"]}
    yes = max((lp[t] for t in YES if t in lp), default=None)
    no = max((lp[t] for t in NO if t in lp), default=None)
    if yes is None and no is None:
        return 0.5  # the model dodged the question: no signal
    if no is None:
        return 1.0
    if yes is None:
        return 0.0
    ey, en = np.exp(yes), np.exp(no)
    return float(ey / (ey + en))


# --- Same data, same shortlists as experiment.py ---------------------------------

corpus = {}
with open(DATA / "corpus.jsonl", encoding="utf-8") as f:
    for line in f:
        doc = json.loads(line)
        corpus[doc["_id"]] = {"title": doc.get("title") or "", "text": doc["text"]}
docs = list(corpus)
queries, qrels = {}, {}
with open(DATA / "queries.jsonl", encoding="utf-8") as f:
    for line in f:
        q = json.loads(line)
        queries[q["_id"]] = q["text"]
with open(DATA / "qrels" / "test.tsv", encoding="utf-8") as f:
    next(f)
    for line in f:
        qid, did, score = line.split("\t")
        qrels.setdefault(qid, {})[did.strip()] = int(score)

eval_qids = [q for q in qrels if q in queries and any(d in corpus for d in qrels[q])]
llm_qids = sorted(eval_qids)[:N_QUERIES]
print(f"corpus: {len(corpus)} docs | pointwise queries: {len(llm_qids)}\n", flush=True)

full_text = {d: f"{corpus[d]['title']}. {corpus[d]['text']}" for d in docs}
tok = lambda s: re.findall(r"[a-z0-9]+", s.lower())
bm25 = BM25Okapi([tok(full_text[d]) for d in docs])

dense_model = max(
    (r for r in json.loads((HERE / "results" / "results.json").read_text())
     if "rerank" not in r["method"] and "BM25" not in r["method"]
     and "hybrid" not in r["method"]),
    key=lambda r: r["ndcg10"],
)["method"].split()[1]
E = np.load(CACHE / f"{dense_model}-{len(docs)}.npy")  # doc order == docs

shortlists = {}
for qid in llm_qids:
    bm25_top = [docs[i] for i in np.argsort(-bm25.get_scores(tok(queries[qid])))][:K]
    q = np.asarray(
        ollama.embed(model=dense_model,
                     input=[QUERY_PREFIX.get(dense_model, "") + queries[qid]])
        ["embeddings"][0], dtype=np.float32)
    q /= np.linalg.norm(q)
    dense_top = [docs[i] for i in np.argsort(-(E @ q))][:K]
    shortlists[qid] = list(dict.fromkeys(bm25_top + dense_top))

for m in (dense_model,):  # free the GPU for the LLMs
    ollama.embed(model=m, input=["."], keep_alive=0)


# --- Era 4, take 2: score, don't generate ---------------------------------------

def pointwise(model: str, query: str, shortlist: list[str]) -> list[str]:
    """Rank the shortlist by the model's P(yes), ties keep retriever order."""
    scores = []
    for d in shortlist:
        resp = ollama.chat(
            model=model,
            messages=[
                {"role": "system", "content": "You are a search reranker."},
                {
                    "role": "user",
                    "content": f"Claim: {query}\n\nAbstract: {full_text[d]}\n\n"
                    "Is this abstract relevant to the claim? "
                    "Answer with exactly one word: yes or no.",
                },
            ],
            logprobs=True,
            top_logprobs=10,
            options=LLM_OPTS,
        )
        scores.append(p_yes(resp))
    order = np.argsort(-np.asarray(scores), kind="stable")
    return [shortlist[i] for i in order], scores


RESULTS = []
for llm in LLM_MODELS:
    t0 = time.perf_counter()
    ndcg, dodged = [], 0
    for n, qid in enumerate(llm_qids):
        for attempt in range(3):
            try:
                ranked, scores = pointwise(llm, queries[qid], shortlists[qid])
                break
            except ollama.ResponseError:  # runner hiccup: give it a moment
                print(f"  retry {attempt + 1} for query {qid}", flush=True)
                time.sleep(10)
        else:  # all retries failed: fall back to the retriever order
            ranked = shortlists[qid]
        dodged += sum(s == 0.5 for s in scores)
        ndcg.append(ndcg_at_k(ranked, qrels[qid]))
        if n < 3:  # exhibit material: gold vs distractors for the first queries
            gold = [d for d in qrels[qid] if d in shortlists[qid]]
            print(f"  query {n + 1}: {queries[qid][:60]!r}", flush=True)
            for d, s in sorted(zip(shortlists[qid], scores), key=lambda x: -x[1])[:3]:
                mark = " <- gold" if d in gold else ""
                print(f"    P(yes) {s:.2f}  {full_text[d][:60]!r}{mark}", flush=True)
    llm_time = time.perf_counter() - t0
    ollama.generate(model=llm, keep_alive=0)  # free memory for the next model

    mean_ndcg = float(np.mean(ndcg))
    RESULTS.append(
        {
            "method": f"4 {llm} pointwise P(yes) on {dense_model}+BM25 top-{K}",
            "ndcg10": mean_ndcg,
            "search_ms": 1000 * llm_time / len(llm_qids),
            "dodged": dodged,
        }
    )
    print(
        f"era 4 {llm:<15} pointwise nDCG@10 {mean_ndcg:.3f}"
        f"  ({llm_time / len(llm_qids):.1f} s/query,"
        f" {dodged} dodged answers)\n",
        flush=True,
    )

Path(HERE / "results" / "run05_results_pointwise.json").write_text(json.dumps(RESULTS, indent=2))
"""Fine-tune llama3.2 3B into a pointwise reranker with QLoRA.

Era 4, take 3: the model that failed the listwise prompt (0.295) and scored
0.668 with prompted P(yes) is now trained for that same yes/no judgment on
the SciFact training split. The test split is never used for training.

Training: (claim, abstract) pairs -> "yes"/"no", loss on the answer token
only, QLoRA (4-bit base + LoRA adapters). Hard negatives: top BM25 hits that
are not gold. Evaluation: the same 50 test queries and shortlists as
run04/run05 (BM25 + mxbai top-10 union), scored with the same P(yes) -
read from the fine-tuned model's first answer token.

Needs ollama running (50 query embeddings) and downloads the base model
from Hugging Face. Run with `uv run python -I src/lora_rerank.py`
(LORA_SMOKE=1 for a quick pipeline check on a tiny slice).
"""

import json
import os
import re
import time
from pathlib import Path

import numpy as np
import ollama
import torch
import torch.nn.functional as F
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from rank_bm25 import BM25Okapi
from torch.utils.data import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    Trainer,
    TrainingArguments,
)

HERE = Path(__file__).resolve().parent.parent  # scripts live in src/, data one level up
DATA = HERE / "data/scifact/scifact"
SRC = HERE / "src"  # scratch artifacts: .emb_cache/, lora_reranker/
CACHE = SRC / ".emb_cache"
K = 10
SMOKE = os.environ.get("LORA_SMOKE") == "1"
BASE = "unsloth/Llama-3.2-3B-Instruct"  # ungated mirror of llama3.2 3B
SYSTEM = "You are a search reranker."
N_NEG = 2  # hard negatives per gold document
MAX_CHARS = 1500  # abstract truncation: ~400-token prompts fit the time budget
TRAIN_CLAIMS = 400  # subsample of the ~900 train claims (laptop-GPU budget)
BATCH = 2  # eval batch: full-vocab logits for a batch are ~180 MB each


def ndcg_at_k(ranked, qrels, k=K):
    dcg = sum(qrels.get(d, 0) / np.log2(r + 2) for r, d in enumerate(ranked[:k]))
    ideal = sorted(qrels.values(), reverse=True)[:k]
    idcg = sum(g / np.log2(r + 2) for r, g in enumerate(ideal))
    return dcg / idcg if idcg else 0.0


# --- Data (same loaders as experiment.py) -----------------------------------------

corpus = {}
with open(DATA / "corpus.jsonl", encoding="utf-8") as f:
    for line in f:
        doc = json.loads(line)
        corpus[doc["_id"]] = {"title": doc.get("title") or "", "text": doc["text"]}
docs = list(corpus)
queries = {}
with open(DATA / "queries.jsonl", encoding="utf-8") as f:
    for line in f:
        q = json.loads(line)
        queries[q["_id"]] = q["text"]
qrels = {}  # split -> qid -> {doc id: relevance}
for split in ("train", "test"):
    qrels[split] = {}
    with open(DATA / "qrels" / f"{split}.tsv", encoding="utf-8") as f:
        next(f)
        for line in f:
            qid, did, score = line.split("\t")
            qrels[split].setdefault(qid, {})[did.strip()] = int(score)

full_text = {d: f"{corpus[d]['title']}. {corpus[d]['text']}" for d in docs}
doc_text = {d: full_text[d][:MAX_CHARS] for d in docs}  # for LLM prompts only
tok_text = lambda s: re.findall(r"[a-z0-9]+", s.lower())
bm25 = BM25Okapi([tok_text(full_text[d]) for d in docs])

# --- The 50 evaluation queries & shortlists (identical to run04/run05) -------------

eval_qids = [
    q for q in qrels["test"]
    if q in queries and any(d in corpus for d in qrels["test"][q])
]
llm_qids = sorted(eval_qids)[:50]
E = np.load(CACHE / f"mxbai-embed-large-{len(docs)}.npy")  # doc order == docs
shortlists = {}
for qid in llm_qids:
    bm25_top = [docs[i] for i in
                np.argsort(-bm25.get_scores(tok_text(queries[qid])))][:K]
    q = np.asarray(
        ollama.embed(
            model="mxbai-embed-large",
            input=["Represent this sentence for searching relevant passages: "
                   + queries[qid]],
        )["embeddings"][0], dtype=np.float32)
    q /= np.linalg.norm(q)
    shortlists[qid] = list(dict.fromkeys(
        bm25_top + [docs[i] for i in np.argsort(-(E @ q))][:K]))
ollama.embed(model="mxbai-embed-large", input=["."], keep_alive=0)  # free GPU


# --- Training examples: the pointwise prompt, answered from the train qrels --------

def user_msg(query, doc):
    return (f"Claim: {query}\n\nAbstract: {doc_text[doc]}\n\n"
            "Is this abstract relevant to the claim? "
            "Answer with exactly one word: yes or no.")


train_pairs = []  # (qid, doc, label)
train_qids = [
    q for q in qrels["train"]
    if q in queries and any(d in corpus for d in qrels["train"][q])
]
rng = np.random.default_rng(42)
train_qids = list(rng.permutation(train_qids)[:TRAIN_CLAIMS])
for qid in train_qids:
    gold = {d for d in qrels["train"][qid] if d in corpus}
    for d in gold:
        train_pairs.append((qid, d, 1))
    ranked = [docs[i] for i in np.argsort(-bm25.get_scores(tok_text(queries[qid])))]
    negatives = [d for d in ranked if d not in gold][: N_NEG * max(len(gold), 1)]
    for d in negatives:
        train_pairs.append((qid, d, 0))
print(f"training examples: {len(train_pairs)}"
      f" ({sum(l for _, _, l in train_pairs)} yes)", flush=True)

tok = AutoTokenizer.from_pretrained(BASE)
tok.padding_side = "right"
YES_IDS = [tok(s, add_special_tokens=False)["input_ids"][0]
           for s in (" yes", "Yes", "yes", " Yes")]
NO_IDS = [tok(s, add_special_tokens=False)["input_ids"][0]
          for s in (" no", "No", "no", " No")]
assert not set(YES_IDS) & set(NO_IDS)
# train on "Yes"/"No": the base model refuses " yes" (untrained CE ~23 vs ~3)
yes_id, no_id = YES_IDS[1], NO_IDS[1]


class Pairs(Dataset):
    def __len__(self):
        return len(train_pairs)

    def __getitem__(self, i):
        qid, d, label = train_pairs[i]
        prompt = tok.apply_chat_template(
            [{"role": "system", "content": SYSTEM},
             {"role": "user", "content": user_msg(queries[qid], d)}],
            add_generation_prompt=True, return_dict=False,
        )
        ids = prompt + [yes_id if label else no_id]
        return {"input_ids": ids, "labels": [-100] * (len(ids) - 1) + [ids[-1]]}


def collate(batch):
    pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
    maxlen = max(len(b["input_ids"]) for b in batch)
    return {
        "input_ids": torch.tensor(
            [b["input_ids"] + [pad] * (maxlen - len(b["input_ids"]))
             for b in batch]),
        "labels": torch.tensor(
            [b["labels"] + [-100] * (maxlen - len(b["labels"])) for b in batch]),
        "attention_mask": torch.tensor(
            [[1] * len(b["input_ids"]) + [0] * (maxlen - len(b["input_ids"]))
             for b in batch]),
    }


# --- QLoRA -------------------------------------------------------------------------

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float32,  # fp16 compute: nan gradients
)
model = AutoModelForCausalLM.from_pretrained(
    BASE, quantization_config=bnb, device_map={"": 0},
    dtype=torch.float16, attn_implementation="sdpa",
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
))
model.print_trainable_parameters()

with torch.no_grad():  # sanity: untrained answer-token loss should be ~3
    b = collate([Pairs()[i] for i in range(2)])
    labels = b["labels"].cuda()
    logits = model(input_ids=b["input_ids"].cuda(),
                   attention_mask=b["attention_mask"].cuda()).logits
    idx = (labels != -100).int().argmax(dim=1)
    rows = torch.arange(len(idx), device=logits.device)
    pred = logits[rows, idx - 1].float()
    print(f"sanity: logits [{float(pred.min()):.1f}, {float(pred.max()):.1f}]"
          f" | untrained CE {float(F.cross_entropy(pred, labels[rows, idx])):.2f}",
          flush=True)

class AnswerTrainer(Trainer):
    """Cross-entropy on the answer token only. The default causal-LM loss
    materializes full-vocab fp32 logits for every position - more than the
    6 GiB GPU can hold next to the model."""

    def compute_loss(self, model, inputs, return_outputs=False,
                     num_items_in_batch=None):
        labels = inputs.pop("labels")
        logits = model(**inputs).logits
        idx = (labels != -100).int().argmax(dim=1)  # answer token position
        rows = torch.arange(len(idx), device=logits.device)
        pred = logits[rows, idx - 1].float()  # distribution predicting it
        loss = F.cross_entropy(pred, labels[rows, idx])
        return (loss, logits) if return_outputs else loss


args = TrainingArguments(
    output_dir=str(SRC / "lora_reranker"),
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    num_train_epochs=1,
    max_steps=4 if SMOKE else -1,
    learning_rate=1e-4,
    warmup_steps=20,
    lr_scheduler_type="cosine",
    fp16=False,
    optim="paged_adamw_8bit",
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    max_grad_norm=1.0,
    logging_steps=10,
    save_strategy="steps",
    save_steps=40,
    save_total_limit=1,
    report_to=[],
    seed=42,
)
AnswerTrainer(model=model, train_dataset=Pairs(), args=args,
              data_collator=collate).train()
model.eval()
model.save_pretrained(SRC / "lora_reranker" / "adapter")  # LoRA weights only
print("adapter saved", flush=True)


# --- Same evaluation as run05: P(yes) on the first answer token --------------------

@torch.no_grad()
def p_yes(query, batch_docs):
    """P(yes)/(P(yes)+P(no)) per document, batched."""
    prompts = [
        tok.apply_chat_template(
            [{"role": "system", "content": SYSTEM},
             {"role": "user", "content": user_msg(query, d)}],
            add_generation_prompt=True, return_dict=False,
        )
        for d in batch_docs
    ]
    pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
    maxlen = max(len(p) for p in prompts)
    ids = torch.full((len(prompts), maxlen), pad, dtype=torch.long)
    attn = torch.zeros((len(prompts), maxlen), dtype=torch.long)
    for i, p in enumerate(prompts):
        ids[i, : len(p)] = torch.tensor(p)
        attn[i, : len(p)] = 1
    logits = model(input_ids=ids.cuda(), attention_mask=attn.cuda()).logits
    last = attn.sum(1) - 1  # next-token distribution after each prompt
    rows = logits[torch.arange(len(prompts)), last.cuda()].float()
    ly = torch.logsumexp(rows[:, YES_IDS], dim=1)  # all yes/no case variants
    ln = torch.logsumexp(rows[:, NO_IDS], dim=1)
    return torch.sigmoid(ly - ln).cpu().numpy()


t0 = time.perf_counter()
ndcg = []
eval_qids = llm_qids[:5] if SMOKE else llm_qids
for n, qid in enumerate(eval_qids):
    sl = shortlists[qid]
    scores = np.concatenate(
        [p_yes(queries[qid], sl[i:i + BATCH]) for i in range(0, len(sl), BATCH)]
    )
    order = np.argsort(-scores, kind="stable")
    ndcg.append(ndcg_at_k([sl[i] for i in order], qrels["test"][qid]))
    if n < 3:  # exhibit material: gold vs distractors
        gold = [d for d in qrels["test"][qid] if d in sl]
        print(f"  query {n + 1}: {queries[qid][:60]!r}", flush=True)
        for i in order[:3]:
            mark = " <- gold" if sl[i] in gold else ""
            print(f"    P(yes) {scores[i]:.2f}  {full_text[sl[i]][:60]!r}{mark}",
                  flush=True)
eval_time = time.perf_counter() - t0
mean_ndcg = float(np.mean(ndcg))
print(
    f"era 4 llama3.2+LoRA   pointwise nDCG@10 {mean_ndcg:.3f}"
    f"  ({eval_time / len(eval_qids):.2f} s/query,"
    f" {len(eval_qids)} queries)",
    flush=True,
)
(HERE / "results" / "run06_results_lora.json").write_text(json.dumps(
    [{"method": "4 llama3.2 3B + LoRA pointwise on mxbai+BM25 top-10",
      "ndcg10": mean_ndcg,
      "search_ms": 1000 * eval_time / len(eval_qids)}],
    indent=2,
))

References

1.
Wadden D, Lin S, Lo K, et al (2020) Fact or fiction: Verifying scientific claims
2.
Thakur N, Reimers N, Rücklé A, Srivastava A, Gurevych I (2021) BEIR: A heterogeneous benchmark for zero-shot evaluation of information retrieval models. In: Thirty-fifth conference on neural information processing systems datasets and benchmarks track (round 2)
3.
Cormack GV, Clarke CLA, Buettcher S (2009) Reciprocal rank fusion outperforms condorcet and individual rank learning methods. In: Proceedings of the 32nd international ACM SIGIR conference on research and development in information retrieval. Association for Computing Machinery, New York, NY, USA, pp 758–759
4.
Sun W, Yan L, Ma X, et al (2023) Is ChatGPT good at search? Investigating large language models as re-ranking agents. In: Bouamor H, Pino J, Bali K (eds) Proceedings of the 2023 conference on empirical methods in natural language processing. Association for Computational Linguistics, Singapore, pp 14918–14937
5.
Hu EJ, Shen Y, Wallis P, et al (2021) LoRA: Low-rank adaptation of large language models