bge_small_en

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

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_small_en")

# 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_small_en")

# 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]}...")