Tevatron

https://texttron.github.io/tevatron
Admin
Tevatron is a simple and efficient toolkit for training and running dense retrievers with deep...
Tokens:3,413
Snippets:46
Trust Score:9.7
Update:1 month ago
Show doc for...
Tokens:
Raw
### Install GradCache

Source: https://texttron.github.io/tevatron

Optional installation to enable gradient cache features during training.

```bash
git clone https://github.com/luyug/GradCache
cd GradCache
pip install .
```

--------------------------------

### Install Tevatron for development

Source: https://texttron.github.io/tevatron

Clones the repository and installs it in editable mode for research and development purposes.

```bash
git https://github.com/texttron/tevatron
cd tevatron
pip install --editable .
```

--------------------------------

### Install Tevatron via pip

Source: https://texttron.github.io/tevatron

Standard installation method using the Python package manager.

```bash
pip install tevatron
```

--------------------------------

### Train with Raw Data

Source: https://texttron.github.io/tevatron/datasets

Example command for training using a dataset formatted like Wikipedia-NQ.

```bash
python -m tevatron.driver.train \
  ... \
  --dataset_name Tevatron/wikipedia-nq \
  --train_dir train_dir \
  ...
```

--------------------------------

### Encode Query with Pre-tokenized Data

Source: https://texttron.github.io/tevatron/datasets

Example command for encoding queries using pre-tokenized data.

```bash
python tevatron.driver.encode \
  ... \
  --encode_in_path dev_dir/dev_data.jsonl \
  --encode_is_qry \
  ...
```

--------------------------------

### Train with Pre-tokenized Data

Source: https://texttron.github.io/tevatron/datasets

Example command for training using pre-tokenized data by omitting or setting --dataset_name to json.

```bash
python -m tevatron.driver.train \
  ... \
  --train_dir train_dir \
  ...
```

--------------------------------

### Encode Corpus with Pre-tokenized Data

Source: https://texttron.github.io/tevatron/datasets

Example command for encoding a corpus using pre-tokenized data.

```bash
python tevatron.driver.encode \
  ... \
  --encode_in_path corpus_dir/corpus_jsonl \
  ...
```

--------------------------------

### Encode Corpus with Raw Data

Source: https://texttron.github.io/tevatron/datasets

Example command for encoding a corpus using the Wikipedia-NQ format.

```bash
python tevatron.driver.encode \
  ... \
  --dataset_name Tevatron/wikipedia-nq-corpus \
  --encode_in_path corpus_dir/corpus_jsonl \
  ...
```

--------------------------------

### Encode Query with Raw Data

Source: https://texttron.github.io/tevatron/datasets

Example command for encoding queries using the Wikipedia-NQ format.

```bash
python tevatron.driver.encode \
  ... \
  --dataset_name Tevatron/wikipedia-nq \
  --encode_in_path dev_dir/dev_data.jsonl \
  --encode_is_qry \
  ...
```

--------------------------------

### Train with Gradient Cache

Source: https://texttron.github.io/tevatron/training

Enables gradient caching to support larger batch sizes on memory-limited GPU setups.

```bash
CUDA_VISIBLE_DEVICES=0 python -m tevatron.driver.train \
  --output_dir model_nq \
  --dataset_name Tevatron/wikipedia-nq \
  --model_name_or_path bert-base-uncased \
  --do_train \
  --save_steps 20000 \
  --fp16 \
  --per_device_train_batch_size 128 \
  --train_n_passages 2 \
  --learning_rate 1e-5 \
  --q_max_len 32 \
  --p_max_len 156 \
  --num_train_epochs 40 \
  --grad_cache \
  --gc_q_chunk_size 32 \
  --gc_p_chunk_size 16
```

--------------------------------

### Train on TPU using Jax/Flax

Source: https://texttron.github.io/tevatron/training

Executes the training process on a V3-8 TPU VM using the tevatron.driver.jax_train module.

```bash
python -m tevatron.driver.jax_train \
  --output_dir model_nq \
  --dataset_name Tevatron/wikipedia-nq \
  --model_name_or_path bert-base-uncased \
  --do_train \
  --per_device_train_batch_size 16 \
  --train_n_passages 2 \
  --learning_rate 1e-5 \
  --q_max_len 32 \
  --p_max_len 156 \
  --num_train_epochs 40
```

--------------------------------

### Train with self-contained dataset

Source: https://texttron.github.io/tevatron/datasets

Command to initiate training using a self-contained dataset from HuggingFace.

```bash
python -m tevatron.driver.train \
  --do_train \
  --output_dir model_scifact \
  --dataset_name Tevatron/scifact \
  --model_name_or_path bert-base-uncased \
  --per_device_train_batch_size 16 \
  --learning_rate 1e-5 \
  --num_train_epochs 5
```

--------------------------------

### Encode on TPU with JAX

Source: https://texttron.github.io/tevatron/encoding

Use the tevatron.driver.jax_encode module to perform encoding on TPU hardware.

```bash
python -m tevatron.driver.jax_encode \
  --output_dir=temp \
  --model_name_or_path model_nq \
  --per_device_eval_batch_size 156 \
  --p_max_len 128 \
  --dataset_name Tevatron/wikipedia-nq-corpus \
  --encoded_save_path corpus_emb.pkl
```

--------------------------------

### Perform Basic Retrieval

Source: https://texttron.github.io/tevatron/retrieval

Executes a brute-force search using the tevatron.faiss_retriever module. Ensure query and corpus embeddings are available as pickle files.

```bash
python -m tevatron.faiss_retriever \
--query_reps query.pkl \
--passage_reps corpus_emb*.pkl \
--depth 100 \
--batch_size -1 \
--save_text \
--save_ranking_to rank.txt
```

--------------------------------

### Train a Dense Retriever

