NeuML / txtai-wikipedia

Wikipedia txtai embeddings index

This is a txtai embeddings index for the English edition of Wikipedia .

This index is built from the Wikipedia January 2024 dataset . Only the first paragraph of the lead section from each article is included in the index. This is similar to an abstract of the article.

It also uses Wikipedia Page Views data to add a percentile field. The percentile field can be used to only match commonly visited pages.

txtai must be installed to use this model.

Example

Version 5.4 added support for loading embeddings indexes from the Hugging Face Hub. See the example below.

from txtai.embeddings import Embeddings

# Load the index from the HF Hub
embeddings = Embeddings()
embeddings.load(provider="huggingface-hub", container="neuml/txtai-wikipedia")

# Run a search
embeddings.search("Roman Empire")

# Run a search matching only the Top 1% of articles
embeddings.search("""
   SELECT id, text, score, percentile FROM txtai WHERE similar('Boston') AND
   percentile >= 0.99
""")

Use Cases

An embeddings index generated by txtai is a fully encapsulated index format. It doesn't require a database server or dependencies outside of the Python install.

The Wikipedia index works well as a fact-based context source for retrieval augmented generation (RAG). In other words, search results from this model can be passed to LLM prompts as the context in which to answer questions.

See this article for additional examples on how to use this model.

Evaluation Results

Performance was evaluated using the NDCG@10 score with a custom question-answer evaluation set . Results are shown below.

e5-base is the best performing model for the evaluation set. This highlights the importance of testing models as e5-base is far from the leading model on the MTEB leaderboard . Benchmark datasets are only a guide.

Build the index

The following steps show how to build this index. These scripts are using the latest data available as of 2024-01-01, update as appropriate.

• Install required build dependencies

pip install txtchat mwparserfromhell datasets

• Download and build pageviews database

mkdir -p pageviews/data
wget -P pageviews/data https://dumps.wikimedia.org/other/pageview_complete/monthly/2023/2023-12/pageviews-202312-user.bz2
python -m txtchat.data.wikipedia.views -p en.wikipedia -v pageviews

• Build Wikipedia dataset

from datasets import load_dataset

# Data dump date from https://dumps.wikimedia.org/enwiki/
date = "20240101"

# Build and save dataset
ds = load_dataset("neuml/wikipedia", language="en", date=date)
ds.save_to_disk(f"wikipedia-{date}")

• Build txtai-wikipedia index

python -m txtchat.data.wikipedia.index \
       -d wikipedia-20240101 \
       -o txtai-wikipedia \
       -v pageviews/pageviews.sqlite