"""
Runs benchmark evaluations with the BEIR dataset.
Install txtai and the following dependencies to run:
pip install txtai pytrec_eval rank-bm25 bm25s elasticsearch psutil
"""
import argparse
import csv
import json
import os
import pickle
import sqlite3
import time
import psutil
import yaml
import numpy as np
from bm25s import BM25 as BM25Sparse
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
from pytrec_eval import RelevanceEvaluator
from rank_bm25 import BM25Okapi
from txtai.embeddings import Embeddings
from txtai.pipeline import LLM, RAG, Tokenizer
from txtai.scoring import ScoringFactory
class Index:
"""
Base index definition. Defines methods to index and search a dataset.
"""
def __init__(self, path, config, output, refresh):
"""
Creates a new index.
Args:
path: path to dataset
config: path to config file
output: path to store index
refresh: overwrites existing index if True, otherwise existing index is loaded
"""
self.path = path
self.config = config
self.output = output
self.refresh = refresh
# Build and save index
self.backend = self.index()
def __call__(self, limit, filterscores=True):
"""
Main evaluation logic. Loads an index, runs the dataset queries and returns the results.
Args:
limit: maximum results
filterscores: if exact matches should be filtered out
Returns:
search results
"""
uids, queries = self.load()
# Run queries in batches
offset, results = 0, {}
for batch in self.batch(queries, 256):
for i, r in enumerate(self.search(batch, limit + 1)):
# Get result as list of (id, score) tuples
r = list(r)
r = [(x["id"], x["score"]) for x in r] if r and isinstance(r[0], dict) else r
if filterscores:
r = [(uid, score) for uid, score in r if uid != uids[offset + i]][:limit]
results[uids[offset + i]] = dict(r)
# Increment offset
offset += len(batch)
return results
def search(self, queries, limit):
"""
Runs a search for a set of queries.
Args:
queries: list of queries to run
limit: maximum results
Returns:
search results
"""
return self.backend.batchsearch(queries, limit)
def index(self):
"""
Indexes a dataset.
"""
raise NotImplementedError
def rows(self):
"""
Iterates over the dataset yielding a row at a time for indexing.
"""
with open(f"{self.path}/corpus.jsonl", encoding="utf-8") as f:
for line in f:
row = json.loads(line)
text = f'{row["title"]}. {row["text"]}' if row.get("title") else row["text"]
if text:
yield (row["_id"], text, None)
def load(self):
"""
Loads queries for the dataset. Returns a list of expected result ids and input queries.
Returns:
(result ids, input queries)
"""
with open(f"{self.path}/queries.jsonl", encoding="utf-8") as f:
data = [json.loads(query) for query in f]
uids, queries = [x["_id"] for x in data], [x["text"] for x in data]
return uids, queries
def batch(self, data, size):
"""
Splits data into equal sized batches.
Args:
data: input data
size: batch size
Returns:
data split into equal size batches
"""
return [data[x : x + size] for x in range(0, len(data), size)]
def readconfig(self, key, default):
"""
Reads configuration from a config file. Returns default configuration
if config file is not found or config key isn't present.
Args:
key: configuration key to lookup
default: default configuration
Returns:
config if found, otherwise returns default config
"""
if self.config and os.path.exists(self.config):
# Read configuration
with open(self.config, "r", encoding="utf-8") as f:
# Check for config
config = yaml.safe_load(f)
if key in config:
return config[key]
return default
class Embed(Index):
"""
Embeddings index using txtai.
"""
def index(self):
if os.path.exists(self.output) and not self.refresh:
embeddings = Embeddings()
embeddings.load(self.output)
else:
# Read configuration
config = self.readconfig("embeddings", {"batch": 8192, "encodebatch": 128, "faiss": {"quantize": True, "sample": 0.05}})
# Build index
embeddings = Embeddings(config)
embeddings.index(self.rows())
embeddings.save(self.output)
return embeddings
class Hybrid(Index):
"""
Hybrid embeddings + BM25 index using txtai.
"""
def index(self):
if os.path.exists(self.output) and not self.refresh:
embeddings = Embeddings()
embeddings.load(self.output)
else:
# Read configuration
config = self.readconfig(
"hybrid",
{
"batch": 8192,
"encodebatch": 128,
"faiss": {"quantize": True, "sample": 0.05},
"scoring": {"method": "bm25", "terms": True, "normalize": True},
},
)
# Build index
embeddings = Embeddings(config)
embeddings.index(self.rows())
embeddings.save(self.output)
return embeddings
class RetrievalAugmentedGeneration(Embed):
"""
Retrieval augmented generation (RAG) using txtai.
"""
def __init__(self, path, config, output, refresh):
# Parent logic
super().__init__(path, config, output, refresh)
# Read LLM configuration
llm = self.readconfig("llm", {})
# Read RAG configuration
rag = self.readconfig("rag", {})
# Load RAG pipeline
self.rag = RAG(self.backend, LLM(**llm), output="reference", **rag)
def search(self, queries, limit):
# Set context window size to limit and run
self.rag.context = limit
return [[(x["reference"], 1)] for x in self.rag(queries, maxlength=4096)]
class Score(Index):
"""
BM25 index using txtai.
"""
def index(self):
# Read configuration
config = self.readconfig("scoring", {"method": "bm25", "terms": True})
# Create scoring instance
scoring = ScoringFactory.create(config)
output = os.path.join(self.output, "scoring")
if os.path.exists(output) and not self.refresh:
scoring.load(output)
else:
scoring.index(self.rows())
scoring.save(output)
return scoring
class RankBM25(Index):
"""
BM25 index using rank-bm25.
"""
def search(self, queries, limit):
ids, backend = self.backend
tokenizer, results = Tokenizer(), []
for query in queries:
scores = backend.get_scores(tokenizer(query))
topn = np.argsort(scores)[::-1][:limit]
results.append([(ids[x], scores[x]) for x in topn])
return results
def index(self):
output = os.path.join(self.output, "rank")
if os.path.exists(output) and not self.refresh:
with open(output, "rb") as f:
ids, model = pickle.load(f)
else:
# Tokenize data
tokenizer, data = Tokenizer(), []
for uid, text, _ in self.rows():
data.append((uid, tokenizer(text)))
ids = [uid for uid, _ in data]
model = BM25Okapi([text for _, text in data])
# Save model
with open(output, "wb") as out:
pickle.dump(model, out)
return ids, model
class BM25S(Index):
"""
BM25 as implemented by bm25s
"""
def __init__(self, path, config, output, refresh):
# Corpus ids
self.ids = None
# Parent logic
super().__init__(path, config, output, refresh)
def search(self, queries, limit):
tokenizer = Tokenizer()
results, scores = self.backend.retrieve([tokenizer(x) for x in queries], corpus=self.ids, k=limit)
# List of queries => list of matches (id, score)
x = []
for a, b in zip(results, scores):
x.append([(str(c), float(d)) for c, d in zip(a, b)])
return x
def index(self):
tokenizer = Tokenizer()
ids, texts = [], []
for uid, text, _ in self.rows():
ids.append(uid)
texts.append(text)
self.ids = ids
if os.path.exists(self.output) and not self.refresh:
model = BM25Sparse.load(self.output)
else:
model = BM25Sparse(method="lucene", k1=1.2, b=0.75)
model.index([tokenizer(x) for x in texts], leave_progress=False)
model.save(self.output)
return model
class SQLiteFTS(Index):
"""
BM25 index using SQLite's FTS extension.
"""
def search(self, queries, limit):
tokenizer, results = Tokenizer(), []
for query in queries:
query = tokenizer(query)
query = " OR ".join([f'"{q}"' for q in query])
self.backend.execute(
f"SELECT id, bm25(textindex) * -1 score FROM textindex WHERE text MATCH ? ORDER BY bm25(textindex) LIMIT {limit}", [query]
)
results.append(list(self.backend))
return results
def index(self):
if os.path.exists(self.output) and not self.refresh:
# Load existing database
connection = sqlite3.connect(self.output)
else:
# Delete existing database
if os.path.exists(self.output):
os.remove(self.output)
# Create new database
connection = sqlite3.connect(self.output)
# Tokenize data
tokenizer, data = Tokenizer(), []
for uid, text, _ in self.rows():
data.append((uid, " ".join(tokenizer(text))))
# Create table
connection.execute("CREATE VIRTUAL TABLE textindex using fts5(id, text)")
# Load data and build index
connection.executemany("INSERT INTO textindex VALUES (?, ?)", data)
connection.commit()
return connection.cursor()
class Elastic(Index):
"""
BM25 index using Elasticsearch.
"""
def search(self, queries, limit):
# Generate bulk queries
request = []
for query in queries:
req_head = {"index": "textindex", "search_type": "dfs_query_then_fetch"}
req_body = {
"_source": False,
"query": {"multi_match": {"query": query, "type": "best_fields", "fields": ["text"], "tie_breaker": 0.5}},
"size": limit,
}
request.extend([req_head, req_body])
# Run ES query
response = self.backend.msearch(body=request, request_timeout=600)
# Read responses
results = []
for resp in response["responses"]:
result = resp["hits"]["hits"]
results.append([(r["_id"], r["_score"]) for r in result])
return results
def index(self):
es = Elasticsearch("http://localhost:9200")
# Delete existing index
# pylint: disable=W0702
try:
es.indices.delete(index="textindex")
except:
pass
bulk(es, ({"_index": "textindex", "_id": uid, "text": text} for uid, text, _ in self.rows()))
es.indices.refresh(index="textindex")
return es
def relevance(path):
"""
Loads relevance data for evaluation.
Args:
path: path to dataset test file
Returns:
relevance data
"""
rel = {}
with open(f"{path}/qrels/test.tsv", encoding="utf-8") as f:
reader = csv.reader(f, delimiter="\t", quoting=csv.QUOTE_MINIMAL)
next(reader)
for row in reader:
queryid, corpusid, score = row[0], row[1], int(row[2])
if queryid not in rel:
rel[queryid] = {corpusid: score}
else:
rel[queryid][corpusid] = score
return rel
def create(method, path, config, output, refresh):
"""
Creates a new index.
Args:
method: indexing method
path: path to dataset
config: path to config file
output: path to store index
refresh: overwrites existing index if True, otherwise existing index is loaded
Returns:
Index
"""
if method == "hybrid":
return Hybrid(path, config, output, refresh)
if method == "rag":
return RetrievalAugmentedGeneration(path, config, output, refresh)
if method == "scoring":
return Score(path, config, output, refresh)
if method == "rank":
return RankBM25(path, config, output, refresh)
if method == "bm25s":
return BM25S(path, config, output, refresh)
if method == "sqlite":
return SQLiteFTS(path, config, output, refresh)
if method == "es":
return Elastic(path, config, output, refresh)
# Default
return Embed(path, config, output, refresh)
def compute(results):
"""
Computes metrics using the results from an evaluation run.
Args:
results: evaluation results
Returns:
metrics
"""
metrics = {}
for r in results:
for metric in results[r]:
if metric not in metrics:
metrics[metric] = []
metrics[metric].append(results[r][metric])
return {metric: round(np.mean(values), 5) for metric, values in metrics.items()}
def evaluate(methods, path, args):
"""
Runs an evaluation.
Args:
methods: list of indexing methods to test
path: path to dataset
args: command line arguments
Returns:
{calculated performance metrics}
"""
print(f"------ {os.path.basename(path)} ------")
# Performance stats
performance = {}
# Calculate stats for each model type
topk = args.topk
evaluator = RelevanceEvaluator(relevance(path), {f"ndcg_cut.{topk}", f"map_cut.{topk}", f"recall.{topk}", f"P.{topk}"})
for method in methods:
# Stats for this source
stats = {}
performance[method] = stats
# Create index and get results
start = time.time()
output = args.output if args.output else f"{path}/{method}"
index = create(method, path, args.config, output, args.refresh)
# Add indexing metrics
stats["index"] = round(time.time() - start, 2)
stats["memory"] = int(psutil.Process().memory_info().rss / (1024 * 1024))
stats["disk"] = int(sum(d.stat().st_size for d in os.scandir(output) if d.is_file()) / 1024) if os.path.isdir(output) else 0
print("INDEX TIME =", time.time() - start)
print(f"MEMORY USAGE = {stats['memory']} MB")
print(f"DISK USAGE = {stats['disk']} KB")
start = time.time()
results = index(topk)
# Add search metrics
stats["search"] = round(time.time() - start, 2)
print("SEARCH TIME =", time.time() - start)
# Calculate stats
metrics = compute(evaluator.evaluate(results))
# Add accuracy metrics
for stat in [f"ndcg_cut_{topk}", f"map_cut_{topk}", f"recall_{topk}", f"P_{topk}"]:
stats[stat] = metrics[stat]
# Print model stats
print(f"------ {method} ------")
print(f"NDCG@{topk} =", metrics[f"ndcg_cut_{topk}"])
print(f"MAP@{topk} =", metrics[f"map_cut_{topk}"])
print(f"Recall@{topk} =", metrics[f"recall_{topk}"])
print(f"P@{topk} =", metrics[f"P_{topk}"])
print()
return performance
def benchmarks(args):
"""
Main benchmark execution method.
Args:
args: command line arguments
"""
# Directory where BEIR datasets are stored
directory = args.directory if args.directory else "beir"
if args.sources and args.methods:
sources, methods = args.sources.split(","), args.methods.split(",")
mode = "a"
else:
# Default sources and methods
sources = [
"trec-covid",
"nfcorpus",
"nq",
"hotpotqa",
"fiqa",
"arguana",
"webis-touche2020",
"quora",
"dbpedia-entity",
"scidocs",
"fever",
"climate-fever",
"scifact",
]
methods = ["embed", "hybrid", "rag", "scoring", "rank", "bm25s", "sqlite", "es"]
mode = "w"
# Run and save benchmarks
with open("benchmarks.json", mode, encoding="utf-8") as f:
for source in sources:
# Run evaluations
results = evaluate(methods, f"{directory}/{source}", args)
# Save as JSON lines output
for method, stats in results.items():
stats["source"] = source
stats["method"] = method
stats["name"] = args.name if args.name else method
json.dump(stats, f)
f.write("\n")
if __name__ == "__main__":
# Command line parser
parser = argparse.ArgumentParser(description="Benchmarks")
parser.add_argument("-c", "--config", help="path to config file", metavar="CONFIG")
parser.add_argument("-d", "--directory", help="root directory path with datasets", metavar="DIRECTORY")
parser.add_argument("-m", "--methods", help="comma separated list of methods", metavar="METHODS")
parser.add_argument("-n", "--name", help="name to assign to this run, defaults to method name", metavar="NAME")
parser.add_argument("-o", "--output", help="index output directory path", metavar="OUTPUT")
parser.add_argument(
"-r",
"--refresh",
help="refreshes index if set, otherwise uses existing index if available",
action="store_true",
)
parser.add_argument("-s", "--sources", help="comma separated list of data sources", metavar="SOURCES")
parser.add_argument("-t", "--topk", help="top k results to use for the evaluation", metavar="TOPK", type=int, default=10)
# Calculate benchmarks
benchmarks(parser.parse_args())