library_name: keras-hub
Model Overview
BGE
BGE (BAAI General Embedding) models for dense text retrieval and semantic similarity tasks, implemented in Keras.
Model Overview
BGE (BAAI General Embedding) is a family of bi-directional, transformer-based text embedding models developed by the Beijing Academy of Artificial Intelligence (BAAI). Built on the BERT encoder architecture, BGE models are fine-tuned specifically for dense retrieval, semantic similarity, and clustering tasks.
For embedding generation, the model outputs L2-normalized embeddings of the [CLS] token's hidden state, producing fixed-dimensional dense vectors suitable for cosine similarity comparisons.
These models can be used with KerasHub through the BgeTextEmbedder task API.
Architecture
BGE models follow the standard BERT encoder architecture:
Tokenizer: WordPiece tokenizer with BERT-compatible special tokens ([CLS], [SEP], [PAD]).
Encoder: Multi-layer bi-directional Transformer encoder.
Embedding output: L2-normalized [CLS] token hidden states.
Intended Use
- Semantic search and information retrieval
- Document similarity and clustering
- Retrieval-Augmented Generation (RAG) pipelines
- Question-answer matching
Training Data
BGE models are trained on large-scale text pair datasets for contrastive learning. See the original paper and BAAI's Hugging Face page for full training details.
Links
- [BGE Quickstart Notebook](coming soon..)
- BGE API Documentation
- BGE Model Card
- Original Paper
- BGE
- KerasHub Beginner Guide
- KerasHub Model Publishing Guide
Installation
Keras and KerasHub can be installed with:
pip install -U -q keras-hub
pip install -U -q keras
Jax, TensorFlow, and Torch come preinstalled in Kaggle Notebooks. For instructions on installing them in another environment see the Keras Getting Started page.
Presets
| Preset | Architecture | Pooling | Normalize | Languages |
|---|---|---|---|---|
bge_small_en |
BERT | CLS | L2 | English |
bge_base_en |
BERT | CLS | L2 | English |
bge_large_en |
BERT | CLS | L2 | English |
bge_small_v1.5_en |
BERT | CLS | L2 | English |
bge_base_v1.5_en |
BERT | CLS | L2 | English |
bge_large_v1.5_en |
BERT | CLS | L2 | English |
bge_base_zh |
BERT | CLS | L2 | Chinese |
bge_large_zh |
BERT | CLS | L2 | Chinese |
bge_small_v1.5_zh |
BERT | CLS | L2 | Chinese |
bge_base_v1.5_zh |
BERT | CLS | L2 | Chinese |
bge_large_v1.5_zh |
BERT | CLS | L2 | Chinese |
bge_llm_embedder |
BERT | CLS | L2 | English |
bge_m3 |
XLM-RoBERTa | CLS | L2 | 100+ |
Example Usage
# Install and setup
!pip install -q keras-hub
import os
os.environ["KERAS_BACKEND"] = "jax" # or "tensorflow" or "torch"
import keras_hub
import numpy as np
# Load a BGE model from the Kaggle preset
embedder = keras_hub.models.BertTextEmbedder.from_preset("bge_large_zh")
# Encode text into embeddings
embeddings = embedder.encode_text(["The weather is lovely today."])
print(f"Shape: {embeddings.shape}") # (1, 384)
# Compute similarity between sentences
query = ["What is deep learning?"]
passages = [
"Deep learning is a subset of machine learning using neural networks with many layers.",
"The Eiffel Tower is located in Paris, France.",
"Neural networks learn representations of data through backpropagation.",
]
# Encode the texts into embeddings before passing to similarity
query_embeddings = embedder.encode_text(query)
passage_embeddings = embedder.encode_documents(passages)
# Calculate similarity
scores = embedder.similarity(query_embeddings, passage_embeddings)
print("Similarity scores:")
# Access the first row of scores [0] since we have 1 query
for passage, score in zip(passages, np.array(scores)[0]):
print(f" {float(score):.4f} → {passage[:60]}...")
Example Usage with Hugging Face URI
# Install and setup
!pip install -q keras-hub
import os
os.environ["KERAS_BACKEND"] = "jax" # or "tensorflow" or "torch"
import keras_hub
import numpy as np
# Load a BGE model from the Kaggle preset
embedder = keras_hub.models.BertTextEmbedder.from_preset("hf://keras/bge_large_zh")
# Encode text into embeddings
embeddings = embedder.encode_text(["The weather is lovely today."])
print(f"Shape: {embeddings.shape}") # (1, 384)
# Compute similarity between sentences
query = ["What is deep learning?"]
passages = [
"Deep learning is a subset of machine learning using neural networks with many layers.",
"The Eiffel Tower is located in Paris, France.",
"Neural networks learn representations of data through backpropagation.",
]
# Encode the texts into embeddings before passing to similarity
query_embeddings = embedder.encode_text(query)
passage_embeddings = embedder.encode_documents(passages)
# Calculate similarity
scores = embedder.similarity(query_embeddings, passage_embeddings)
print("Similarity scores:")
# Access the first row of scores [0] since we have 1 query
for passage, score in zip(passages, np.array(scores)[0]):
print(f" {float(score):.4f} → {passage[:60]}...")