Rebrand to baa.ai Merino-Mini-v2 (backbone-only attribution)

LICENSE

@@ -1,202 +1,24 @@
-
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright [yyyy] [name of copyright owner]
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.
+Merino-Mini-v2 — Proprietary License
+Copyright (c) 2026 BAA AI (Black Sheep AI). All rights reserved.
+
+1. SCOPE. This license governs the "BAA Contributions" in this package: the
+   shared word-embedding architecture and configuration, the router / loader
+   code (modeling_baa.py), the model packaging, BAA AI's weight contributions,
+   the model card, and associated documentation.
+
+2. GRANT. No right to use, reproduce, modify, distribute, sublicense, or create
+   derivative works of the BAA Contributions is granted except under a separate
+   written agreement with BAA AI (Black Sheep AI).
+
+3. THIRD-PARTY COMPONENT. This package incorporates the MiniLM-L6-H384-uncased backbone,
+   provided under the MIT License — see LICENSE-minilm.txt. The MIT terms govern that
+   backbone component only; nothing in this license limits any rights you have
+   under the MIT License with respect to it.
+
+4. NO WARRANTY. THE PACKAGE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO MERCHANTABILITY, FITNESS FOR
+   A PARTICULAR PURPOSE, AND NONINFRINGEMENT. IN NO EVENT SHALL BAA AI BE LIABLE
+   FOR ANY CLAIM, DAMAGES, OR OTHER LIABILITY ARISING FROM OR IN CONNECTION WITH
+   THE PACKAGE OR ITS USE.
+
+Contact: BAA AI (Black Sheep AI) — baa.ai

LICENSE-minilm.txt

@@ -0,0 +1,23 @@
+Backbone component: MiniLM-L6-H384-uncased — MIT License
+
+MIT License
+
+Copyright (c) Microsoft Corporation.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MODEL_CARD.md

@@ -0,0 +1,61 @@
+---
+license: other
+license_name: baa-proprietary
+library_name: sentence-transformers
+tags:
+- retrieval
+- embeddings
+- reranker
+- cross-encoder
+- rag
+- sentence-similarity
+pipeline_tag: sentence-similarity
+---
+
+# baa.ai · Merino-Mini-v2
+
+**One model that does both halves of RAG retrieval — bi-encoder embedding *and* cross-encoder reranking — over a single shared word-embedding table.** A 384-dimensional English model, ~44M parameters, by BAA AI (Black Sheep AI).
+
+## Get the optimal model for *your* data
+
+Merino-Mini-v2 is a strong, cost-efficient **default**. But the best embedder + reranker is **corpus-specific** — the ideal choice depends on your documents and your notion of relevance. **baa.ai offers exclusive tooling that identifies the optimal embedding and reranking models for your specific data**, so you ship the smallest models that maximize document recovery on your corpus. For a tailored recommendation, **reach out to baa.ai**.
+
+## What it is
+
+A two-role retrieval model over a **shared input word-embedding matrix** (stored once). The bi-encoder embedder and a deeper 12-layer cross-encoder reranker are built on the same `MiniLM-L6-H384-uncased` backbone, so their word-embedding table is stored a single time and injected into the reranker at load — a smaller download at **no measured quality loss**, with no retraining.
+
+- **Embed role:** bi-encoder, 384-d, L2-normalized.
+- **Rerank role:** cross-encoder, single relevance logit per (query, document) pair.
+- **Router:** call `.embed(...)` or `.rerank(...)`.
+
+## Usage
+
+```python
+from modeling_baa import BaaEmbeddingReranker   # included in this repo
+
+m = BaaEmbeddingReranker("baa-ai/Merino-Mini-v2")
+qv = m.embed(["how does a cross-encoder reranker work?"], is_query=True)[0]
+dv = m.embed(["a cross-encoder scores a (query, document) pair jointly",
+              "bi-encoders embed query and document separately for fast retrieval"])
+ranked = m.rerank("how does a cross-encoder reranker work?",
+                  ["a cross-encoder scores a (query, document) pair jointly",
+                   "the mitochondria is the powerhouse of the cell"])
+# -> [(doc, score), ...] sorted best-first
+```
+
+## Specs
+
+| | |
+|---|---|
+| Embedding dim | 384 |
+| Parameters | ~44M (embedder + reranker, shared word-embedding table) |
+| Languages | English |
+| Max sequence length | 512 |
+| Hardware | CPU / edge / GPU |
+
+## License & attribution
+
+- **BAA Contributions** (shared-embedding architecture, router/loader code, packaging, weights, docs) are **proprietary to BAA AI (Black Sheep AI)** — see `LICENSE`.
+- Incorporates the `MiniLM-L6-H384-uncased` backbone under the **MIT License** — see `LICENSE-minilm.txt`.
+
+© 2026 BAA AI (Black Sheep AI) — baa.ai. Provided "as is" without warranty.

