File size: 5,484 Bytes
06356cc 5a80bb9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | ---
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]}...")
```
|