bagel.py

import argparse
import json
import os

import torch
import wandb
from tqdm import tqdm

from src.bagel import bagel_retrieval
from src.evaluate import evaluate
from src.llm import handle_llm
from src.utils import load_dataset, seed_everything

def main(dataset_name, model_name, acq_func, beta, llm_budget, k_warm_start, kernel, args):
    dtype = torch.float64
    if not args.wandb_disable:
        configs = dict(vars(args))
        run = wandb.init(
            project="bagel",
            config=configs,
            group=args.wandb_group,
        )
    else:
        run = None

    base_path = os.path.dirname(os.path.abspath(__file__))
    dataset, cache, relevance_map, queries, passages, query_ids, passage_ids, query_embeddings, passage_embeddings = (
        load_dataset(base_path, dataset_name, model_name, args.llm_name, args.prompt_type))

    if args.debug:
        print("\n > DEBUG MODE")
        query_ids = query_ids[[0]]
        queries = queries[[0]]
        query_embeddings = query_embeddings[[0]]

    k_retrieval = max(args.cutoff)
    llm = handle_llm(args.llm_name, args.prompt_type, args.score_type)

    print("\n")
    results = {}
    passage_embeddings = torch.from_numpy(passage_embeddings).to(dtype=dtype)
    for query, q_id, q_emb in tqdm(zip(queries, query_ids, query_embeddings), desc=" > Bandit Ranking", total=len(queries)):
        query_embedding = torch.from_numpy(q_emb).to(dtype=dtype)
        preds, scores, warm, bandit, length_scale, _ = bagel_retrieval(
            query=query,
            query_id=q_id,
            query_embedding=query_embedding,
            passages=passages,
            passage_ids=passage_ids.copy(),
            passage_embeddings=passage_embeddings,

            llm=llm,
            llm_budget=llm_budget,
            score_type=args.score_type,
            k_warm_start=k_warm_start,

            kernel=kernel,
            acq_func=acq_func,
            alpha=args.alpha,
            length_scale=args.length_scale,
            beta=beta,

            k_retrieval=k_retrieval,
            return_score=True,
            cache=cache,
            update_cache=dataset.cache_path,
            verbose=args.verbose,
            dtype=dtype
        )

        results[q_id] = {
            "pred": preds,
            "score": scores,
            "warm_passages": warm,
            "bandit_passages": bandit,
        }
        if run is not None:
            run.log({"length_scale": length_scale})
    metric, results = evaluate(results, relevance_map, args.cutoff, threshold=dataset.relevance_threshold)

    if run is not None:
        updated_dict = {}
        for k, v in metric.items():
            new_key = str(k).replace("@", "/")
            updated_dict[new_key] = v
        wandb.log(updated_dict)

        os.makedirs(f"results/{dataset_name}/", exist_ok=True)
        with open(f"results/{dataset_name}/{run.name}.json", "w", encoding="utf-8") as f:
            json.dump(results, f, indent=1, ensure_ascii=False)

def arg_parser():
    parser = argparse.ArgumentParser(description='Implementation of BAGEL')
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility.")

    parser.add_argument('--dataset_name', type=str, required=True,
                        help='Dataset to evaluate.', choices=['covid', 'nfcorpus', 'robust', 'traveldest'])

    parser.add_argument("--llm_name", type=str, default='unsloth/Qwen3-14B-unsloth-bnb-4bit',
                        choices=['unsloth/Qwen3-14B-unsloth-bnb-4bit', 'openai/gpt4o'],
                        help="LLM backend used for relevance scoring.")
    parser.add_argument("--prompt_type", type=str, default="zeroshot", choices=["zeroshot"],
                        help="Prompt template type.")
    parser.add_argument("--score_type", type=str, choices=['er', 'pr'], default='er',
                        help="Score to use from LLM outputs: expected relevance (er) or predicted relevance class (pr).")

    parser.add_argument('--llm_budget', type=int, default=50,
                        help='Total number of LLM scoring calls per query.')
    parser.add_argument('--warm_start', type=int, default=25,
                        help='Number of top-dense passages scored before exploration.')
    parser.add_argument('--acq_func', type=str, default='ucb',
                        choices=['ucb', 'ei', 'pi', 'thompson', 'random', 'dense'],
                        help='Acquisition strategy for selecting next passage.')
    parser.add_argument('--kernel', type=str, default='rbf', help='Gaussian Process kernel type.',
                        choices=['rbf', 'linear', 'matern'])

    parser.add_argument("--alpha", type=float, default=1e-3,
                        help='Noise term for GP likelihood (numerical stability / observation noise).')
    parser.add_argument("--length_scale", type=float, default=1.0,
                        help='Initial kernel length scale for Gaussian Process.')
    parser.add_argument('--beta', type=float, default=2,
                        help='Exploration coefficient for UCB (used only when --acq_func ucb).')

    parser.add_argument('--emb_model', type=str, default='all-MiniLM-L6-v2',
                        help='Sentence embedding model name.')
    parser.add_argument("--cutoff", type=int, nargs="+", default=[1, 5, 10, 50, 100],
                        help='Evaluation cutoffs for ranking metrics (e.g., --cutoff 1 5 10 50).')

    parser.add_argument("--wandb_disable", action="store_true", help="Disable Weights & Biases logging.")
    parser.add_argument("--wandb_group", type=str, default=None, help="W&B group name.")

    parser.add_argument("--debug", action="store_true", help="Run only the first query.")
    parser.add_argument("--verbose", action="store_true", help="Print detailed bandit progress logs.")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = arg_parser()
    seed_everything(args.seed)
    main(dataset_name=args.dataset_name, model_name=args.emb_model, acq_func=args.acq_func, beta=args.beta,
         llm_budget=args.llm_budget, k_warm_start=args.warm_start, kernel=args.kernel, args=args)