|
| 1 | +from keras import ops |
| 2 | + |
| 3 | +from keras_hub.src.api_export import keras_hub_export |
| 4 | +from keras_hub.src.models.bert.bert_backbone import BertBackbone |
| 5 | +from keras_hub.src.models.bert.bert_text_embedder_preprocessor import ( |
| 6 | + BertTextEmbedderPreprocessor, |
| 7 | +) |
| 8 | +from keras_hub.src.models.text_embedder import TextEmbedder |
| 9 | + |
| 10 | + |
| 11 | +@keras_hub_export("keras_hub.models.BertTextEmbedder") |
| 12 | +class BertTextEmbedder(TextEmbedder): |
| 13 | + """An end-to-end BERT model for generating sentence embeddings. |
| 14 | +
|
| 15 | + This model attaches a mean pooling and L2 normalization head to a |
| 16 | + `keras_hub.models.BertBackbone` instance, mapping from the backbone |
| 17 | + outputs to fixed-size sentence embeddings suitable for semantic |
| 18 | + similarity, clustering, and retrieval tasks. |
| 19 | +
|
| 20 | + This is the architecture used by sentence-transformers models like |
| 21 | + `all-MiniLM-L6-v2`. For usage of this model with pre-trained weights, |
| 22 | + use the `from_preset()` constructor. |
| 23 | +
|
| 24 | + This model can optionally be configured with a `preprocessor` layer, in |
| 25 | + which case it will automatically apply preprocessing to raw inputs during |
| 26 | + `fit()`, `predict()`, and `evaluate()`. This is done by default when |
| 27 | + creating the model with `from_preset()`. |
| 28 | +
|
| 29 | + Disclaimer: Pre-trained models are provided on an "as is" basis, without |
| 30 | + warranties or conditions of any kind. |
| 31 | +
|
| 32 | + Args: |
| 33 | + backbone: A `keras_hub.models.BertBackbone` instance. |
| 34 | + preprocessor: A `keras_hub.models.BertTextEmbedderPreprocessor` or |
| 35 | + `None`. If `None`, this model will not apply preprocessing, and |
| 36 | + inputs should be preprocessed before calling the model. |
| 37 | + pooling_mode: str. The pooling strategy to use. One of `"mean"`, |
| 38 | + `"cls"`, or `"max"`. Defaults to `"mean"`. |
| 39 | + - `"mean"`: Attention-mask-aware mean pooling over all tokens. |
| 40 | + - `"cls"`: Use the `[CLS]` token representation. |
| 41 | + - `"max"`: Max pooling over all tokens. |
| 42 | + normalize: bool. Whether to L2 normalize the output embeddings. |
| 43 | + Defaults to `True`. |
| 44 | +
|
| 45 | + Examples: |
| 46 | +
|
| 47 | + Raw string data. |
| 48 | + ```python |
| 49 | + embedder = keras_hub.models.BertTextEmbedder.from_preset( |
| 50 | + "all_minilm_l6_v2_en", |
| 51 | + ) |
| 52 | +
|
| 53 | + # Semantic search. |
| 54 | + query = "Which planet is known as the Red Planet?" |
| 55 | + documents = [ |
| 56 | + "Mars is often referred to as the Red Planet.", |
| 57 | + "Venus is often called Earth's twin.", |
| 58 | + ] |
| 59 | + q_emb = embedder.encode_text(query) |
| 60 | + d_embs = embedder.encode_documents(documents) |
| 61 | + sims = embedder.similarity(q_emb, d_embs) |
| 62 | + print("Best match:", documents[sims.argmax()]) |
| 63 | + ``` |
| 64 | +
|
| 65 | + Preprocessed integer data. |
| 66 | + ```python |
| 67 | + features = { |
| 68 | + "token_ids": np.ones(shape=(2, 12), dtype="int32"), |
| 69 | + "segment_ids": np.array([[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0]] * 2), |
| 70 | + "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2), |
| 71 | + } |
| 72 | +
|
| 73 | + embedder = keras_hub.models.BertTextEmbedder.from_preset( |
| 74 | + "all_minilm_l6_v2_en", |
| 75 | + preprocessor=None, |
| 76 | + ) |
| 77 | + embeddings = embedder.predict(features) |
| 78 | + ``` |
| 79 | + """ |
| 80 | + |
| 81 | + backbone_cls = BertBackbone |
| 82 | + preprocessor_cls = BertTextEmbedderPreprocessor |
| 83 | + |
| 84 | + def __init__( |
| 85 | + self, |
| 86 | + backbone, |
| 87 | + preprocessor=None, |
| 88 | + pooling_mode="mean", |
| 89 | + normalize=True, |
| 90 | + **kwargs, |
| 91 | + ): |
| 92 | + # === Layers === |
| 93 | + self.backbone = backbone |
| 94 | + self.preprocessor = preprocessor |
| 95 | + |
| 96 | + # === Functional Model === |
| 97 | + inputs = backbone.input |
| 98 | + backbone_outputs = backbone(inputs) |
| 99 | + sequence_output = backbone_outputs["sequence_output"] |
| 100 | + padding_mask = inputs["padding_mask"] |
| 101 | + |
| 102 | + # Apply pooling. |
| 103 | + if pooling_mode == "mean": |
| 104 | + pooled = self._mean_pooling(sequence_output, padding_mask) |
| 105 | + elif pooling_mode == "cls": |
| 106 | + pooled = sequence_output[:, 0, :] |
| 107 | + elif pooling_mode == "max": |
| 108 | + pooled = self._max_pooling(sequence_output, padding_mask) |
| 109 | + else: |
| 110 | + raise ValueError( |
| 111 | + f"Invalid pooling_mode: '{pooling_mode}'. " |
| 112 | + "Expected one of 'mean', 'cls', or 'max'." |
| 113 | + ) |
| 114 | + |
| 115 | + # Apply L2 normalization. |
| 116 | + if normalize: |
| 117 | + pooled = self._l2_normalize(pooled) |
| 118 | + |
| 119 | + super().__init__( |
| 120 | + inputs=inputs, |
| 121 | + outputs=pooled, |
| 122 | + **kwargs, |
| 123 | + ) |
| 124 | + |
| 125 | + # === Config === |
| 126 | + self.pooling_mode = pooling_mode |
| 127 | + self.normalize = normalize |
| 128 | + |
| 129 | + @staticmethod |
| 130 | + def _mean_pooling(sequence_output, padding_mask): |
| 131 | + """Attention-mask-aware mean pooling over token embeddings.""" |
| 132 | + # Expand mask: [batch, seq_len] -> [batch, seq_len, 1] |
| 133 | + mask = ops.cast( |
| 134 | + ops.expand_dims(padding_mask, axis=-1), sequence_output.dtype |
| 135 | + ) |
| 136 | + # Sum token embeddings, masked. |
| 137 | + sum_embeddings = ops.sum(sequence_output * mask, axis=1) |
| 138 | + # Sum mask for normalization. |
| 139 | + sum_mask = ops.maximum(ops.sum(mask, axis=1), 1e-9) |
| 140 | + return sum_embeddings / sum_mask |
| 141 | + |
| 142 | + @staticmethod |
| 143 | + def _max_pooling(sequence_output, padding_mask): |
| 144 | + """Max pooling over token embeddings, ignoring padding.""" |
| 145 | + mask = ops.cast(ops.expand_dims(padding_mask, axis=-1), dtype="bool") |
| 146 | + # Set padding positions to -inf so they don't affect max. |
| 147 | + fill_value = ops.cast( |
| 148 | + ops.convert_to_tensor(float("-inf")), sequence_output.dtype |
| 149 | + ) |
| 150 | + masked_output = ops.where(mask, sequence_output, fill_value) |
| 151 | + return ops.max(masked_output, axis=1) |
| 152 | + |
| 153 | + @staticmethod |
| 154 | + def _l2_normalize(embeddings): |
| 155 | + """L2 normalize embeddings to unit length.""" |
| 156 | + return ops.nn.normalize(embeddings, axis=-1, order=2) |
| 157 | + |
| 158 | + def get_config(self): |
| 159 | + config = super().get_config() |
| 160 | + config.update( |
| 161 | + { |
| 162 | + "pooling_mode": self.pooling_mode, |
| 163 | + "normalize": self.normalize, |
| 164 | + } |
| 165 | + ) |
| 166 | + return config |
| 167 | + |
| 168 | + def encode_documents(self, documents, **kwargs): |
| 169 | + """Encode a string or list of documents into embeddings. |
| 170 | +
|
| 171 | + This is a convenience method that wraps `predict()` for a |
| 172 | + single document or batch of documents. The output embeddings |
| 173 | + are suitable for computing similarity against query embeddings. |
| 174 | +
|
| 175 | + Args: |
| 176 | + documents: A string or list of strings to encode. |
| 177 | + **kwargs: Additional keyword arguments passed to |
| 178 | + `predict()`. |
| 179 | +
|
| 180 | + Returns: |
| 181 | + A tensor of shape `(batch_size, embedding_dim)`. |
| 182 | + """ |
| 183 | + if isinstance(documents, str): |
| 184 | + documents = [documents] |
| 185 | + return self.predict(documents, **kwargs) |
| 186 | + |
| 187 | + def similarity(self, query_embeddings, document_embeddings): |
| 188 | + """Compute similarity between query and document embeddings. |
| 189 | +
|
| 190 | + Computes the dot product between query and document embeddings. |
| 191 | + When embeddings are L2 normalized (the default), this is |
| 192 | + equivalent to cosine similarity. Returns a similarity matrix |
| 193 | + of shape `(num_queries, num_documents)`. |
| 194 | +
|
| 195 | + Args: |
| 196 | + query_embeddings: An array or tensor of shape |
| 197 | + `(num_queries, embedding_dim)`. |
| 198 | + document_embeddings: An array or tensor of shape |
| 199 | + `(num_documents, embedding_dim)`. |
| 200 | +
|
| 201 | + Returns: |
| 202 | + A tensor of shape `(num_queries, num_documents)` containing |
| 203 | + similarity scores. |
| 204 | + """ |
| 205 | + query_tensor = ops.convert_to_tensor(query_embeddings) |
| 206 | + document_tensor = ops.convert_to_tensor(document_embeddings) |
| 207 | + return ops.matmul(query_tensor, ops.transpose(document_tensor)) |
0 commit comments