NOTICE

@@ -1,27 +1,5 @@
-baa.ai · Embedding-Reranker-MiniLM-L12-v1
-Copyright (c) 2026 baa.ai
+Merino-Mini-v2
+Copyright (c) 2026 BAA AI (Black Sheep AI). All rights reserved.
 
-This product is a derivative work licensed under the Apache License, Version 2.0.
-It re-packages and modifies the following upstream models:
-
-  1. sentence-transformers/all-MiniLM-L6-v2  (license: Apache-2.0)
-       Used as the bi-encoder embedder. Provides the canonical (shared)
-       word-embedding table for the combined model.
-       https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
-
-  2. cross-encoder/ms-marco-MiniLM-L-12-v2  (license: Apache-2.0)
-       Used as the cross-encoder reranker. Its word-embedding table has been
-       removed on disk and is injected at load time from the shared table
-       above, reducing the combined footprint.
-       https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-12-v2
-
-Both upstream models derive from the Microsoft MiniLM-L6-H384-uncased architecture.
-
-Modifications by baa.ai:
-  - Unified the two models into a single artifact over one shared
-    word-embedding table (the reranker's word-embedding matrix is stored
-    once, in the embedder, and injected at load).
-  - Added a combined loader (modeling_baa.py) exposing embed() and rerank().
-
-This NOTICE file is provided in accordance with Section 4(d) of the
-Apache License, Version 2.0. See the LICENSE file for the full license text.
+BAA Contributions: proprietary — see LICENSE.
+Backbone: MiniLM-L6-H384-uncased — MIT License — see LICENSE-minilm.txt.

README.md

@@ -1,5 +1,6 @@
 ---
-license: apache-2.0
+license: other
+license_name: baa-proprietary
 library_name: sentence-transformers
 tags:
 - retrieval
@@ -11,26 +12,28 @@ tags:
 pipeline_tag: sentence-similarity
 ---
 
-# baa.ai · Embedding-Reranker-MiniLM-L12-v1
+# baa.ai · Merino-Mini-v2
 
-**A single model that does both halves of RAG retrieval — bi-encoder embedding *and* cross-encoder reranking — over one shared word-embedding table.**
+**One model that does both halves of RAG retrieval — bi-encoder embedding *and* cross-encoder reranking — over a single shared word-embedding table.** A 384-dimensional English model, ~44M parameters, by BAA AI (Black Sheep AI).
 
-The embedder and reranker share their word-embedding table (stored once), so the packaged model is **~20.9% smaller on disk than shipping the two components separately**, with **no measured loss in retrieval quality**.
+## Get the optimal model for *your* data
 
-## Why this model
+Merino-Mini-v2 is a strong, cost-efficient **default**. But the best embedder + reranker is **corpus-specific** — the ideal choice depends on your documents and your notion of relevance. **baa.ai offers exclusive tooling that identifies the optimal embedding and reranking models for your specific data**, so you ship the smallest models that maximize document recovery on your corpus. For a tailored recommendation, **reach out to baa.ai**.
 
-Most RAG stacks bolt an embedder onto a reranker and pay for both. A well-matched embedder + reranker that share a backbone recover the right documents just as well — in a single, smaller download.
+## What it is
 
