Skip to content

Commit b7f1b57

Browse files
vjdaburdevvijay-dabur
vjdaburdev
and
vijay-dabur
authored
Add BGE Embedding Model to Hub (#2755)
* Register bge-small-en-v1.5 Model * Updated formatting for pylint * Removed BGE Model directory and refer BERT model files in conversion * Added preset changes and updated conversion script * Corrected review comments from code assist * Fix gemini code assist review comments * Fix gemini code assist review comments * Fix review comments to get preprosessor config * Update code to support xlm_roberta architecture and merge changes from multilingual model * Fix Gemini Code assist review comments * Fix review comments * Fix review comments for preset map and gpu test --------- Co-authored-by: Vijay <vijay_dabur@epam.com>
1 parent 852625a commit b7f1b57

13 files changed

Lines changed: 1637 additions & 1 deletion

‎keras_hub/api/models/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -937,6 +937,12 @@
937937
from keras_hub.src.models.xlm_roberta.xlm_roberta_text_classifier_preprocessor import (
938938
XLMRobertaTextClassifierPreprocessor as XLMRobertaTextClassifierPreprocessor,
939939
)
940+
from keras_hub.src.models.xlm_roberta.xlm_roberta_text_embedder import (
941+
XLMRobertaTextEmbedder as XLMRobertaTextEmbedder,
942+
)
943+
from keras_hub.src.models.xlm_roberta.xlm_roberta_text_embedder_preprocessor import (
944+
XLMRobertaTextEmbedderPreprocessor as XLMRobertaTextEmbedderPreprocessor,
945+
)
940946
from keras_hub.src.models.xlm_roberta.xlm_roberta_tokenizer import (
941947
XLMRobertaTokenizer as XLMRobertaTokenizer,
942948
)

‎keras_hub/src/models/bert/bert_presets.py‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,4 +227,165 @@
227227
},
228228
"kaggle_handle": "kaggle://keras/sentence-transformers/keras/msmarco_minilm_l12_cos_v5_en/1",
229229
},
230+
# BGE family: BAAI General Embedding models optimized for dense retrieval.
231+
"bge_small_en": {
232+
"metadata": {
233+
"description": (
234+
"12-layer BGE small English embedding model (v1). Maps "
235+
"sentences to 384-dimensional L2-normalized dense vectors. "
236+
"Optimized for dense retrieval and semantic similarity."
237+
),
238+
"params": 33360000,
239+
"path": "bert",
240+
},
241+
"kaggle_handle": "kaggle://keras/bge/keras/bge_small_en/1",
242+
},
243+
"bge_base_en": {
244+
"metadata": {
245+
"description": (
246+
"12-layer BGE base English embedding model (v1). Maps "
247+
"sentences to 768-dimensional L2-normalized dense vectors. "
248+
"Optimized for dense retrieval and semantic similarity."
249+
),
250+
"params": 109482240,
251+
"path": "bert",
252+
},
253+
"kaggle_handle": "kaggle://keras/bge/keras/bge_base_en/1",
254+
},
255+
"bge_large_en": {
256+
"metadata": {
257+
"description": (
258+
"24-layer BGE large English embedding model (v1). Maps "
259+
"sentences to 1024-dimensional L2-normalized dense vectors. "
260+
"Highest accuracy in the BGE English v1 family."
261+
),
262+
"params": 335141888,
263+
"path": "bert",
264+
},
265+
"kaggle_handle": "kaggle://keras/bge/keras/bge_large_en/1",
266+
},
267+
"bge_small_v1.5_en": {
268+
"metadata": {
269+
"description": (
270+
"12-layer BGE small English embedding model (v1.5). Maps "
271+
"sentences to 384-dimensional L2-normalized dense vectors. "
272+
"Optimized for dense retrieval and semantic similarity."
273+
),
274+
"params": 33360000,
275+
"path": "bert",
276+
},
277+
"kaggle_handle": "kaggle://keras/bge/keras/bge_small_v1.5_en/1",
278+
},
279+
"bge_base_v1.5_en": {
280+
"metadata": {
281+
"description": (
282+
"12-layer BGE base English embedding model (v1.5). Maps "
283+
"sentences to 768-dimensional L2-normalized dense vectors. "
284+
"Optimized for dense retrieval and semantic similarity."
285+
),
286+
"params": 109482240,
287+
"path": "bert",
288+
},
289+
"kaggle_handle": "kaggle://keras/bge/keras/bge_base_v1.5_en/1",
290+
},
291+
"bge_large_v1.5_en": {
292+
"metadata": {
293+
"description": (
294+
"24-layer BGE large English embedding model (v1.5). Maps "
295+
"sentences to 1024-dimensional L2-normalized dense vectors. "
296+
"Highest accuracy in the BGE English family."
297+
),
298+
"params": 335141888,
299+
"path": "bert",
300+
},
301+
"kaggle_handle": "kaggle://keras/bge/keras/bge_large_v1.5_en/1",
302+
},
303+
"bge_base_zh": {
304+
"metadata": {
305+
"description": (
306+
"12-layer BGE base Chinese embedding model (v1). Maps "
307+
"sentences to 768-dimensional L2-normalized dense vectors. "
308+
"Optimized for dense retrieval on Chinese text."
309+
),
310+
"params": 102267648,
311+
"path": "bert",
312+
},
313+
"kaggle_handle": "kaggle://keras/bge/keras/bge_base_zh/1",
314+
},
315+
"bge_large_zh": {
316+
"metadata": {
317+
"description": (
318+
"24-layer BGE large Chinese embedding model (v1). Maps "
319+
"sentences to 1024-dimensional L2-normalized dense vectors. "
320+
"Highest accuracy in the BGE Chinese v1 family."
321+
),
322+
"params": 325522432,
323+
"path": "bert",
324+
},
325+
"kaggle_handle": "kaggle://keras/bge/keras/bge_large_zh/1",
326+
},
327+
"bge_small_v1.5_zh": {
328+
"metadata": {
329+
"description": (
330+
"12-layer BGE small Chinese embedding model (v1.5). Maps "
331+
"sentences to 384-dimensional L2-normalized dense vectors. "
332+
"Optimized for dense retrieval on Chinese text."
333+
),
334+
"params": 23953920,
335+
"path": "bert",
336+
},
337+
"kaggle_handle": "kaggle://keras/bge/keras/bge_small_v1.5_zh/1",
338+
},
339+
"bge_base_v1.5_zh": {
340+
"metadata": {
341+
"description": (
342+
"12-layer BGE base Chinese embedding model (v1.5). Maps "
343+
"sentences to 768-dimensional L2-normalized dense vectors. "
344+
"Optimized for dense retrieval on Chinese text."
345+
),
346+
"params": 102267648,
347+
"path": "bert",
348+
},
349+
"kaggle_handle": "kaggle://keras/bge/keras/bge_base_v1.5_zh/1",
350+
},
351+
"bge_large_v1.5_zh": {
352+
"metadata": {
353+
"description": (
354+
"24-layer BGE large Chinese embedding model (v1.5). Maps "
355+
"sentences to 1024-dimensional L2-normalized dense vectors. "
356+
"Highest accuracy in the BGE Chinese v1.5 family."
357+
),
358+
"params": 325522432,
359+
"path": "bert",
360+
},
361+
"kaggle_handle": "kaggle://keras/bge/keras/bge_large_v1.5_zh/1",
362+
},
363+
"bge_llm_embedder": {
364+
"metadata": {
365+
"description": (
366+
"BGE-LLM-Embedder: 12-layer embedding model for "
367+
"retrieval-augmented language model applications. Maps text to "
368+
"768-dimensional dense vectors and supports knowledge, memory, "
369+
"demonstration, and tool retrieval tasks."
370+
),
371+
"params": 109482240,
372+
"path": "bert",
373+
},
374+
"kaggle_handle": "kaggle://keras/bge/keras/bge_llm_embedder/1",
375+
},
376+
# "multilingual-e5-*" family: multilingual dense retrieval models.
377+
"multilingual_e5_small": {
378+
"metadata": {
379+
"description": (
380+
"12-layer multilingual E5 embedding model with 384-dimensional "
381+
"vectors. Fine-tuned for dense retrieval across 100+ languages "
382+
"using weakly-supervised contrastive pre-training. "
383+
"Prefix inputs with 'query: ' for queries and 'passage: ' "
384+
"for documents."
385+
),
386+
"params": 117653760,
387+
"path": "bert",
388+
},
389+
"kaggle_handle": "kaggle://keras/multilingual-e5/keras/multilingual_e5_small/1",
390+
},
230391
}

