Skip to content

fix: Added way more training dataset annotations #1765

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
KennethEnevoldsen merged 19 commits into from
Jan 17, 2025
Merged

fix: Added way more training dataset annotations #1765

Show file tree
Hide file tree
Changes from 16 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
a168496
fix: Leaderboard: `K` instead of `M`
KennethEnevoldsen Jan 11, 2025
e61d7f2
format
KennethEnevoldsen Jan 11, 2025
e1b89e3
fixed existing annotations to refer to task name instead of hf dataset
KennethEnevoldsen Jan 11, 2025
9ffeae4
added annotation to nvidia
KennethEnevoldsen Jan 11, 2025
0495d32
added voyage
KennethEnevoldsen Jan 11, 2025
5f7ef65
added uae annotations
KennethEnevoldsen Jan 11, 2025
ac48012
Added stella annotations
KennethEnevoldsen Jan 11, 2025
c1c7eb6
sentence trf models
KennethEnevoldsen Jan 11, 2025
4ec9121
added salesforce and e5
KennethEnevoldsen Jan 11, 2025
c54859d
jina
KennethEnevoldsen Jan 11, 2025
d7f5684
bge + model2vec
KennethEnevoldsen Jan 11, 2025
9ea60ff
added llm2vec annotations
KennethEnevoldsen Jan 11, 2025
b123d92
add jasper
KennethEnevoldsen Jan 11, 2025
aa728d1
format
KennethEnevoldsen Jan 11, 2025
87b5d9d
Merge remote-tracking branch 'origin' into add-more-annotations
KennethEnevoldsen Jan 11, 2025
121bf0e
format
KennethEnevoldsen Jan 12, 2025
b2b9cca
Updated annotations and moved jina models
KennethEnevoldsen Jan 13, 2025
3aab7ec
fix: add even more training dataset annotations (#1793)
KennethEnevoldsen Jan 15, 2025
580b3b5
Merge remote-tracking branch 'origin' into add-more-annotations
KennethEnevoldsen Jan 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 33 additions & 87 deletions mteb/models/bge_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,36 @@

model_prompts = {"query": "Represent this sentence for searching relevant passages: "}

bge_training_data = {
# source: https://data.baai.ac.cn/details/BAAI-MTP
"NQ": ["test"],
"NQHardNegatives": ["test"],
"AmazonReviewsClassification": [
"validation",
"test",
], # assumed from: amazon_reviews_multi
"MLQARetrieval": [
"validation",
"test",
], # assumed from mlqa (question, context)
# not in mteb
# Dataset Pairs
# wudao (title, passage)
# cmrc2018 (query, context)
# dureader (query, context)
# simclue (sentence_a, sentence_b)
# csl (title, abstract)
# amazon_reviews_multi (title, body)
# wiki_atomic_edits (base_sentence, edited_sentence)
# mlqa (question, context)
# xlsum (title, summary) (title, text)
# "sentence-transformers data": [], # https://huggingface.co/datasets/sentence-transformers/embedding-training-data # TODO check this further
# "wikipedia": [], # title + section title, passage
# "reddit": [], # title, body
# "stackexchange": [], # (title, upvoted answer) (title+body, upvoted answer)
# "s2orc": [], # (title, abstract) (title, citation title) (abstract, citation abstract)
}

bge_small_en_v1_5 = ModelMeta(
loader=partial( # type: ignore
sentence_transformers_loader,
Expand All @@ -29,35 +59,7 @@
use_instructions=True,
public_training_data=True, # https://data.baai.ac.cn/details/BAAI-MTP
public_training_code=None, # seemingly released (at least for some models, but the link is broken
training_datasets={
# source: https://data.baai.ac.cn/details/BAAI-MTP
"NQ": ["test"],
"NQHardNegatives": ["test"],
"AmazonReviewsClassification": [
"validation",
"test",
], # assumed from: amazon_reviews_multi
"MLQARetrieval": [
"validation",
"test",
], # assumed from mlqa (question, context)
# not in mteb
# Dataset Pairs
# wudao (title, passage)
# cmrc2018 (query, context)
# dureader (query, context)
# simclue (sentence_a, sentence_b)
# csl (title, abstract)
# amazon_reviews_multi (title, body)
# wiki_atomic_edits (base_sentence, edited_sentence)
# mlqa (question, context)
# xlsum (title, summary) (title, text)
# "sentence-transformers data": [], # https://huggingface.co/datasets/sentence-transformers/embedding-training-data # TODO check this further
# "wikipedia": [], # title + section title, passage
# "reddit": [], # title, body
# "stackexchange": [], # (title, upvoted answer) (title+body, upvoted answer)
# "s2orc": [], # (title, abstract) (title, citation title) (abstract, citation abstract)
},
training_datasets=bge_training_data,
)

bge_base_en_v1_5 = ModelMeta(
Expand All @@ -83,35 +85,7 @@
use_instructions=True,
public_training_data=True, # https://data.baai.ac.cn/details/BAAI-MTP
public_training_code=None, # seemingly released (at least for some models, but the link is broken
training_datasets={
# source: https://data.baai.ac.cn/details/BAAI-MTP
"NQ": ["test"],
"NQHardNegatives": ["test"],
"AmazonReviewsClassification": [
"validation",
"test",
], # assumed from: amazon_reviews_multi
"MLQARetrieval": [
"validation",
"test",
], # assumed from mlqa (question, context)
# not in mteb
# Dataset Pairs
# wudao (title, passage)
# cmrc2018 (query, context)
# dureader (query, context)
# simclue (sentence_a, sentence_b)
# csl (title, abstract)
# amazon_reviews_multi (title, body)
# wiki_atomic_edits (base_sentence, edited_sentence)
# mlqa (question, context)
# xlsum (title, summary) (title, text)
# "sentence-transformers data": [], # https://huggingface.co/datasets/sentence-transformers/embedding-training-data # TODO check this further
# "wikipedia": [], # title + section title, passage
# "reddit": [], # title, body
# "stackexchange": [], # (title, upvoted answer) (title+body, upvoted answer)
# "s2orc": [], # (title, abstract) (title, citation title) (abstract, citation abstract)
},
training_datasets=bge_training_data,
)

bge_large_en_v1_5 = ModelMeta(
Expand All @@ -137,33 +111,5 @@
use_instructions=True,
public_training_data=True, # https://data.baai.ac.cn/details/BAAI-MTP
public_training_code=None, # seemingly released (at least for some models, but the link is broken
training_datasets={
# source: https://data.baai.ac.cn/details/BAAI-MTP
"NQ": ["test"],
"NQHardNegatives": ["test"],
"AmazonReviewsClassification": [
"validation",
"test",
], # assumed from: amazon_reviews_multi
"MLQARetrieval": [
"validation",
"test",
], # assumed from mlqa (question, context)
# not in mteb
# Dataset Pairs
# wudao (title, passage)
# cmrc2018 (query, context)
# dureader (query, context)
# simclue (sentence_a, sentence_b)
# csl (title, abstract)
# amazon_reviews_multi (title, body)
# wiki_atomic_edits (base_sentence, edited_sentence)
# mlqa (question, context)
# xlsum (title, summary) (title, text)
# "sentence-transformers data": [], # https://huggingface.co/datasets/sentence-transformers/embedding-training-data # TODO check this further
# "wikipedia": [], # title + section title, passage
# "reddit": [], # title, body
# "stackexchange": [], # (title, upvoted answer) (title+body, upvoted answer)
# "s2orc": [], # (title, abstract) (title, citation title) (abstract, citation abstract)
},
training_datasets=bge_training_data,
)
8 changes: 7 additions & 1 deletion mteb/models/e5_instruct.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

from mteb.model_meta import ModelMeta

from .e5_models import E5_PAPER_RELEASE_DATE, XLMR_LANGUAGES
from .e5_models import E5_PAPER_RELEASE_DATE, E5_TRAINING_DATA, XLMR_LANGUAGES
from .instruct_wrapper import instruct_wrapper

MISTRAL_LANGUAGES = ["eng_Latn", "fra_Latn", "deu_Latn", "ita_Latn", "spa_Latn"]
Expand Down Expand Up @@ -40,6 +40,9 @@
embed_dim=1024,
license="mit",
max_tokens=514,
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_mistral = ModelMeta(
Expand Down Expand Up @@ -69,4 +72,7 @@
embed_dim=4096,
license="mit",
max_tokens=32768,
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)
148 changes: 40 additions & 108 deletions mteb/models/e5_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,19 @@
PromptType.passage.value: "passage: ",
}

E5_TRAINING_DATA = {
# from 4.2 in https://arxiv.org/pdf/2212.03533
# also pre-training data from a variety of sources (stackexchange, semantic scholar, reddit, CC, ...)
"MSMARCO": ["train"],
"MSMARCOHardNegatives": ["train"],
"NanoMSMARCORetrieval": ["train"],
"MSMARCO-PL": ["train"], # translation not trained on
"NQ": ["train"],
"NQHardNegatives": ["train"],
"NanoNQRetrieval": ["train"],
"NQ-PL": ["train"], # translation not trained on
}

e5_mult_small = ModelMeta(
loader=partial( # type: ignore
sentence_transformers_loader,
Expand All @@ -134,26 +147,9 @@
similarity_fn_name="cosine",
framework=["Sentence Transformers", "PyTorch"],
use_instructions=True,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2212.03533
# table 1:
# Wikipedia 150M
# mC4 160M
# Multilingual CC News 160M
# NLLB 160M
# Reddit 160M
# S2ORC 50M
# Stackexchange 50M
# xP3 80M
# Misc. SBERT Data 10M
# ----
# from Misc. SBERT Data 10M:
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_mult_base = ModelMeta(
Expand All @@ -176,26 +172,9 @@
similarity_fn_name="cosine",
framework=["Sentence Transformers", "PyTorch"],
use_instructions=True,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2402.05672
# table 1:
# Wikipedia 150M
# mC4 160M
# Multilingual CC News 160M
# NLLB 160M
# Reddit 160M
# S2ORC 50M
# Stackexchange 50M
# xP3 80M
# Misc. SBERT Data 10M
# ----
# from Misc. SBERT Data 10M:
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_mult_large = ModelMeta(
Expand All @@ -219,26 +198,9 @@
similarity_fn_name="cosine",
framework=["Sentence Transformers", "PyTorch"],
use_instructions=True,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2402.05672
# table 1:
# Wikipedia 150M
# mC4 160M
# Multilingual CC News 160M
# NLLB 160M
# Reddit 160M
# S2ORC 50M
# Stackexchange 50M
# xP3 80M
# Misc. SBERT Data 10M
# ----
# from Misc. SBERT Data 10M:
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_eng_small_v2 = ModelMeta(
Expand All @@ -261,14 +223,9 @@
similarity_fn_name="cosine",
framework=["Sentence Transformers", "PyTorch"],
use_instructions=True,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2212.03533
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_eng_small = ModelMeta(
Expand All @@ -292,14 +249,9 @@
similarity_fn_name="cosine",
framework=["Sentence Transformers", "PyTorch"],
use_instructions=True,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2212.03533
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_eng_base_v2 = ModelMeta(
Expand All @@ -325,14 +277,9 @@
use_instructions=True,
superseded_by=None,
adapted_from=None,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2212.03533
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_eng_large_v2 = ModelMeta(
Expand All @@ -358,14 +305,9 @@
use_instructions=True,
superseded_by=None,
adapted_from=None,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2212.03533
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_large = ModelMeta(
Expand All @@ -391,14 +333,9 @@
use_instructions=True,
superseded_by="intfloat/e5-large-v2",
adapted_from=None,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2212.03533
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)

e5_base = ModelMeta(
Expand All @@ -424,12 +361,7 @@
use_instructions=True,
superseded_by="intfloat/e5-base-v2",
adapted_from=None,
public_training_data=False, # couldn't find
public_training_code=False, # couldn't find
training_datasets={
# source: https://arxiv.org/pdf/2212.03533
"NQ": ["test"],
"NQHardNegatives": ["test"],
"MSMARCO": ["train"], # dev?
},
public_training_data=False,
public_training_code=False,
training_datasets=E5_TRAINING_DATA,
)
Loading