Skip to content
This repository was archived by the owner on Feb 12, 2025. It is now read-only.

Commit 7e62e77

Browse files
ibhatiibhati
authored
Restructure and add WIT (#2)
* Restructure the DPR dataset to its own directory and added WIT
1 parent 182363f commit 7e62e77

12 files changed

+387
-92
lines changed

‎README.md

+3-92
Original file line numberDiff line numberDiff line change
@@ -1,94 +1,5 @@
1-
# DPR Dataset Generator
1+
# Vector Search Datasets
22

3-
This repository provides code to generate base and query vector datasets for similarity search benchmarking and evaluation on high-dimensional vectors stemming from large language models.
4-
With the dense passage retriever (DPR) [[1]](#1), we encode text snippets from the C4 dataset [[2]](#2) to generate 768-dimensional vectors:
5-
- context DPR embeddings for the base set and
6-
- question DPR embeddings for the query set.
3+
This repository provides code to generate several datasets for similarity search benchmarking and evaluation on high-dimensional vectors stemming from recent deep learning models.
4+
Please see the details of each dataset in the respective README files.
75

8-
The metric for similarity search is inner product [[1]](#1).
9-
10-
The number of base and query embedding vectors is parametrizable.
11-
12-
## DPR10M
13-
14-
A specific instance with 10 million base vectors and 10,000 query vectors is introduced in [[3]](#3). Use the script [dpr_dataset_10M.py](dpr_dataset_10M.py) to generate this dataset. The corresponding ground-truth (available [here](gtruth_dpr10M_innerProduct.ivecs)) is generated conducting an exhaustive search with the inner product
15-
metric.
16-
17-
Here is a summary of the **steps to generate this dataset**:
18-
19-
1. **Download the files** corresponding to the `en` variant of the C4 dataset accesible [here](https://huggingface.co/datasets/allenai/c4).
20-
The complete set of files requires 350GB of storage, so you might want to follow the instructions to download only a subset. For example, to generate 10M embeddings
21-
we used the first 2 files from the train set (i.e., files `c4-train.00000-of-01024.json.gz` and `c4-train.00001-of-01024.json.gz` in `c4/en/train`).
22-
23-
2. **Execute** the `generate_dpr_embeddings` function to generate a `.fvecs` file containing the new embeddings.
24-
Note that different settings should be used to generate the **base vectors** and the **query set**, as they use the
25-
DPR context and query encoders respectively.
26-
See the script [dpr_dataset_10M.py](dpr_dataset_10M.py) for details.
27-
28-
```
29-
# Example code to generate base vectors
30-
31-
from dpr_dataset_generate import generate_dpr_embeddings
32-
33-
base_C4_folder = '/home/username/research/datasets/c4/en' # Set this path to where your c4/en folder is located
34-
cache_folder = f'/home/username/.cache/huggingface/datasets/' # Set to the hugginface datasets cache path
35-
dataset_dir = f'{base_C4_folder}/train/'
36-
37-
num_embd = 10000000
38-
init_file = 0
39-
num_of_files = 2 # Make sure the input files (2 in this case) are enough to generate the
40-
# requested number of embeddings.
41-
# To get an estimate, use the optional parameter get_total_embeddings_only
42-
# to get the number of embeddings that can be generated from a certain
43-
# group of files without actually generating the embeddings.
44-
fname_prefix_out = 'c4-en'
45-
doc_stride = 32
46-
max_length = 64
47-
batch_size = 512
48-
dim = 768
49-
50-
generate_dpr_embeddings(init_file, num_of_files, num_embd, doc_stride, max_length, dim,
51-
batch_size,
52-
dataset_dir, fname_prefix_out, cache_folder)
53-
```
54-
3. **Generate the ground-truth** by conducting an exhaustive search with the inner product metric.
55-
We provide the [ground-truth](gtruth_dpr10M_innerProduct.ivecs) for the dataset generated using
56-
[dpr_dataset_10M.py](dpr_dataset_10M.py).
57-
58-
> **_NOTE:_** Due to floating-point arithmetic precision the vector embeddings generated using the provided
59-
> code in different machines may slightly vary. Keep in mind that this could cause small discrepancies with the provided ground-truth.
60-
61-
4. Functions `read_fvecs` and `read_ivecs` can be used to read `.fvecs` and `.ivecs` files respectively.
62-
63-
## References
64-
Reference to cite when you use datasets generated with this code in a research paper:
65-
66-
```
67-
@article{aguerrebere2023similarity,
68-
title={Similarity search in the blink of an eye with compressed indices},
69-
volume = {16},
70-
number = {11},
71-
pages = {3433--3446},
72-
journal = {Proceedings of the VLDB Endowment},
73-
author={Cecilia Aguerrebere and Ishwar Bhati and Mark Hildebrand and Mariano Tepper and Ted Willke},
74-
year = {2023}
75-
}
76-
```
77-
78-
<a id="1">[1]</a>
79-
Karpukhin, V.; Oguz, B.; Min, S.; Lewis, P.; Wu, L.; Edunov, S.; Chen, D.; Yih, W..: Dense Passage
80-
Retrieval for Open-Domain Question Answering. In: Proceedings of the 2020 Conference on Empirical
81-
Methods in Natural Language Processing (EMNLP). 6769–6781. (2020)
82-
83-
<a id="2">[2]</a>
84-
Raffel, C.; Shazeer, N.; Roberts, A.; Lee, K.; Narang, S.; Matena, M.; Zhou, Y.; Li, W.; Liu,
85-
P.J.: Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.
86-
In: The Journal of Machine Learning Research 21,140:1–140:67.(2020)
87-
88-
<a id="3">[3]</a>
89-
Aguerrebere, C.; Bhati I.; Hildebrand M.; Tepper M.; Willke T.:Similarity search in the blink of an eye with compressed
90-
indices. In: Proceedings of the VLDB Endowment, 16, 11, 3433 - 3446. (2023)
91-
92-
This "research quality code" is for Non-Commercial purposes provided by Intel "As Is" without any express or implied
93-
warranty of any kind. Please see the dataset's applicable license for terms and conditions. Intel does not own the
94-
rights to this data set and does not confer any rights to it. Intel does not warrant or assume responsibility for the accuracy or completeness of any information, text, graphics, links or other items within the code. A thorough security review has not been performed on this code. Additionally, this repository may contain components that are out of date or contain known security vulnerabilities.

‎dpr/README.md

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
# DPR Dataset Generator
2+
3+
This repository provides code to generate base and query vector datasets for similarity search benchmarking and evaluation on high-dimensional vectors stemming from large language models.
4+
With the dense passage retriever (DPR) [[1]](#1), we encode text snippets from the C4 dataset [[2]](#2) to generate 768-dimensional vectors:
5+
- context DPR embeddings for the base set and
6+
- question DPR embeddings for the query set.
7+
8+
The metric for similarity search is inner product [[1]](#1).
9+
10+
The number of base and query embedding vectors is parametrizable.
11+
12+
## DPR10M
13+
14+
A specific instance with 10 million base vectors and 10,000 query vectors is introduced in [[3]](#3). Use the script [dpr_dataset_10M.py](dpr_dataset_10M.py) to generate this dataset. The corresponding ground-truth (available [here](gtruth_dpr10M_innerProduct.ivecs)) is generated conducting an exhaustive search with the inner product
15+
metric.
16+
17+
Here is a summary of the **steps to generate this dataset**:
18+
19+
1. **Download the files** corresponding to the `en` variant of the C4 dataset accesible [here](https://huggingface.co/datasets/allenai/c4).
20+
The complete set of files requires 350GB of storage, so you might want to follow the instructions to download only a subset. For example, to generate 10M embeddings
21+
we used the first 2 files from the train set (i.e., files `c4-train.00000-of-01024.json.gz` and `c4-train.00001-of-01024.json.gz` in `c4/en/train`).
22+
23+
2. **Execute** the `generate_dpr_embeddings` function to generate a `.fvecs` file containing the new embeddings.
24+
Note that different settings should be used to generate the **base vectors** and the **query set**, as they use the
25+
DPR context and query encoders respectively.
26+
See the script [dpr_dataset_10M.py](dpr_dataset_10M.py) for details.
27+
28+
```
29+
# Example code to generate base vectors
30+
31+
from dpr_dataset_generate import generate_dpr_embeddings
32+
33+
base_C4_folder = '/home/username/research/datasets/c4/en' # Set this path to where your c4/en folder is located
34+
cache_folder = f'/home/username/.cache/huggingface/datasets/' # Set to the hugginface datasets cache path
35+
dataset_dir = f'{base_C4_folder}/train/'
36+
37+
num_embd = 10000000
38+
init_file = 0
39+
num_of_files = 2 # Make sure the input files (2 in this case) are enough to generate the
40+
# requested number of embeddings.
41+
# To get an estimate, use the optional parameter get_total_embeddings_only
42+
# to get the number of embeddings that can be generated from a certain
43+
# group of files without actually generating the embeddings.
44+
fname_prefix_out = 'c4-en'
45+
doc_stride = 32
46+
max_length = 64
47+
batch_size = 512
48+
dim = 768
49+
50+
generate_dpr_embeddings(init_file, num_of_files, num_embd, doc_stride, max_length, dim,
51+
batch_size,
52+
dataset_dir, fname_prefix_out, cache_folder)
53+
```
54+
3. **Generate the ground-truth** by conducting an exhaustive search with the inner product metric.
55+
We provide the [ground-truth](gtruth_dpr10M_innerProduct.ivecs) for the dataset generated using
56+
[dpr_dataset_10M.py](dpr_dataset_10M.py).
57+
58+
> **_NOTE:_** Due to floating-point arithmetic precision the vector embeddings generated using the provided
59+
> code in different machines may slightly vary. Keep in mind that this could cause small discrepancies with the provided ground-truth.
60+
61+
4. Functions `read_fvecs` and `read_ivecs` can be used to read `.fvecs` and `.ivecs` files respectively.
62+
63+
## References
64+
Reference to cite when you use datasets generated with this code in a research paper:
65+
66+
```
67+
@article{aguerrebere2023similarity,
68+
title={Similarity search in the blink of an eye with compressed indices},
69+
volume = {16},
70+
number = {11},
71+
pages = {3433--3446},
72+
journal = {Proceedings of the VLDB Endowment},
73+
author={Cecilia Aguerrebere and Ishwar Bhati and Mark Hildebrand and Mariano Tepper and Ted Willke},
74+
year = {2023}
75+
}
76+
```
77+
78+
<a id="1">[1]</a>
79+
Karpukhin, V.; Oguz, B.; Min, S.; Lewis, P.; Wu, L.; Edunov, S.; Chen, D.; Yih, W..: Dense Passage
80+
Retrieval for Open-Domain Question Answering. In: Proceedings of the 2020 Conference on Empirical
81+
Methods in Natural Language Processing (EMNLP). 6769–6781. (2020)
82+
83+
<a id="2">[2]</a>
84+
Raffel, C.; Shazeer, N.; Roberts, A.; Lee, K.; Narang, S.; Matena, M.; Zhou, Y.; Li, W.; Liu,
85+
P.J.: Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.
86+
In: The Journal of Machine Learning Research 21,140:1–140:67.(2020)
87+
88+
<a id="3">[3]</a>
89+
Aguerrebere, C.; Bhati I.; Hildebrand M.; Tepper M.; Willke T.:Similarity search in the blink of an eye with compressed
90+
indices. In: Proceedings of the VLDB Endowment, 16, 11, 3433 - 3446. (2023)
91+
92+
This "research quality code" is for Non-Commercial purposes provided by Intel "As Is" without any express or implied
93+
warranty of any kind. Please see the dataset's applicable license for terms and conditions. Intel does not own the
94+
rights to this data set and does not confer any rights to it. Intel does not warrant or assume responsibility for the accuracy or completeness of any information, text, graphics, links or other items within the code. A thorough security review has not been performed on this code. Additionally, this repository may contain components that are out of date or contain known security vulnerabilities.

‎dpr_dataset_10M.py renamed to ‎dpr/dpr_dataset_10M.py

File renamed without changes.

‎dpr_dataset_generate.py renamed to ‎dpr/dpr_dataset_generate.py

File renamed without changes.

‎gtruth_dpr10M_innerProduct.ivecs renamed to ‎dpr/gtruth_dpr10M_innerProduct.ivecs

File renamed without changes.

‎requirements.txt renamed to ‎dpr/requirements_dpr.txt

File renamed without changes.

‎wit/README.md

+89
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
# wit-512-1M Dataset Generator
2+
3+
This repository provides code to generate base and query (test and learn sets) embeddings for similarity search benchmarking
4+
and evaluation on high-dimensional vectors. The dataset is designed to benchmark similarity search methods under
5+
scenarios with out-of-distribution (OOD) queries stemming from a text-to-image application [[1]](#1).
6+
7+
The WIT dataset[[2]](#2) is a multimodal multilingual dataset that contains 37 million rich image-text examples
8+
extracted from Wikipedia pages. For each example in the first million training images
9+
(downloaded from [here](https://storage.cloud.google.com/wikimedia-image-caption-public/image_data_train.tar)), we take the image and encode it
10+
using the multimodal OpenAI CLIP-ViT-B32 model [[3]](#3) to generate a database vector.
11+
We create the query set using the first 20K text descriptions in one of the provided test sets (concatenating
12+
the Reference and Attribution description fields) and generating the corresponding embeddings using CLIPViT-B32-multilingual-v1 [[4]](#4).
13+
The use of CLIP-ViT-B32 for images and multi-lingual CLIP-ViT-B32-multilingual-v1 for text follows the protocol suggested
14+
[here]( https://huggingface.co/sentence-transformers/clip-ViT-B-32-multilingual-v1).
15+
Finally, for each query, we compute the 100 ground truth nearest neighbors using maximum inner product.
16+
We use the first 10K queries as a query test set and the remaining 10K as a learn set.
17+
18+
The metric for similarity search used with this dataset is inner product.
19+
20+
21+
## Steps to generate the dataset
22+
23+
The script [wit_dataset_1M.py](wit_dataset_1M.py) generates 1 million base vectors from the provided training images
24+
and two sets of queries, test and learn set, each with 10K vectors from the text descriptions of the provided test set.
25+
26+
Here is a summary of the **steps to generate this dataset**:
27+
28+
1. **Download the WIT training Images and test set**.
29+
We download the training images from [here](https://storage.cloud.google.com/wikimedia-image-caption-public/image_data_train.tar) and extract them in the desired location
30+
> **_NOTE:_** the above link requires Google login authentication to download the training images.
31+
32+
```
33+
tar -xvf image_data_train.tar -C $BASE_PATH
34+
```
35+
The extracted files will be in $BASE_PATH/image_data_train/image_pixels/.
36+
The images are encoded in base64 format, see image/file format details [here](https://www.kaggle.com/c/wikipedia-image-caption/data).
37+
38+
For queries, we download one of the [test set](https://storage.googleapis.com/gresearch/wit/wit_v1.test.all-00000-of-00005.tsv.gz) and extract it
39+
```
40+
mkdir -p $BASE_PATH/test_set
41+
tar -xvfz wit_v1.test.all-00000-of-00005.tsv.gz -C $BASE_PATH/test_set
42+
```
43+
44+
2. **Run** the `wit_dataset_1M.py` script to generate `.fvecs` files containing the base
45+
, query and learn set vectors. **Remember to set the path** to the folder where the
46+
downloaded training images, test files are located.
47+
48+
3. **Generate the ground-truth** by conducting an exhaustive search with the inner product metric.
49+
We provide the ground-truth files `wit_test_gt_1M_innerproduct.ivecs` and `wit_learn_gt_1M_innerproduct.ivecs` for the test and learn sets, respectively.
50+
51+
4. Functions `read_fvecs` and `read_ivecs` can be used to read `.fvecs` and `.ivecs` files respectively.
52+
53+
> **_NOTE:_** Due to floating-point arithmetic precision the vector embeddings generated using the provided
54+
> code in different machines may slightly vary. Keep in mind that this could cause small discrepancies with the provided ground-truth.
55+
56+
57+
## References
58+
Reference to cite when you use datasets generated with this code in a research paper:
59+
60+
```
61+
@article{tepper2023leanvec,
62+
title={LeanVec: Search your vectors faster by making them fit},
63+
author={Mariano Tepper and Ishwar Singh Bhati and Cecilia Aguerrebere and Mark Hildebrand and Ted Willke},
64+
year={2023},
65+
journal={arXiv},
66+
doi={https://doi.org/10.48550/arXiv.2312.16335}
67+
}
68+
```
69+
<a id="1">[1]</a>
70+
Mariano Tepper, Ishwar Singh Bhati, Cecilia Aguerrebere, Mark Hildebrand, and Ted Willke:
71+
LeanVec: Search your vectors faster by making them fit. (2023)
72+
73+
<a id="2">[2]</a>
74+
Krishna Srinivasan, Karthik Raman, Jiecao Chen, Michael Bendersky, and Marc Najork:
75+
WIT: Wikipedia-based Image Text Dataset for Multimodal Multilingual Machine Learning. (2021)
76+
77+
<a id="3">[3]</a>
78+
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry,
79+
Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever:
80+
Learning Transferable Visual Models From Natural Language Supervision. (2021)
81+
82+
<a id="4">[4]</a>
83+
Nils Reimers, and Iryna Gurevych:
84+
Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation. (2020)
85+
86+
This "research quality code" is for Non-Commercial purposes provided by Intel "As Is" without any express or implied
87+
warranty of any kind. Please see the dataset's applicable license for terms and conditions. Intel does not own the
88+
rights to this data set and does not confer any rights to it. Intel does not warrant or assume responsibility for the accuracy or completeness of any information, text, graphics, links or other items within the code. A thorough security review has not been performed on this code. Additionally, this repository may contain components that are out of date or contain known security vulnerabilities.
89+

‎wit/requirements_wit.txt

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
sentence-transformers==2.2.2
2+
Pillow==9.4.0
3+
numpy==1.26.2
4+
typing_extensions==4.5.0
5+
natsort==8.4.0

‎wit/wit_dataset_1M.py

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
from wit_dataset_generate import generate_image_embeddings, generate_text_embeddings
2+
3+
4+
if __name__ == "__main__":
5+
# This is an example script to generate the wit-512-1M dataset, containing:
6+
#
7+
# --- 1M base vectors from image embeddings
8+
# --- out-of-distribution (OOD) query and learning sets (10k vectors each) from text embeddings
9+
#
10+
# Introduced in the paper ""LeanVec: Search your vectors faster by making them fit", 2023,
11+
# Tepper, Bhati, Aguerrebere, Hildebrand, Willke (https://arxiv.org/abs/2312.16335)
12+
#
13+
#
14+
# This dataset is created using a subset of Google's multimodal multilingual WIT dataset,
15+
# using image-text examples extracted from Wikipedia pages (https://github.com/google-research-datasets/wit).
16+
# To generate a base vector, we take the image and encode it using OpenAI CLIP-ViT-B32 model.
17+
# For queries, we use text descriptions in one of the provided test sets
18+
# (concatenating the Reference and Attribution description fields) and generating the corresponding
19+
# embeddings using OpenAI CLIP-ViT-B32-multilingual-v1. We followed the steps suggested in
20+
# https://huggingface.co/sentence-transformers/clip-ViT-B-32-multilingual-v1
21+
#
22+
# See the README for more details.
23+
#
24+
# Please see the documentation of the generate_image_embeddings and generate_text_embeddings functions
25+
# for details on the required parameters.
26+
27+
28+
base_path = '/raid0/ishwarsi/datasets/wit' # Path to the location of the WIT Images/Test datasets are located
29+
images_dir = f'{base_path}/image_data_train/image_pixels' # Images files directory storing *.csv.gz files
30+
output_dir = f'{base_path}/output' # Directory where the created dataset will be saved
31+
32+
# Files saved in [output_dir]/embeddings/{fname_prefix}.fvecs'
33+
34+
# Generate image embeddings from the images provided in the images_dir
35+
fname_prefix = 'wit_base_1M'
36+
num_vecs = 1000_000
37+
generate_image_embeddings(images_dir, num_vecs, output_dir, fname_prefix)
38+
39+
test_file = f'{base_path}/test_set/wit_v1.test.all-00000-of-00005.tsv'
40+
41+
# Generate text embeddings from the test file
42+
fname_prefix = 'wit_query_10k'
43+
num_vecs = 10_000
44+
generate_text_embeddings(test_file, num_vecs, output_dir, fname_prefix)
45+
46+
# Learn queries start from an offset (the last parameter)
47+
fname_prefix = 'wit_learn_query_10k'
48+
num_vecs = 10_000
49+
num_skip = 10_000
50+
generate_text_embeddings(test_file, num_vecs, output_dir, fname_prefix, num_skip)

0 commit comments

Comments
 (0)