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