-- **Two jobs, one download** — embed for retrieval, then rerank the candidates.
-- **Smaller footprint** — the shared word-embedding table is stored once.
-- **Strong default** — a sensible starting point for production RAG.
+A two-role retrieval model over a **shared input word-embedding matrix** (stored once). The bi-encoder embedder and a deeper 12-layer cross-encoder reranker are built on the same `MiniLM-L6-H384-uncased` backbone, so their word-embedding table is stored a single time and injected into the reranker at load — a smaller download at **no measured quality loss**, with no retraining.
+
+- **Embed role:** bi-encoder, 384-d, L2-normalized.
+- **Rerank role:** cross-encoder, single relevance logit per (query, document) pair.
+- **Router:** call `.embed(...)` or `.rerank(...)`.
 
 ## Usage
 
 ```python
 from modeling_baa import BaaEmbeddingReranker   # included in this repo
 
-m = BaaEmbeddingReranker("baa-ai/Embedding-Reranker-MiniLM-L12-v1")
+m = BaaEmbeddingReranker("baa-ai/Merino-Mini-v2")
 qv = m.embed(["how does a cross-encoder reranker work?"], is_query=True)[0]
 dv = m.embed(["a cross-encoder scores a (query, document) pair jointly",
               "bi-encoders embed query and document separately for fast retrieval"])
@@ -40,25 +43,19 @@ ranked = m.rerank("how does a cross-encoder reranker work?",
 # -> [(doc, score), ...] sorted best-first
 ```
 
-## Get the optimal models for *your* data
-
-This model is a great **default**. But the best embedder and reranker are **corpus-specific**. **baa.ai offers exclusive tooling that identifies the optimal embedding and reranking models for your specific data** — if you want that tailored recommendation, **reach out to baa.ai**.
-
 ## Specs
 
 | | |
 |---|---|
 | Embedding dim | 384 |
-| Vocab | 30522 |
+| Parameters | ~44M (embedder + reranker, shared word-embedding table) |
+| Languages | English |
 | Max sequence length | 512 |
-| Combined params | ~44.4M (vs ~56.1M separate) |
-| Footprint vs separate models | ~20.9% smaller on disk, no measured quality loss |
 | Hardware | CPU / edge / GPU |
 
 ## License & attribution
 
-Released under the **Apache License 2.0**. Derivative work re-packaging two upstream models into a single shared-backbone artifact (see `NOTICE`):
-- `sentence-transformers/all-MiniLM-L6-v2` (embedder)
-- `cross-encoder/ms-marco-MiniLM-L-12-v2` (reranker)
+- **BAA Contributions** (shared-embedding architecture, router/loader code, packaging, weights, docs) are **proprietary to BAA AI (Black Sheep AI)** — see `LICENSE`.
+- Incorporates the `MiniLM-L6-H384-uncased` backbone under the **MIT License** — see `LICENSE-minilm.txt`.
 
-© baa.ai. Provided "as is" without warranty; see `LICENSE`.
+© 2026 BAA AI (Black Sheep AI) — baa.ai. Provided "as is" without warranty.

config.json

@@ -1,23 +1,15 @@
 {
-  "model_type": "baa-embed-rerank",
-  "description": "Unified bi-encoder embedder + cross-encoder reranker sharing one word-embedding table.",
-  "components": {
-    "embedder": "embedder/",
-    "reranker": "reranker/"
-  },
-  "loader": "modeling_baa.BaaEmbeddingReranker",
+  "model_type": "baa-embedding-reranker",
+  "name": "Merino-Mini-v2",
+  "version": "1",
+  "license": "Proprietary \u2014 BAA AI (Black Sheep AI); MiniLM-L6-H384-uncased backbone under MIT",
+  "architecture": "shared-word-embedding: one MiniLM-L6-H384-uncased word-embedding table shared across the embedder and reranker stacks",
+  "embedding_dim": 384,
   "embed_query_prompt": "",
   "embed_doc_prompt": "",
-  "embedding_dim": 384,
-  "vocab_size": 30522,
   "max_seq_length": 512,
-  "separate_params_millions": 56.1,
-  "combined_params_millions": 44.4,
-  "disk_saving_pct": 20.9,
-  "license": "apache-2.0",
+  "params_millions": 44,
+  "backbone": "MiniLM-L6-H384-uncased (MIT)",
   "trust_remote_code": false,
-  "upstream": {
-    "embedder": "sentence-transformers/all-MiniLM-L6-v2",
-    "reranker": "cross-encoder/ms-marco-MiniLM-L-12-v2"
-  }
+  "loader": "modeling_baa.BaaEmbeddingReranker"
 }

