README.md

---
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 ](https://arxiv.org/pdf/2309.07597)and [BAAI's Hugging Face page](https://huggingface.co/BAAI) for full training details.
## Links

* [BGE Quickstart Notebook](coming soon..)
* [BGE API Documentation](https://keras.io/keras_hub/api/models/bert/)
* [BGE Model Card](https://huggingface.co/BAAI/bge-small-en)
* [Original Paper](https://arxiv.org/pdf/2309.07597)
* [BGE](https://huggingface.co/BAAI)
* [KerasHub Beginner Guide](https://keras.io/guides/keras_hub/getting_started/)
* [KerasHub Model Publishing Guide](https://keras.io/guides/keras_hub/upload/)

## 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](https://keras.io/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_llm_embedder")

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

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


```