‎keras_hub/src/models/xlm_roberta/xlm_roberta_presets.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,45 @@
2323
},
2424
"kaggle_handle": "kaggle://keras/xlm_roberta/keras/xlm_roberta_large_multi/3",
2525
},
26+
"bge_m3": {
27+
"metadata": {
28+
"description": (
29+
"568M-parameter multilingual text embedding model supporting "
30+
"100+ languages with sequences up to 8192 tokens. Uses CLS "
31+
"token pooling with L2 normalization. Supports dense, sparse, "
32+
"and multi-vector (ColBERT-style) retrieval. From BAAI."
33+
),
34+
"params": 566702080,
35+
"path": "xlm_roberta",
36+
},
37+
"kaggle_handle": "kaggle://keras/bge/keras/bge_m3/1",
38+
},
39+
"multilingual_e5_base": {
40+
"metadata": {
41+
"description": (
42+
"12-layer multilingual E5 embedding model with 768-dimensional "
43+
"vectors. Fine-tuned for dense retrieval across 100+ languages "
44+
"using weakly-supervised contrastive pre-training. "
45+
"Prefix inputs with 'query: ' for queries and 'passage: ' "
46+
"for documents."
47+
),
48+
"params": 277450752,
49+
"path": "xlm_roberta",
50+
},
51+
"kaggle_handle": "kaggle://keras/multilingual-e5/keras/multilingual_e5_base/1",
52+
},
53+
"multilingual_e5_large": {
54+
"metadata": {
55+
"description": (
56+
"24-layer multilingual E5 embedding model with "
57+
"1024-dimensional vectors. Fine-tuned for dense retrieval "
58+
"across 100+ languages using weakly-supervised contrastive "
59+
"pre-training. Prefix inputs with 'query: ' for queries "
60+
"and 'passage: ' for documents."
61+
),
62+
"params": 558837760,
63+
"path": "xlm_roberta",
64+
},
65+
"kaggle_handle": "kaggle://keras/multilingual-e5/keras/multilingual_e5_large/1",
66+
},
2667
}

0 commit comments

Comments
 (0)