embedder/README.md

@@ -1,175 +0,0 @@
----
-base_model:
-  - nreimers/MiniLM-L6-H384-uncased
-language: en
-license: apache-2.0
-library_name: sentence-transformers
-tags:
-- sentence-transformers
-- feature-extraction
-- sentence-similarity
-- transformers
-datasets:
-- s2orc
-- flax-sentence-embeddings/stackexchange_xml
-- ms_marco
-- gooaq
-- yahoo_answers_topics
-- code_search_net
-- search_qa
-- eli5
-- snli
-- multi_nli
-- wikihow
-- natural_questions
-- trivia_qa
-- embedding-data/sentence-compression
-- embedding-data/flickr30k-captions
-- embedding-data/altlex
-- embedding-data/simple-wiki
-- embedding-data/QQP
-- embedding-data/SPECTER
-- embedding-data/PAQ_pairs
-- embedding-data/WikiAnswers
-pipeline_tag: sentence-similarity
----
-
-
-# all-MiniLM-L6-v2
-This is a [sentence-transformers](https://www.SBERT.net) model: It maps sentences & paragraphs to a 384 dimensional dense vector space and can be used for tasks like clustering or semantic search.
-
-## Usage (Sentence-Transformers)
-Using this model becomes easy when you have [sentence-transformers](https://www.SBERT.net) installed:
-
-```
-pip install -U sentence-transformers
-```
-
-Then you can use the model like this:
-```python
-from sentence_transformers import SentenceTransformer
-sentences = ["This is an example sentence", "Each sentence is converted"]
-
-model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-embeddings = model.encode(sentences)
-print(embeddings)
-```
-
-## Usage (HuggingFace Transformers)
-Without [sentence-transformers](https://www.SBERT.net), you can use the model like this: First, you pass your input through the transformer model, then you have to apply the right pooling-operation on-top of the contextualized word embeddings.
-
-```python
-from transformers import AutoTokenizer, AutoModel
-import torch
-import torch.nn.functional as F
-
-#Mean Pooling - Take attention mask into account for correct averaging
-def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-
-
-# Sentences we want sentence embeddings for
-sentences = ['This is an example sentence', 'Each sentence is converted']
-
-# Load model from HuggingFace Hub
-tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-
-# Tokenize sentences
-encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
-
-# Compute token embeddings
-with torch.no_grad():
-    model_output = model(**encoded_input)
-
-# Perform pooling
-sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
-
-# Normalize embeddings
-sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
-
-print("Sentence embeddings:")
-print(sentence_embeddings)
-```
-
-------
-
-## Background
-
-The project aims to train sentence embedding models on very large sentence level datasets using a self-supervised 
-contrastive learning objective. We used the pretrained [`nreimers/MiniLM-L6-H384-uncased`](https://huggingface.co/nreimers/MiniLM-L6-H384-uncased) model and fine-tuned in on a 
-1B sentence pairs dataset. We use a contrastive learning objective: given a sentence from the pair, the model should predict which out of a set of randomly sampled other sentences, was actually paired with it in our dataset.
-
-We developed this model during the 
-[Community week using JAX/Flax for NLP & CV](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/7104), 
-organized by Hugging Face. We developed this model as part of the project:
-[Train the Best Sentence Embedding Model Ever with 1B Training Pairs](https://discuss.huggingface.co/t/train-the-best-sentence-embedding-model-ever-with-1b-training-pairs/7354). We benefited from efficient hardware infrastructure to run the project: 7 TPUs v3-8, as well as intervention from Googles Flax, JAX, and Cloud team member about efficient deep learning frameworks.
-
-## Intended uses
-
-Our model is intended to be used as a sentence and short paragraph encoder. Given an input text, it outputs a vector which captures 
-the semantic information. The sentence vector may be used for information retrieval, clustering or sentence similarity tasks.
-
-By default, input text longer than 256 word pieces is truncated.
-
-
-## Training procedure
-
-### Pre-training 
-
-We use the pretrained [`nreimers/MiniLM-L6-H384-uncased`](https://huggingface.co/nreimers/MiniLM-L6-H384-uncased) model. Please refer to the model card for more detailed information about the pre-training procedure.
-
-### Fine-tuning 
-
-We fine-tune the model using a contrastive objective. Formally, we compute the cosine similarity from each possible sentence pairs from the batch.
-We then apply the cross entropy loss by comparing with true pairs.
-
-#### Hyper parameters
-
-We trained our model on a TPU v3-8. We train the model during 100k steps using a batch size of 1024 (128 per TPU core).
-We use a learning rate warm up of 500. The sequence length was limited to 128 tokens. We used the AdamW optimizer with
-a 2e-5 learning rate. The full training script is accessible in this current repository: `train_script.py`.
-
-#### Training data
-
-We use the concatenation from multiple datasets to fine-tune our model. The total number of sentence pairs is above 1 billion sentences.
-We sampled each dataset given a weighted probability which configuration is detailed in the `data_config.json` file.
-
-
-| Dataset                                                  | Paper                                    | Number of training tuples  |
-|--------------------------------------------------------|:----------------------------------------:|:--------------------------:|
-| [Reddit comments (2015-2018)](https://github.com/PolyAI-LDN/conversational-datasets/tree/master/reddit) | [paper](https://arxiv.org/abs/1904.06472) | 726,484,430 |
-| [S2ORC](https://github.com/allenai/s2orc) Citation pairs (Abstracts) | [paper](https://aclanthology.org/2020.acl-main.447/) | 116,288,806 |
-| [WikiAnswers](https://github.com/afader/oqa#wikianswers-corpus) Duplicate question pairs | [paper](https://doi.org/10.1145/2623330.2623677) | 77,427,422 |
-| [PAQ](https://github.com/facebookresearch/PAQ) (Question, Answer) pairs | [paper](https://arxiv.org/abs/2102.07033) | 64,371,441 |
-| [S2ORC](https://github.com/allenai/s2orc) Citation pairs (Titles) | [paper](https://aclanthology.org/2020.acl-main.447/) | 52,603,982 |
-| [S2ORC](https://github.com/allenai/s2orc) (Title, Abstract) | [paper](https://aclanthology.org/2020.acl-main.447/) | 41,769,185 |
-| [Stack Exchange](https://huggingface.co/datasets/flax-sentence-embeddings/stackexchange_xml) (Title, Body) pairs  | - | 25,316,456 |
-| [Stack Exchange](https://huggingface.co/datasets/flax-sentence-embeddings/stackexchange_xml) (Title+Body, Answer) pairs  | - | 21,396,559 |
-| [Stack Exchange](https://huggingface.co/datasets/flax-sentence-embeddings/stackexchange_xml) (Title, Answer) pairs  | - | 21,396,559 |
-| [MS MARCO](https://microsoft.github.io/msmarco/) triplets | [paper](https://doi.org/10.1145/3404835.3462804) | 9,144,553 |
-| [GOOAQ: Open Question Answering with Diverse Answer Types](https://github.com/allenai/gooaq) | [paper](https://arxiv.org/pdf/2104.08727.pdf) | 3,012,496 |
-| [Yahoo Answers](https://www.kaggle.com/soumikrakshit/yahoo-answers-dataset) (Title, Answer) | [paper](https://proceedings.neurips.cc/paper/2015/hash/250cf8b51c773f3f8dc8b4be867a9a02-Abstract.html) | 1,198,260 |
-| [Code Search](https://huggingface.co/datasets/code_search_net) | - | 1,151,414 |
-| [COCO](https://cocodataset.org/#home) Image captions | [paper](https://link.springer.com/chapter/10.1007%2F978-3-319-10602-1_48) | 828,395|
-| [SPECTER](https://github.com/allenai/specter) citation triplets | [paper](https://doi.org/10.18653/v1/2020.acl-main.207) | 684,100 |
-| [Yahoo Answers](https://www.kaggle.com/soumikrakshit/yahoo-answers-dataset) (Question, Answer) | [paper](https://proceedings.neurips.cc/paper/2015/hash/250cf8b51c773f3f8dc8b4be867a9a02-Abstract.html) | 681,164 |
-| [Yahoo Answers](https://www.kaggle.com/soumikrakshit/yahoo-answers-dataset) (Title, Question) | [paper](https://proceedings.neurips.cc/paper/2015/hash/250cf8b51c773f3f8dc8b4be867a9a02-Abstract.html) | 659,896 |
-| [SearchQA](https://huggingface.co/datasets/search_qa) | [paper](https://arxiv.org/abs/1704.05179) | 582,261 |
-| [Eli5](https://huggingface.co/datasets/eli5) | [paper](https://doi.org/10.18653/v1/p19-1346) | 325,475 |
-| [Flickr 30k](https://shannon.cs.illinois.edu/DenotationGraph/) | [paper](https://transacl.org/ojs/index.php/tacl/article/view/229/33) | 317,695 |
-| [Stack Exchange](https://huggingface.co/datasets/flax-sentence-embeddings/stackexchange_xml) Duplicate questions (titles) | | 304,525 |
-| AllNLI ([SNLI](https://nlp.stanford.edu/projects/snli/) and [MultiNLI](https://cims.nyu.edu/~sbowman/multinli/) | [paper SNLI](https://doi.org/10.18653/v1/d15-1075), [paper MultiNLI](https://doi.org/10.18653/v1/n18-1101) | 277,230 | 
-| [Stack Exchange](https://huggingface.co/datasets/flax-sentence-embeddings/stackexchange_xml) Duplicate questions (bodies) | | 250,519 |
-| [Stack Exchange](https://huggingface.co/datasets/flax-sentence-embeddings/stackexchange_xml) Duplicate questions (titles+bodies) | | 250,460 |
-| [Sentence Compression](https://github.com/google-research-datasets/sentence-compression) | [paper](https://www.aclweb.org/anthology/D13-1155/) | 180,000 |
-| [Wikihow](https://github.com/pvl/wikihow_pairs_dataset) | [paper](https://arxiv.org/abs/1810.09305) | 128,542 |
-| [Altlex](https://github.com/chridey/altlex/) | [paper](https://aclanthology.org/P16-1135.pdf) | 112,696 |
-| [Quora Question Triplets](https://quoradata.quora.com/First-Quora-Dataset-Release-Question-Pairs) | - | 103,663 |
-| [Simple Wikipedia](https://cs.pomona.edu/~dkauchak/simplification/) | [paper](https://www.aclweb.org/anthology/P11-2117/) | 102,225 |
-| [Natural Questions (NQ)](https://ai.google.com/research/NaturalQuestions) | [paper](https://transacl.org/ojs/index.php/tacl/article/view/1455) | 100,231 |
-| [SQuAD2.0](https://rajpurkar.github.io/SQuAD-explorer/) | [paper](https://aclanthology.org/P18-2124.pdf) | 87,599 |
-| [TriviaQA](https://huggingface.co/datasets/trivia_qa) | - | 73,346 |
-| **Total** | | **1,170,060,424** |

modeling_baa.py

@@ -7,7 +7,7 @@ at no measured quality cost.
 
 Works for BERT-based and XLM-RoBERTa-based stacks alike: the reranker's encoder submodule is resolved
 generically via `reranker.base_model` (so `.bert` / `.roberta` are both handled). Optional per-model query/doc
-prompts are read from config.json (e.g. arctic uses "query: ").
+prompts are read from config.json (e.g. some models use a "query: " prefix).
 
 Usage:
     from modeling_baa import BaaEmbeddingReranker
@@ -49,6 +49,10 @@ class BaaEmbeddingReranker:
         base.embeddings.word_embeddings.weight.data = shared_wemb.to(self.reranker.dtype).clone()
         self.reranker.to(self.device).eval()
         self.rr_tok = AutoTokenizer.from_pretrained(rr_dir, trust_remote_code=trc)
+        # Weights may be stored fp16 on disk (smaller artifact); CPU can't compute in half -> upcast to fp32.
+        if str(self.device) == "cpu":
+            self.embedder = self.embedder.to(torch.float32)
+            self.reranker = self.reranker.float()
 
     def embed(self, texts, is_query=False, batch_size=32):
         """Return L2-normalized bi-encoder vectors. Applies the model's query/doc prompt if configured."""