Source: https://texttron.github.io/tevatron/training

Executes a basic training run for a dense retriever using distributed training across multiple GPUs.

```bash
python -m torch.distributed.launch --nproc_per_node=4 -m tevatron.driver.train \
  --output_dir model_nq \
  --dataset_name Tevatron/wikipedia-nq \
  --model_name_or_path bert-base-uncased \
  --do_train \
  --save_steps 20000 \
  --fp16 \
  --per_device_train_batch_size 32 \
  --train_n_passages 2 \
  --learning_rate 1e-5 \
  --q_max_len 32 \
  --p_max_len 156 \
  --num_train_epochs 40 \
  --negatives_x_device
```

--------------------------------

### Encode Queries with Tevatron

Source: https://texttron.github.io/tevatron/encoding

Generate embeddings for query datasets using the --encode_is_qry flag.

```bash
python -m tevatron.driver.encode \
  --output_dir=temp \
  --model_name_or_path model_nq \
  --fp16 \
  --per_device_eval_batch_size 156 \
  --dataset_name Tevatron/wikipedia-nq/test \
  --encoded_save_path query.pkl \
  --q_max_len 32 \
  --encode_is_qry
```

--------------------------------

### Pre-tokenized Data Formats

Source: https://texttron.github.io/tevatron/datasets

Required JSONL structures for training and encoding when using pre-tokenized data.

```json
{'query': TEXT_TYPE, 'positives': List[TEXT_TYPE], 'negatives': List[TEXT_TYPE]}
```

```json
{text_id: "xxx", 'text': TEXT_TYPE}
```

--------------------------------

### Encode self-contained topics

Source: https://texttron.github.io/tevatron/datasets

Command to encode query topics for a self-contained dataset.

```bash
python tevatron.driver.encode \
  --do_encode \
  --output_dir=temp_out \
  --model_name_or_path model_scifact \
  --per_device_eval_batch_size 64 \
  --dataset_name Tevatron/scifact/dev \
  --encode_is_qry \
  --q_max_len 64 \
  --encoded_save_path queries_emb.pkl
```

--------------------------------

### Reduce Sharded Results

Source: https://texttron.github.io/tevatron/retrieval

Combines search results from multiple shards into a single ranking file. Requires document and query IDs to be integers.

```bash
python -m tevatron.faiss_retriever.reducer \
--score_dir ${INTERMEDIATE_DIR} \
--query query.pkl \
--save_ranking_to rank.txt
```

--------------------------------

### Encode Corpus with Tevatron

Source: https://texttron.github.io/tevatron/encoding

Use the tevatron.driver.encode module to generate embeddings for a corpus dataset.

```bash
python -m tevatron.driver.encode \
  --output_dir=temp \
  --model_name_or_path model_nq \
  --fp16 \
  --per_device_eval_batch_size 156 \
  --p_max_len 128 \
  --dataset_name Tevatron/wikipedia-nq-corpus \
  --encoded_save_path corpus_emb.pkl
```

--------------------------------

### Exactly Matched Dataset Format

Source: https://texttron.github.io/tevatron/datasets

JSON structure for datasets where passage relevancy is determined by exact answer matching.

```json
{
   "query_id": "<query id>",
   "query": "<query text>",
   "answers": ["<answer>"],
   "positive_passages": [
     {"docid": "<passage id>", "title": "<passage title>", "text": "<passage body>"}
   ],
   "negative_passages": [
     {"docid": "<passage id>", "title": "<passage title>", "text": "<passage body>"}
   ]
}
```

--------------------------------

### Sharded Corpus Encoding

Source: https://texttron.github.io/tevatron/encoding

Parallelize the encoding process for large corpora by splitting the data into multiple shards.

```bash
for s in $(seq -f "%02g" 0 19)
do
python -m tevatron.driver.encode \
  --output_dir=temp \
  --model_name_or_path model_nq \
  --fp16 \
  --per_device_eval_batch_size 156 \
  --p_max_len 128 \
  --dataset_name Tevatron/wikipedia-nq-corpus \
  --encoded_save_path corpus_emb_${s}.pkl \
  --encode_num_shard 20 \
  --encode_shard_index ${s}
done
```

--------------------------------

### Encode self-contained corpus

Source: https://texttron.github.io/tevatron/datasets

Command to encode the corpus associated with a self-contained dataset.

```bash
python tevatron.driver.encode \
  --do_encode \
  --output_dir=temp_out \
  --model_name_or_path model_scifact \
  --per_device_eval_batch_size 64 \
  --dataset_name Tevatron/scifact-corpus \
  --p_max_len 512 \
  --encoded_save_path corpus_emb.pkl
```

--------------------------------

### Perform Sharded Search

Source: https://texttron.github.io/tevatron/retrieval

Parallelizes search over corpus shards when embeddings are too large for memory. Requires a subsequent reduction step to combine results.

```bash
INTERMEDIATE_DIR=intermediate
mkdir ${INTERMEDIATE_DIR}
for s in $(seq -f "%02g" 0 19)
do
python -m tevatron.faiss_retriever \
--query_reps query_emb.pkl \
--passage_reps corpus_emb${s}.pkl \
--depth 100 \
--save_ranking_to ${INTERMEDIATE_DIR}/${s}
done
```

--------------------------------

### Relevancy Judged Dataset Format

Source: https://texttron.github.io/tevatron/datasets

JSON structure for datasets where passage relevancy is pre-annotated.

```json
{
   "query_id": "<query id>",
   "query": "<query text>",
   "positive_passages": [
     {"docid": "<passage id>", "title": "<passage title>", "text": "<passage body>"}
   ],
   "negative_passages": [
     {"docid": "<passage id>", "title": "<passage title>", "text": "<passage body>"}
   ]
}
```
Tokens:
Raw