Package upstream model with reproducible local MMEB evaluation

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,434 @@
+Tencent is pleased to support the open source community by making WeMM-Embedding-9B available.
+
+Copyright (C) 2026 Tencent.  All rights reserved. 
+
+The open-source software and/or models included in this distribution may have been modified by Tencent (“Tencent Modifications”). All Tencent Modifications are Copyright (C) Tencent.
+
+WeMM-Embedding-9B is licensed under Apache-2.0, except for the third-party components listed below, which remain licensed under their respective original terms. WeMM-Embedding-9B does not impose any additional restrictions beyond those specified in the original licenses of these third-party components. Users are required to comply with all applicable terms and conditions of the original licenses and to ensure that the use of these third-party components conforms to all relevant laws and regulations.
+
+For the avoidance of doubt, WeMM-Embedding-9B refers to code, parameters, and weights made publicly available by Tencent in accordance with Apache-2.0. 
+
+Terms of Apache-2.0:
+--------------------------------------------------------------------
+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.
+
+
+Dependencies and Licenses:
+
+This open-source project, WeMM-Embedding-9B, builds upon the following open-source models and/or software components, each of which remains licensed under its original license. Certain models or software may include modifications made by Tencent (“Tencent Modifications”), which are Copyright (C) Tencent.
+
+In case you believe there have been errors in the attribution below, you may submit the concerns to us for review and correction.
+
+
+Open Source Model Licensed under the  Apache-2.0:
+--------------------------------------------------------------------
+1. Qwen3.5-9B
+Copyright 2026 Alibaba Cloud
+Please note this model has been modified by Tencent in this distribution.
+
+
+ Terms of the  Apache-2.0：
+--------------------------------------------------------------------
+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.
+
+

README.md

@@ -0,0 +1,22 @@
+---
+license: other
+license_name: apache-2.0
+license_link: LICENSE
+library_name: transformers
+pipeline_tag: feature-extraction
+base_model: tencent/WeMM-Embedding-9B
+tags:
+- multimodal-embedding
+- retrieval
+- evaluation
+---
+
+# WeMM-Embedding-9B with MMEB evaluation
+
+Mirror of [tencent/WeMM-Embedding-9B](https://huggingface.co/tencent/WeMM-Embedding-9B) with the local Nexus MMEB evaluation source, configuration and runtime snapshot, prepared on 2026-10-15. The weights retain their upstream provenance and license. This is not a new fine-tuned model.
+
+- Original model documentation and inference examples: [UPSTREAM_README.md](UPSTREAM_README.md).
+- Evaluation setup and runnable commands: [evaluation/README.md](evaluation/README.md).
+- Integrity and source information: [package_manifest.json](package_manifest.json).
+
+The evaluation package includes image, video, moment and visual-document retrieval configurations. Evaluation data must be obtained separately. No new benchmark scores are asserted by this mirror.

UPSTREAM_README.md

@@ -0,0 +1,228 @@
+---
+library_name: transformers
+license: other
+license_name: apache-2.0
+license_link: https://huggingface.co/tencent/WeMM-Embedding-9B/blob/main/LICENSE
+
+base_model:
+- Qwen/Qwen3.5-9B
+pipeline_tag: feature-extraction
+
+tags:
+- sentence-transformers
+- multimodal-embedding
+- text-embedding
+- image-embedding
+- video-embedding
+- mrl
+
+language:
+- zh
+- en
+---
+
+# WeMM-Embedding-9B
+
+[![Hugging Face](https://img.shields.io/badge/🤗-Hugging%20Face-yellow)](https://huggingface.co/collections/tencent/wemm-embedding)
+[![Technical Report](https://img.shields.io/badge/📄-Technical%20Report-red)](https://arxiv.org/abs/2608.24053)
+[![GitHub](https://img.shields.io/badge/GitHub-WeMM--Embedding-black?logo=github)](https://github.com/Tencent/WeMM-Embedding)
+
+WeMM-Embedding-9B is a universal multimodal embedding model built on Qwen3.5. It accepts text, images, videos, visual documents, and interleaved multimodal inputs, and returns a 4,096-dimensional L2-normalized embedding. Audio input is not supported.
+
+## Installation
+
+```bash
+pip install torch transformers==5.2.0 "qwen-vl-utils[decord]==0.0.14" \
+  "sentence-transformers>=5.7.0" "accelerate>=1.1.0"
+```
+
+## Transformers
+
+```python
+import torch
+from qwen_vl_utils import process_vision_info
+from transformers import AutoModel, AutoProcessor
+
+model_id = "tencent/WeMM-Embedding-9B"
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    model_id, trust_remote_code=True, dtype=torch.bfloat16
+).cuda().eval()
+
+messages = [{"role": "user", "content": [
+    {"type": "image", "image": "/path/to/image.jpg"},
+    {"type": "video", "video": "/path/to/video.mp4"},
+    {"type": "text", "text": "This can be any text input."},
+]}]
+text = processor.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=False
+)
+images, videos, video_kwargs = process_vision_info(
+    messages,
+    image_patch_size=16,
+    return_video_kwargs=True,
+    return_video_metadata=True,
+)
+if videos is not None:
+    videos, video_metadata = zip(*videos)
+    videos, video_metadata = list(videos), list(video_metadata)
+else:
+    video_metadata = None
+inputs = processor(
+    text=text,
+    images=images,
+    videos=videos,
+    video_metadata=video_metadata,
+    return_tensors="pt",
+    **video_kwargs,
+).to("cuda")
+
+with torch.inference_mode():
+    embedding = model.embedding(**inputs)
+```
+
+Use any subset of the content items to encode text, image, or video independently.
+
+## Sentence Transformers
+
+```python
+from sentence_transformers import SentenceTransformer
+
+model_id = "tencent/WeMM-Embedding-9B"
+model = SentenceTransformer(model_id, trust_remote_code=True)
+
+queries = [
+    "Which Llama 4 model variants are available?",
+    "How is mapo tofu prepared?",
+]
+documents = [
+    "Mapo tofu is a Sichuan dish of soft tofu simmered in a spicy, numbing sauce of chili bean paste and Sichuan peppercorn.",
+    {
+        "image": "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/llama4_hgf.png",
+        "text": "Represent this image.",
+    },
+    {
+        "video": "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/mapo_tofu.mp4",
+        "text": "Represent this video.",
+    },
+]
+
+query_embeddings = model.encode_query(queries)
+document_embeddings = model.encode_document(documents)
+print(query_embeddings.shape, document_embeddings.shape)
+# (2, 4096) (3, 4096)
+
+similarities = model.similarity(query_embeddings, document_embeddings)
+print(similarities)
+# tensor([[0.2153, 0.5843, 0.1221],
+#         [0.7665, 0.2604, 0.5366]])
+```
+
+Each input is a string, a URL or path, a `PIL.Image`, or a dict combining `image`,
+`video`, and `text` keys. Put `image` or `video` before `text` so the prompt matches
+the ordering used above. Chat messages such as
+`{"role": "user", "content": [{"type": "image", "image": ...}, {"type": "text", "text": ...}]}`
+are also accepted, which is the way to interleave several images or videos in one input.
+
+## Matryoshka Embeddings
+
+```python
+d = 256
+embedding_d = torch.nn.functional.normalize(embedding[..., :d], dim=-1)
+```
+
+With Sentence Transformers, pass `truncate_dim` and let it renormalize:
+
+```python
+embeddings_d = model.encode_document(documents, truncate_dim=d, normalize_embeddings=True)
+```
+
+Use a dimension listed in `model.config.matryoshka_dimensions`.
+
+## Serving
+
+vLLM `0.27.0`:
+
+```bash
+MODEL_PATH=/path/to/WeMM-Embedding-9B
+vllm serve "$MODEL_PATH" \
+  --runner pooling \
+  --chat-template "$MODEL_PATH/embedding_chat_template.jinja"
+```
+
+SGLang `0.5.9`:
+
+```bash
+MODEL_PATH=/path/to/WeMM-Embedding-9B
+python patch_sglang_video.py
+python -m sglang.launch_server \
+  --model-path "$MODEL_PATH" \
+  --is-embedding \
+  --enable-precise-embedding-interpolation
+```
+
+## Evaluation
+
+### MMEB-v2
+
+Results on 78 datasets from Table 1 of the [technical report](https://github.com/Tencent/WeMM-Embedding/blob/main/assets/WeMM_Embedding_tech_report.pdf). Image and video tasks use Hit@1, while visual-document tasks use NDCG@5. Higher is better.
+
+| Model | Size | AVG | Image | Video | VisDoc |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| VLM2Vec | 2B | 47.8 | 59.7 | 29.0 | 44.0 |
+| GME | 2B | 55.4 | 51.9 | 33.9 | 76.8 |
+| VLM2Vec-V2 | 2B | 59.3 | 64.9 | 34.9 | 69.2 |
+| Qwen3-VL-Embedding | 2B | 73.2 | 75.0 | 61.9 | 79.2 |
+| DME-Small† | 2B | 74.8 | 75.9 | 65.6 | 79.9 |
+| **WeMM-Embedding** | **2B** | **77.9** | **79.6** | **70.8** | **80.7** |
+| **WeMM-Embedding** | **4B** | **79.2** | **80.8** | **72.1** | **82.0** |
+| VLM2Vec | 8B | 53.2 | 65.5 | 34.0 | 49.1 |
+| GME | 8B | 59.2 | 56.0 | 38.6 | 79.3 |
+| Qwen3-VL-Embedding | 8B | 77.8 | 80.1 | 67.1 | 82.4 |
+| DME-Medium† | 9B | 78.4 | 79.8 | 70.8 | 82.0 |
+| **WeMM-Embedding** | **9B** | **80.6** | **81.9** | **74.3** | **83.3** |
+
+† Closed-source leaderboard submission without publicly released model weights or a public inference endpoint.
+
+### MMEB-v3
+
+Results on all 190 tasks from Table 2 of the [technical report](https://github.com/Tencent/WeMM-Embedding/blob/main/assets/WeMM_Embedding_tech_report.pdf). V3-All includes the 78 MMEB-v2 tasks, 53 text tasks, 47 agent tasks, 11 audio tasks, and MCMR. Unsupported tasks are assigned a score of zero.
+
+| Model | Size | V3-All | Text | Agent | MCMR | Audio |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| VLM2Vec-V2 | 2B | 38.3 | 24.5 | 28.7 | 4.1 | 0.0 |
+| Omni-Embed-Nemotron | 3B | 43.5 | 39.2 | 36.5 | 26.1 | 36.5 |
+| E5-Omni | 3B | 44.6 | 26.7 | 36.9 | 31.9 | 30.8 |
+| Qwen3-VL-Embedding | 2B | 50.9 | 39.2 | 39.3 | 42.0 | 0.0 |
+| **WeMM-Embedding** | **2B** | **56.0** | **45.3** | **45.1** | **42.5** | **0.0** |
+| **WeMM-Embedding** | **4B** | **58.2** | **47.9** | **49.0** | **41.9** | **0.0** |
+| WAVE | 7B | 26.3 | 13.7 | 11.3 | 8.9 | 31.8 |
+| VLM2Vec | 8B | 32.9 | 22.2 | 19.7 | 0.9 | 0.0 |
+| LCO-Embedding-Omni | 7B | 40.6 | 32.4 | 27.8 | 20.0 | 43.2 |
+| GME | 8B | 43.6 | 37.1 | 35.6 | 27.3 | 0.0 |
+| E5-Omni | 7B | 47.1 | 26.9 | 36.7 | 41.1 | 43.0 |
+| Tianmu-Emb-Uni | 8B | 53.3 | 43.6 | 39.4 | 38.8 | 38.9 |
+| Qwen3-VL-Embedding | 8B | 53.5 | 42.5 | 38.4 | 38.0 | 0.0 |
+| **WeMM-Embedding** | **9B** | **59.5** | **48.8** | **51.0** | **49.3** | **0.0** |
+
+Text results use NDCG@5; agent, MCMR, and audio results use Hit@1.
+
+## Citation
+If you find this repository useful, please consider giving a star ⭐ and citation
+```bibtex
+@article{wemm-embedding,
+      title={WeMM-Embedding: WeChat Multi-Modal Embedding Technical Report}, 
+      author={Junjie Zhou and Ke Mei and Lei Li and Tianyi Wang and Fengyun Rao and Jing Lyu},
+      year={2026},
+      eprint={2608.24053},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2608.24053}, 
+}
+```
+
+## License
+
+WeMM-Embedding-9B, including the code, model parameters, and weights made publicly
+available by Tencent, is licensed under the [Apache License 2.0](https://huggingface.co/tencent/WeMM-Embedding-9B/blob/main/LICENSE).
+Third-party components remain subject to their respective original licenses.

additional_chat_templates/sentence_transformers.jinja

@@ -0,0 +1,32 @@
+{#- Embedding template used by Sentence Transformers.
+
+    Renders the same surface form as chat_template.jinja for text and images. Video
+    is the exception: it renders as a bare <|video_pad|>, because the processor
+    expands that token into one <seconds><|vision_start|>...<|vision_end|> block per
+    frame and, from transformers 5.3, no longer consumes a surrounding pair. The
+    wrapped form therefore gains a second, unwanted pair of vision boundaries there.
+    This matches embedding_chat_template.jinja, which is bare for the same reason.
+
+    The tokenizer normalizer drops the role newline before a leading image and the
+    newline after the final <|im_end|>, then the post-processor appends <embedding>. #}
+{%- for message in messages %}
+{{- '<|im_start|>' + message.role + '\n' }}
+{%- if message.content is string %}
+{{- message.content }}
+{%- elif message.content is iterable and message.content is not mapping %}
+{%- for item in message.content %}
+{%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+{{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+{%- elif 'video' in item or item.type == 'video' %}
+{{- '<|video_pad|>' }}
+{%- elif 'text' in item %}
+{{- item.text }}
+{%- else %}
+{{- raise_exception('Unexpected item type in content.') }}
+{%- endif %}
+{%- endfor %}
+{%- elif message.content is not none %}
+{{- raise_exception('Unexpected content type.') }}
+{%- endif %}
+{{- '<|im_end|>\n' }}
+{%- endfor %}

chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is true %}
+        {{- '<think>\n' }}
+    {%- else %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,122 @@
+{
+  "architectures": [
+    "Qwen3_5ForConditionalGeneration"
+  ],
+  "dtype": "bfloat16",
+  "image_token_id": 248056,
+  "model_type": "qwen3_5",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_output_gate": true,
+    "bos_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 32,
+    "linear_value_head_dim": 128,
+    "mamba_ssm_dtype": "float32",
+    "max_position_embeddings": 262144,
+    "mlp_only_layers": [],
+    "model_type": "qwen3_5_text",
+    "mtp_num_hidden_layers": 1,
+    "mtp_use_dedicated_embeddings": false,
+    "num_attention_heads": 16,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 4,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        11,
+        11,
+        10
+      ],
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 10000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 248078
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.2.0",
+  "video_token_id": 248057,
+  "vision_config": {
+    "deepstack_visual_indexes": [],
+    "depth": 27,
+    "dtype": "bfloat16",
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "model_type": "qwen3_5",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 4096,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 248054,
+  "vision_start_token_id": 248053,
+  "auto_map": {
+    "AutoModel": "modeling_wemm_embedding.WeMMEmbedding",
+    "AutoModelForCausalLM": "modeling_wemm_embedding.WeMMEmbedding"
+  },
+  "matryoshka_dimensions": [
+    64,
+    128,
+    256,
+    512,
+    1024,
+    2048,
+    4096
+  ]
+}

config_sentence_transformers.json

@@ -0,0 +1,11 @@
+{
+  "__version__": {
+    "pytorch": "2.11.0+cu128",
+    "sentence_transformers": "5.7.0",
+    "transformers": "5.2.0"
+  },
+  "default_prompt_name": null,
+  "model_type": "SentenceTransformer",
+  "prompts": {},
+  "similarity_fn_name": "cosine"
+}

embedding_chat_template.jinja

@@ -0,0 +1,28 @@
+{# Direct-tokenization template for vLLM 0.27 embedding requests.
+   Qwen's processor post-processor removes the role newline before images,
+   keeps it before text/video, and appends <embedding> without a preceding
+   newline. vLLM 0.27 adds each frame's vision boundaries itself, so video
+   uses a bare <video_pad> placeholder. #}
+{%- for message in messages %}
+{{- '<|im_start|>' + message.role }}
+{%- if message.content is string %}
+{{- '\n' + message.content }}
+{%- else %}
+{%- for item in message.content %}
+{%- if item.type in ['image', 'image_url'] or 'image' in item or 'image_url' in item %}
+{{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+{%- elif item.type in ['video', 'video_url'] or 'video' in item or 'video_url' in item %}
+{{- '\n<|video_pad|>' }}
+{%- elif item.type == 'text' or 'text' in item %}
+{{- ('\n' if loop.first else '') + item.text }}
+{%- else %}
+{{- raise_exception('Unsupported embedding content type') }}
+{%- endif %}
+{%- endfor %}
+{%- endif %}
+{{- '<|im_end|>' }}
+{%- if not loop.last %}
+{{- '\n' }}
+{%- endif %}
+{%- endfor %}
+{{- '<embedding>' }}

evaluation/README.md

@@ -0,0 +1,31 @@
+# Reproduce the local MMEB evaluation
+
+This is a snapshot of the local Nexus MMEB evaluator on 2026-10-15, packaged with tencent/WeMM-Embedding-9B weights. These are the upstream released models, not local training checkpoints.
+
+## Setup and execution
+
+Use Python 3.10 and a CUDA-capable PyTorch environment. Install dependencies:
+
+```bash
+python -m pip install -r evaluation/requirements.txt
+DATA_BASEDIR=/path/to/vlm2vec_eval NPROC=8 bash evaluation/run_eval.sh image_retrieval
+DATA_BASEDIR=/path/to/vlm2vec_eval NPROC=8 bash evaluation/run_eval.sh video_ret_5
+DATA_BASEDIR=/path/to/vlm2vec_eval NPROC=8 bash evaluation/run_eval.sh moment_retrieval
+DATA_BASEDIR=/path/to/vlm2vec_eval NPROC=8 bash evaluation/run_eval.sh visdoc_full
+```
+
+Each process loads a full model on its GPU; NPROC is data parallelism, not model sharding. Reduce BATCH_SIZE if memory is insufficient. One GPU is the default. MODEL_PATH, OUTPUT_DIR, GENERATED_MEDIA_BASEDIR, MASTER_PORT, BATCH_SIZE and FORCE_RECOMPUTE may be overridden.
+
+The launcher sets NEXUS_WEMM_OFFICIAL=1 before importing Nexus. WeMM uses its released custom model, user-only prompt, processor settings and embedding() path. Qwen uses the Nexus Qwen-style conversation and last-token pooling. Both normalize embeddings. Task instructions, candidate selection, local/global evaluation, video sampling and metrics are determined by the included dataset parsers and YAML files; preserve them when comparing results. The video configuration uses 64 frames. The evaluator writes *_score.json files into OUTPUT_DIR.
+
+## Data requirements
+
+Benchmark media are not included in this model package. DATA_BASEDIR must contain image-tasks/, video-tasks/ and visdoc-tasks/ as required by the YAMLs. The dataset parsers download annotations/corpora from their configured Hugging Face sources; see nexus/Nexus/evaluation/mmeb_v2/constant.py and data/datasets/. Visual-document page images can be generated under GENERATED_MEDIA_BASEDIR, which must be writable. Internet access or an existing dataset cache is required.
+
+## Runtime snapshot
+
+The original runtime uses transformers 5.6.0.dev0, installed from a development source archive without a recorded Git commit. To avoid depending on a moving main branch, vendor/ contains the installed Transformers source and distribution metadata (excluding bytecode and install records, with an upstream test token redacted in testing_utils.py). The launcher prepends this runtime snapshot to PYTHONPATH. Other installed dependency versions are recorded in requirements.txt. The upstream model's recommended environment may differ; see UPSTREAM_README.md. Vendored source retains its upstream licensing.
+
+## Scope of verification
+
+Packaging checks include weight shard references, JSON/YAML parsing, source compilation, evaluator CLI imports, and file hashes. No new GPU benchmark run is performed during packaging. This package does not claim to reproduce any particular score without the original benchmark data and runtime hardware.

evaluation/configs/image_retrieval.yaml

@@ -0,0 +1,75 @@
+# RET i -> t
+MSCOCO_i2t:
+    dataset_parser: image_i2t
+    dataset_name: MSCOCO_i2t
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+VisualNews_i2t:
+    dataset_parser: image_i2t
+    dataset_name: VisualNews_i2t
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+# RET t -> i
+VisDial:
+    dataset_parser: image_t2i
+    dataset_name: VisDial
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+MSCOCO_t2i:
+    dataset_parser: image_t2i
+    dataset_name: MSCOCO_t2i
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+VisualNews_t2i:
+    dataset_parser: image_t2i
+    dataset_name: VisualNews_t2i
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+WebQA:
+    dataset_parser: image_t2i
+    dataset_name: WebQA
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+EDIS:
+    dataset_parser: image_t2i
+    dataset_name: EDIS
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+Wiki-SS-NQ:
+    dataset_parser: image_t2i
+    dataset_name: Wiki-SS-NQ
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+# RET i -> i
+CIRR:
+    dataset_parser: image_i2i_vg
+    dataset_name: CIRR
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+NIGHTS:
+    dataset_parser: image_i2i_vg
+    dataset_name: NIGHTS
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+OVEN:
+    dataset_parser: image_i2i_vg
+    dataset_name: OVEN
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local
+FashionIQ:
+    dataset_parser: image_i2i_vg
+    dataset_name: FashionIQ
+    dataset_split: test
+    image_root: image-tasks/
+    eval_type: local

evaluation/configs/moment_retrieval.yaml

@@ -0,0 +1,35 @@
+QVHighlight:
+    dataset_parser: moment_retrieval
+    dataset_name: QVHighlight
+    data_path: video-tasks/data/qvhighlight.jsonl
+    video_root: ""
+    clip_root: ""
+    frame_root: video-tasks/frames/video_mret/QVHighlight
+    num_negative_clips: 9
+    max_video_frames_saved: 64
+    max_clip_frames_saved: 8
+    num_video_frames: 64
+    num_clip_frames: 8
+    eval_type: local
+Charades-STA:
+    dataset_parser: moment_retrieval
+    dataset_name: Charades-STA
+    data_path: video-tasks/data/charades_sta.jsonl
+    video_root: ""
+    clip_root: ""
+    frame_root: video-tasks/frames/video_mret/Charades-STA
+    num_negative_clips: 9
+    max_video_frames_saved: 64
+    max_clip_frames_saved: 8
+    num_video_frames: 64
+    num_clip_frames: 8
+    eval_type: local
+MomentSeeker:
+    dataset_parser: momentseeker
+    dataset_name: MomentSeeker
+    video_root: ""
+    clip_root: ""
+    frame_root: video-tasks/frames/video_mret/MomentSeeker
+    num_negative_clips: 9
+    num_video_frames: 64
+    eval_type: local

evaluation/configs/video_ret_5.yaml

@@ -0,0 +1,45 @@
+MSR-VTT:
+    dataset_parser: msrvtt
+    dataset_name: MSR-VTT
+    video_root: video-tasks/videos/video_ret/MSR-VTT/
+    frame_root: video-tasks/frames/video_ret/data/ziyan/video_retrieval/MSR-VTT/frames
+    num_sample_per_subset: None
+    max_frames_saved: 64
+    num_frames: 64
+    eval_type: global
+MSVD:
+    dataset_parser: msvd
+    dataset_name: MSVD
+    video_root: video-tasks/videos/video_ret/MSVD/
+    frame_root: video-tasks/frames/video_ret/data/ziyan/video_retrieval/MSVD/frames
+    num_sample_per_subset: None
+    max_frames_saved: 64
+    num_frames: 64
+    eval_type: global
+DiDeMo:
+    dataset_parser: didemo
+    dataset_name: DiDeMo
+    video_root: video-tasks/videos/video_ret/DiDeMo/
+    frame_root: video-tasks/frames/video_ret/data/ziyan/video_retrieval/DiDeMo/frames
+    num_sample_per_subset: None
+    max_frames_saved: 64
+    num_frames: 64
+    eval_type: global
+YouCook2:
+    dataset_parser: youcook2
+    dataset_name: YouCook2
+    video_root: video-tasks/videos/video_ret/YouCook2/
+    frame_root: video-tasks/frames/video_ret/data/ziyan/video_retrieval/YouCook2/frames
+    num_sample_per_subset: None
+    max_frames_saved: 64
+    num_frames: 64
+    eval_type: global
+VATEX:
+    dataset_parser: vatex
+    dataset_name: VATEX
+    video_root: video-tasks/videos/video_ret/VATEX/
+    frame_root: video-tasks/frames/video_ret/data/ziyan/video_retrieval/VATEX/frames
+    num_sample_per_subset: None
+    max_frames_saved: 64
+    num_frames: 64
+    eval_type: global

evaluation/configs/visdoc_full.yaml

@@ -0,0 +1,144 @@
+ViDoRe_arxivqa:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_arxivqa
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_arxivqa
+    eval_type: global
+ViDoRe_docvqa:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_docvqa
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_docvqa
+    eval_type: global
+ViDoRe_infovqa:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_infovqa
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_infovqa
+    eval_type: global
+ViDoRe_tabfquad:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_tabfquad
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_tabfquad
+    eval_type: global
+ViDoRe_tatdqa:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_tatdqa
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_tatdqa
+    eval_type: global
+ViDoRe_shiftproject:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_shiftproject
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_shiftproject
+    eval_type: global
+ViDoRe_syntheticDocQA_artificial_intelligence:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_syntheticDocQA_artificial_intelligence
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_syntheticDocQA_artificial_intelligence
+    eval_type: global
+ViDoRe_syntheticDocQA_energy:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_syntheticDocQA_energy
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_syntheticDocQA_energy
+    eval_type: global
+ViDoRe_syntheticDocQA_government_reports:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_syntheticDocQA_government_reports
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_syntheticDocQA_government_reports
+    eval_type: global
+ViDoRe_syntheticDocQA_healthcare_industry:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_syntheticDocQA_healthcare_industry
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoRe_syntheticDocQA_healthcare_industry
+    eval_type: global
+ViDoRe_esg_reports_human_labeled_v2:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_esg_reports_human_labeled_v2
+    dataset_split: test
+    image_root: visdoc-tasks/esg_reports_human_labeled_v2
+    eval_type: global
+ViDoRe_biomedical_lectures_v2_multilingual:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_biomedical_lectures_v2_multilingual
+    dataset_split: test
+    image_root: visdoc-tasks/biomedical_lectures_v2_multilingual
+    eval_type: global
+ViDoRe_economics_reports_v2_multilingual:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_economics_reports_v2_multilingual
+    dataset_split: test
+    image_root: visdoc-tasks/economics_reports_v2_multilingual
+    eval_type: global
+ViDoRe_esg_reports_v2_multilingual:
+    dataset_parser: vidore
+    dataset_name: ViDoRe_esg_reports_v2_multilingual
+    dataset_split: test
+    image_root: visdoc-tasks/esg_reports_v2_multilingual
+    eval_type: global
+VisRAG_ArxivQA:
+    dataset_parser: visrag
+    dataset_name: VisRAG_ArxivQA
+    dataset_split: test
+    image_root: visdoc-tasks/VisRAG_ArxivQA
+    eval_type: global
+VisRAG_ChartQA:
+    dataset_parser: visrag
+    dataset_name: VisRAG_ChartQA
+    dataset_split: test
+    image_root: visdoc-tasks/VisRAG_ChartQA
+    eval_type: global
+VisRAG_MP-DocVQA:
+    dataset_parser: visrag
+    dataset_name: VisRAG_MP-DocVQA
+    dataset_split: test
+    image_root: visdoc-tasks/VisRAG_MP-DocVQA
+    eval_type: global
+VisRAG_SlideVQA:
+    dataset_parser: visrag
+    dataset_name: VisRAG_SlideVQA
+    dataset_split: test
+    image_root: visdoc-tasks/VisRAG_SlideVQA
+    eval_type: global
+VisRAG_InfoVQA:
+    dataset_parser: visrag
+    dataset_name: VisRAG_InfoVQA
+    dataset_split: test
+    image_root: visdoc-tasks/VisRAG_InfoVQA
+    eval_type: global
+VisRAG_PlotQA:
+    dataset_parser: visrag
+    dataset_name: VisRAG_PlotQA
+    dataset_split: test
+    image_root: visdoc-tasks/VisRAG_PlotQA
+    eval_type: global
+ViDoSeek-page:
+    dataset_parser: vidore
+    dataset_name: ViDoSeek-page
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoSeek-page
+    eval_type: global
+ViDoSeek-doc:
+    dataset_parser: vidore
+    dataset_name: ViDoSeek-doc
+    dataset_split: test
+    image_root: visdoc-tasks/ViDoSeek-doc
+    eval_type: global
+MMLongBench-page:
+    dataset_parser: vidore
+    dataset_name: MMLongBench-page
+    dataset_split: test
+    image_root: visdoc-tasks/MMLongBench-page
+    eval_type: global
+MMLongBench-doc:
+    dataset_parser: vidore
+    dataset_name: MMLongBench-doc
+    dataset_split: test
+    image_root: visdoc-tasks/MMLongBench-doc
+    eval_type: global

evaluation/nexus/LICENSE

@@ -0,0 +1,22 @@
+Copyright [2024] [USTC]
+
+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.
+
+Additional Terms:
+1. Users are prohibited from utilizing Nexus to provide any commercial services to Third-Party Entities without obtaining explicit prior written authorization from USTC.
+
+Definition of "Third-Party Entities":
+Refers to any enterprises, social organizations, institutions, or other
+legal entities not affiliated with the individual or organizational entity
+that utilizes the Nexus source code, excluding the said user and its
+directly controlled subsidiaries.

evaluation/nexus/Nexus/__init__.py

@@ -0,0 +1,24 @@
+from importlib import import_module
+
+__all__ = [
+    "AbsInferenceArguments",
+    "TextEmbedder",
+    "BaseEmbedderInferenceEngine",
+    "TextReranker",
+    "BaseRerankerInferenceEngine",
+    "MultiModalEmbedder",
+    "MultimodalEmbedder",
+]
+
+
+def __getattr__(name):
+    if name == "AbsInferenceArguments":
+        module = import_module("Nexus.abc")
+        return getattr(module, name)
+    if name in {"TextEmbedder", "BaseEmbedderInferenceEngine", "TextReranker", "BaseRerankerInferenceEngine"}:
+        module = import_module("Nexus.inference")
+        return getattr(module, name)
+    if name in {"MultiModalEmbedder", "MultimodalEmbedder"}:
+        module = import_module("Nexus.inference.embedder.multimodal_retrieval")
+        return getattr(module, name, getattr(module, "MultimodalEmbedder"))
+    raise AttributeError(f"module 'Nexus' has no attribute {name!r}")

evaluation/nexus/Nexus/abc/__init__.py

@@ -0,0 +1,10 @@
+from importlib import import_module
+
+__all__ = ["AbsInferenceArguments"]
+
+
+def __getattr__(name):
+    if name == "AbsInferenceArguments":
+        module = import_module("Nexus.abc.inference.arguments")
+        return getattr(module, name)
+    raise AttributeError(f"module 'Nexus.abc' has no attribute {name!r}")

evaluation/nexus/Nexus/abc/arguments.py

@@ -0,0 +1,197 @@
+import json
+import os
+import yaml
+import logging
+from pathlib import Path
+from typing import Any, Union, get_args, get_origin
+from dataclasses import dataclass, asdict, fields
+
+logger = logging.getLogger(__name__)
+
+DEFINITE_LOCAL_PATH_KEYS = {
+    "train_data",
+    "eval_data",
+    "output_dir",
+    "dataset_dir",
+    "media_root",
+    "image_root",
+    "video_root",
+    "cache_path",
+    "cache_dir",
+    "corpus_embd_save_dir",
+    "eval_output_dir",
+    "eval_output_path",
+}
+MAYBE_LOCAL_OR_REMOTE_KEYS = {
+    "model_name_or_path",
+    "processor_name_or_path",
+    "embedder_name_or_path",
+}
+
+
+def init_argument(type_, x):
+    if x is None:
+        return None
+    origin = get_origin(type_)
+    args = get_args(type_)
+
+    if type_ in [Any, object]:
+        return x
+
+    if origin is Union:
+        for tmp_type_ in args:
+            if tmp_type_ is type(None):
+                continue
+            try:
+                return init_argument(tmp_type_, x)
+            except TypeError:
+                continue
+        raise TypeError(f"Failed to init argument {x} ({type(x)}) to {type_}.")
+
+    if isinstance(x, dict):
+        try:
+            # is subclass of AbsArguments
+            return type_.from_dict(x)
+        except:
+            # dict
+            return dict(x)
+    elif origin in [list, tuple]:
+        item_type = args[0] if len(args) > 0 else Any
+        converted = [init_argument(item_type, i) for i in x]
+        return converted if origin is list else tuple(converted)
+    elif isinstance(x, list):
+        return list(x)
+    else:
+        tmp_x = None
+        try:
+            if isinstance(type_, type) and isinstance(x, type_):
+                tmp_x = x
+            else:
+                tmp_x = type_(x)
+                logger.warning(f"Init argument: Convert {x} ({type(x)}) to {tmp_x} ({type(tmp_x)}).")
+        except:
+            try:
+                for tmp_type_ in args:
+                    if tmp_type_ is None:
+                        continue
+                    if tmp_type_ is type(None):
+                        continue
+                    if isinstance(x, tmp_type_):
+                        tmp_x = x
+                        break
+                    else:
+                        tmp_x = tmp_type_(x)
+                        logger.warning(f"Init argument: Convert {x} ({type(x)}) to {tmp_x} ({type(tmp_x)}).")
+            except (AttributeError, TypeError):
+                raise TypeError(f"Failed to init argument {x} ({type(x)}) to {type_}.")
+        if tmp_x is None:
+            raise TypeError(f"Failed to init argument {x} ({type(x)}) to {type_}.")
+        return tmp_x
+
+
+def _looks_like_remote_reference(value: str) -> bool:
+    return "://" in value
+
+
+def _resolve_local_path_value(value: str, base_dir: Path, force: bool = False) -> str:
+    if value in [None, ""]:
+        return value
+
+    expanded = os.path.expanduser(value)
+    if os.path.isabs(expanded) or _looks_like_remote_reference(expanded):
+        return expanded
+
+    if not force:
+        candidate = base_dir / expanded
+        if not (expanded.startswith(".") or expanded.startswith("..") or candidate.exists()):
+            return value
+
+    return str((base_dir / expanded).resolve())
+
+
+def _resolve_config_paths(config_dict: dict, base_dir: Path) -> dict:
+    resolved = dict(config_dict)
+    for key, value in list(resolved.items()):
+        if value in [None, ""]:
+            continue
+
+        if key in DEFINITE_LOCAL_PATH_KEYS:
+            if isinstance(value, list):
+                resolved[key] = [_resolve_local_path_value(item, base_dir, force=True) for item in value]
+            elif isinstance(value, str):
+                resolved[key] = _resolve_local_path_value(value, base_dir, force=True)
+            continue
+
+        if key in MAYBE_LOCAL_OR_REMOTE_KEYS and isinstance(value, str):
+            resolved[key] = _resolve_local_path_value(value, base_dir, force=False)
+    return resolved
+
+
+@dataclass
+class AbsArguments:
+
+    def to_dict(self):
+        return asdict(self)
+
+    def to_json(self, save_path: Union[str, Path], overwrite: bool = False):
+        if isinstance(save_path, str):
+            save_path = Path(save_path)
+
+        if save_path.exists() and not overwrite:
+            raise FileExistsError(f"{save_path} already exists. Set `overwrite=True` to overwrite the file.")
+
+        save_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(save_path, "w", encoding="utf-8") as f:
+            json.dump(self.to_dict(), f, indent=4)
+
+    @classmethod
+    def from_dict(cls, _dict: dict):
+        _fields_w_type = [_field.name for _field in fields(cls)]
+        
+        _fields_wo_type = [
+            x for x in cls.__dict__.keys()
+            if not x.startswith("__") and not x.endswith("__") and x not in _fields_w_type
+        ]
+        _fields_wo_type_dict = {
+            _field_name: _dict.pop(_field_name)
+            for _field_name in _fields_wo_type if _field_name in _dict
+        }
+        
+        for k in _dict.keys():
+            if k not in [_field.name for _field in fields(cls)]:
+                raise ValueError(f'{k} is not in fields({cls}).')
+
+        for _field in fields(cls):
+            if _field.name not in _dict:
+                continue
+            _dict[_field.name] = init_argument(_field.type, _dict[_field.name])
+        
+        _instance = cls(**_dict)
+        
+        for _field_name, _field_value in _fields_wo_type_dict.items():
+            setattr(_instance, _field_name, _field_value)
+        return _instance
+
+    @classmethod
+    def from_json(cls, load_path: Union[str, Path]):
+        if isinstance(load_path, str):
+            load_path = Path(load_path)
+
+        if not load_path.exists():
+            raise FileNotFoundError(f"{load_path} does not exist.")
+
+        with open(load_path, "r", encoding="utf-8") as f:
+            _dict = _resolve_config_paths(json.load(f), load_path.parent)
+            return cls.from_dict(_dict)
+
+    @classmethod
+    def from_yaml(cls, load_path: Union[str, Path]):
+        if isinstance(load_path, str):
+            load_path = Path(load_path)
+
+        if not load_path.exists():
+            raise FileNotFoundError(f"{load_path} does not exist.")
+
+        with open(load_path, "r", encoding="utf-8") as f:
+            _dict = _resolve_config_paths(yaml.safe_load(f), load_path.parent)
+            return cls.from_dict(_dict)

evaluation/nexus/Nexus/abc/evaluation/__init__.py

@@ -0,0 +1,13 @@
+from .arguments import AbsEvalArguments, AbsEvalDataLoaderArguments
+from .data_loader import AbsEvalDataLoader
+from .evaluator import AbsEvaluator
+from .runner import AbsEvalRunner
+
+
+__all__ = [
+    "AbsEvalArguments",
+    "AbsEvalDataLoaderArguments",
+    "AbsEvalDataLoader",
+    "AbsEvaluator",
+    "AbsEvalRunner",
+]

evaluation/nexus/Nexus/abc/evaluation/arguments.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass, field
+
+from Nexus.abc.arguments import AbsArguments
+
+
+@dataclass
+class AbsEvalArguments(AbsArguments):
+    eval_name: str = field(
+        default=None,
+        metadata={"help": "Name of the evaluation task."}
+    )
+    eval_output_dir: str = field(
+        default=None,
+        metadata={"help": "Directory to save the evaluation results."}
+    )
+    overwrite: bool = field(
+        default=False,
+        metadata={"help": "Overwrite the output directory if it exists."}
+    )
+
+
+@dataclass
+class AbsEvalDataLoaderArguments(AbsArguments):
+    eval_dataset_name_or_path: str = field(
+        default=None,
+        metadata={"help": "Name or path to the evaluation dataset."}
+    )
+    eval_batch_size: int = field(
+        default=1,
+        metadata={"help": "Batch size for evaluation."}
+    )

evaluation/nexus/Nexus/abc/evaluation/data_loader.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+
+
+class AbsEvalDataLoader(ABC):
+    @abstractmethod
+    def __init__(
+        self,
+        *args,
+        **kwargs
+    ):
+        pass

evaluation/nexus/Nexus/abc/evaluation/evaluator.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+
+from .data_loader import AbsEvalDataLoader
+
+
+class AbsEvaluator(ABC):
+    def __init__(
+        self,
+        eval_name: str,
+        data_loader: AbsEvalDataLoader,
+        overwrite: bool = False
+    ):
+        self.eval_name = eval_name
+        self.data_loader = data_loader
+        self.overwrite = overwrite
+
+    @abstractmethod
+    def evaluate(self, *args, **kwargs):
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def compute_metrics(*args, **kwargs):
+        pass

evaluation/nexus/Nexus/abc/evaluation/runner.py

@@ -0,0 +1,41 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from .arguments import AbsEvalArguments, AbsEvalDataLoaderArguments
+from .data_loader import AbsEvalDataLoader
+from .evaluator import AbsEvaluator
+
+
+class AbsEvalRunner(ABC):
+    def __init__(
+        self,
+        eval_args: AbsEvalArguments,
+        eval_data_loader_args: AbsEvalDataLoaderArguments
+    ):
+        self.eval_args = eval_args
+        self.data_loader = self.load_data_loader(eval_data_loader_args)
+        self.evaluator = self.load_evaluator()
+
+    def load_data_loader(
+        self,
+        eval_data_loader_args: AbsEvalDataLoaderArguments
+    ) -> AbsEvalDataLoader:
+        return AbsEvalDataLoader(
+            **eval_data_loader_args.to_dict()
+        )
+
+    def load_evaluator(self):
+        return AbsEvaluator(
+            eval_name=self.eval_args.eval_name,
+            data_loader=self.data_loader,
+            overwrite=self.eval_args.overwrite
+        )
+
+    @abstractmethod
+    def run(
+        self,
+        model: Any,
+        *args,
+        **kwargs
+    ):
+        pass

evaluation/nexus/Nexus/abc/inference/__init__.py

@@ -0,0 +1,24 @@
+from importlib import import_module
+
+__all__ = [
+    "AbsEmbedder",
+    "AbsReranker",
+    "AbsInferenceArguments",
+    "InferenceEngine",
+]
+
+
+def __getattr__(name):
+    if name == "AbsEmbedder":
+        module = import_module("Nexus.abc.inference.embedder")
+        return getattr(module, name)
+    if name == "AbsReranker":
+        module = import_module("Nexus.abc.inference.reranker")
+        return getattr(module, name)
+    if name == "AbsInferenceArguments":
+        module = import_module("Nexus.abc.inference.arguments")
+        return getattr(module, name)
+    if name == "InferenceEngine":
+        module = import_module("Nexus.abc.inference.inference_engine")
+        return getattr(module, name)
+    raise AttributeError(f"module 'Nexus.abc.inference' has no attribute {name!r}")

evaluation/nexus/Nexus/abc/inference/arguments.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass, field
+from typing import List, Union
+
+from Nexus.abc.arguments import AbsArguments
+
+
+@dataclass
+class AbsInferenceArguments(AbsArguments):
+    stage: str = field(
+        default=None,
+        metadata={"help": "Stage of the inference process."}
+    )
+    use_fp16: bool = field(
+        default=True,
+        metadata={"help": "use fp16"}
+    )
+    model_name_or_path: str = field(
+        default=None,
+        metadata={"help": "Name or path to the model."}
+    )
+    onnx_model_path: str= field(
+        default=None,
+        metadata={
+            'help':'Path to onnx model'
+        }
+    )
+    trt_model_path: str= field(
+        default=None,
+        metadata={
+            'help':'Path to trt model'
+        }
+    )
+    output_topk: int = field(
+        default=10,
+        metadata={"help": "Number of top-k results to output."}
+    )
+    infer_device: Union[str, int, List[int], List[str]] = field(
+        default="cpu",
+        metadata={"help": "Device to perform inference.", "nargs": "+"}
+    )
+    infer_mode: str = field(
+        default="normal",
+        metadata={"help": "Inference mode: normal, onnx, tensorrt.", "choices": ["normal", "onnx", "tensorrt"]}
+    )
+    infer_batch_size: int = field(
+        default=1,
+        metadata={"help": "Batch size for inference."}
+    )
+    max_workspace_size: int = field(
+        default=1 << 30,
+        metadata={
+            'help':'Max workspace size for tesorrt session, default = 1 GB'
+        }
+    )
+    normalize: bool = field(
+        default=True,
+        metadata={
+            'help':'normalize embeddings'
+        }
+    )
+    query_max_length: int=field(
+        default=512,
+        metadata={
+            'help':'Query max length'
+        }
+    )
+    passage_max_length: int=field(
+        default=512,
+        metadata={
+            'help':'Passage max length'
+        }
+    )

evaluation/nexus/Nexus/abc/inference/embedder/AbsEmbedder.py

@@ -0,0 +1,264 @@
+import logging
+from tqdm import tqdm, trange
+from abc import ABC, abstractmethod
+from typing import Any, Union, List, Dict, Literal, Optional
+
+import queue
+import multiprocessing as mp
+from multiprocessing import Queue
+
+import math
+import gc
+import torch
+import numpy as np
+from transformers import is_torch_npu_available
+
+logger = logging.getLogger(__name__)
+
+
+class AbsEmbedder(ABC):
+    """
+    Base class for embedder.
+    Extend this class and implement :meth:`encode_queries`, :meth:`encode_corpus`, :meth:`encode` for custom embedders.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        
+    def stop_self_pool(self):
+        if self.pool is not None:
+            self.stop_multi_process_pool(self.pool)
+            self.pool = None
+        try:
+            self.model.to('cpu')
+            torch.cuda.empty_cache()
+        except:
+            pass
+        gc.collect()
+
+    @staticmethod
+    def get_target_devices(devices: Union[str, int, List[str], List[int]]) -> List[str]:
+        """
+
+        Args:
+            devices (Union[str, int, List[str], List[int]]): specified devices, can be `str`, `int`, list of `str`, or list of `int`.
+
+        Raises:
+            ValueError: Devices should be a string or an integer or a list of strings or a list of integers.
+
+        Returns:
+            List[str]: A list of target devices in format.
+        """
+        if devices is None:
+            if torch.cuda.is_available():
+                return [f"cuda:{i}" for i in range(torch.cuda.device_count())]
+            elif is_torch_npu_available():
+                return [f"npu:{i}" for i in range(torch.npu.device_count())]
+            elif torch.backends.mps.is_available():
+                return [f"mps:{i}" for i in range(torch.mps.device_count())]
+            else:
+                return ["cpu"]
+        elif isinstance(devices, str):
+            return [devices]
+        elif isinstance(devices, int):
+            return [f"cuda:{devices}"]
+        elif isinstance(devices, list):
+            if isinstance(devices[0], str):
+                return devices
+            elif isinstance(devices[0], int):
+                return [f"cuda:{device}" for device in devices]
+            else:
+                raise ValueError("devices should be a string or an integer or a list of strings or a list of integers.")
+        else:
+            raise ValueError("devices should be a string or an integer or a list of strings or a list of integers.")
+
+    @abstractmethod
+    def encode_query(
+        self,
+        *args,
+        **kwargs: Any
+    ):
+        """encode queries
+        """
+        return self.encode(*args, **kwargs)
+    
+    @abstractmethod
+    def encode_info(
+        self,
+        *args,
+        **kwargs: Any
+    ):
+        """
+        encode info (corpus)
+        """
+        return self.encode(*args, **kwargs)
+
+
+    @abstractmethod
+    def encode(
+        self,
+        *args, **kwargs
+    ):
+        """encode
+        """
+        pass
+
+        
+    def __del__(self):
+        self.stop_self_pool()
+
+    @abstractmethod
+    def encode_single_device(
+        self,
+        *args,
+        **kwargs
+    ):
+        """
+        This method should encode on a single device.
+        """
+        pass
+
+    # adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L807
+    def start_multi_process_pool(
+        self,
+        process_target_func: Any,
+    ) -> Dict[Literal["input", "output", "processes"], Any]:
+        """
+        Starts a multi-process pool to process the encoding with several independent processes
+        via :meth:`SentenceTransformer.encode_multi_process <sentence_transformers.SentenceTransformer.encode_multi_process>`.
+
+        This method is recommended if you want to encode on multiple GPUs or CPUs. It is advised
+        to start only one process per GPU. This method works together with encode_multi_process
+        and stop_multi_process_pool.
+
+        Returns:
+            Dict[str, Any]: A dictionary with the target processes, an input queue, and an output queue.
+        """
+        if self.model is None:
+            raise ValueError("Model is not initialized.")
+
+        logger.info("Start multi-process pool on devices: {}".format(", ".join(map(str, self.target_devices))))
+
+        self.model.to("cpu")
+        self.model.share_memory()
+        ctx = mp.get_context("spawn")
+        input_queue = ctx.Queue()
+        output_queue = ctx.Queue()
+        processes = []
+
+        for device_id in tqdm(self.target_devices, desc='initial target device'):
+            p = ctx.Process(
+                target=process_target_func,
+                args=(device_id, self, input_queue, output_queue),
+                daemon=True,
+            )
+            p.start()
+            processes.append(p)
+
+        return {"input": input_queue, "output": output_queue, "processes": processes}
+
+    # adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L976
+    @staticmethod
+    def _encode_multi_process_worker(
+        target_device: str, model: 'AbsEmbedder', input_queue: Queue, results_queue: Queue
+    ) -> None:
+        """
+        Internal working process to encode sentences in multi-process setup
+        """
+        while True:
+            try:
+                chunk_id, sentences, kwargs = (
+                    input_queue.get()
+                )
+                embeddings = model.encode_single_device(
+                    sentences,
+                    device=target_device,
+                    **kwargs
+                )
+
+                results_queue.put([chunk_id, embeddings])
+            except queue.Empty:
+                break
+
+    # copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
+    @staticmethod
+    def stop_multi_process_pool(pool: Dict[Literal["input", "output", "processes"], Any]) -> None:
+        """
+        Stops all processes started with start_multi_process_pool.
+
+        Args:
+            pool (Dict[str, object]): A dictionary containing the input queue, output queue, and process list.
+
+        Returns:
+            None
+        """
+        for p in pool["processes"]:
+            p.terminate()
+
+        for p in pool["processes"]:
+            p.join()
+            p.close()
+
+        pool["input"].close()
+        pool["output"].close()
+        try:
+            pool["input"].join_thread()
+        except Exception:
+            pass
+        try:
+            pool["output"].join_thread()
+        except Exception:
+            pass
+        pool = None
+
+    # adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L877
+    def encode_multi_process(
+        self,
+        sentences: List[str],
+        pool: Dict[Literal["input", "output", "processes"], Any],
+        **kwargs
+    ):
+        chunk_size = math.ceil(len(sentences) / len(pool["processes"]))
+
+        input_queue = pool["input"]
+        last_chunk_id = 0
+        chunk = []
+
+        for sentence in sentences:
+            chunk.append(sentence)
+            if len(chunk) >= chunk_size:
+                input_queue.put(
+                    [last_chunk_id, chunk, kwargs]
+                )
+                last_chunk_id += 1
+                chunk = []
+
+        if len(chunk) > 0:
+            input_queue.put([last_chunk_id, chunk, kwargs])
+            last_chunk_id += 1
+
+        output_queue = pool["output"]
+        results_list = sorted(
+            [output_queue.get() for _ in trange(last_chunk_id, desc="Chunks")],
+            key=lambda x: x[0],
+        )
+        embeddings = self._concatenate_results_from_multi_process([result[1] for result in results_list])
+        return embeddings
+
+    def _concatenate_results_from_multi_process(self, results_list: List[Union[torch.Tensor, np.ndarray, Any]]):
+        """concatenate and return the results from all the processes
+
+        Args:
+            results_list (List[Union[torch.Tensor, np.ndarray, Any]]): A list of results from all the processes.
+
+        Raises:
+            NotImplementedError: Unsupported type for results_list
+
+        Returns:
+            Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
+        """
+        if isinstance(results_list[0], torch.Tensor):
+            return torch.cat(results_list, dim=0)
+        elif isinstance(results_list[0], np.ndarray):
+            return np.concatenate(results_list, axis=0)
+        else:
+            raise NotImplementedError("Unsupported type for results_list")

evaluation/nexus/Nexus/abc/inference/embedder/__init__.py

@@ -0,0 +1 @@
+from .AbsEmbedder import AbsEmbedder

evaluation/nexus/Nexus/abc/inference/inference_engine.py

@@ -0,0 +1,87 @@
+import logging
+from pathlib import Path
+from abc import ABC, abstractmethod
+from typing import Union, List, Tuple, Any
+
+import pandas as pd
+
+try:
+    import onnx
+except ImportError:
+    onnx = None
+
+try:
+    import onnxruntime as ort
+except ImportError:
+    ort = None
+
+try:
+    import tensorrt as trt
+except ImportError:
+    trt = None
+
+from .arguments import AbsInferenceArguments
+
+logger = logging.getLogger(__name__)
+
+class InferenceEngine(ABC):
+    def __init__(self, infer_args: AbsInferenceArguments):
+        self.config = infer_args.to_dict()
+
+    # @abstractmethod
+    # def load_model(self):
+    #     pass
+
+    @staticmethod
+    def load_onnx_model(onnx_model_path: Union[str, Path]):
+        if onnx is None:
+            raise ImportError("onnx is required to load ONNX models.")
+        onnx_model = onnx.load(onnx_model_path)
+        onnx.checker.check_model(onnx_model)
+        logger.info(f"====== Loaded ONNX model from {onnx_model_path} ======")
+        logger.info(onnx.helper.printable_graph(onnx_model.graph))
+        return onnx_model
+
+    # @abstractmethod
+    # def get_normal_session(self):
+    #     pass
+
+    @abstractmethod
+    def get_ort_session(self) -> Any:
+        pass
+
+    @abstractmethod
+    def get_trt_session(self) -> Any:
+        pass
+
+    def get_inference_session(self):
+        if self.config["infer_mode"] == "normal":
+            return self.get_normal_session()
+        if self.config["infer_mode"] == "onnx":
+            return self.get_ort_session()
+        elif self.config["infer_mode"] == "tensorrt":
+            return self.get_trt_session()
+        else:
+            raise ValueError(f"Invalid inference mode: {self.config['infer_mode']}")
+
+    @abstractmethod
+    def convert_to_onnx(self):
+        pass
+
+    # @abstractmethod
+    # def inference(
+    #     self,
+    #     inputs: Union[List[str], List[Tuple[str, str]], pd.DataFrame, Any],
+    #     *args,
+    #     **kwargs
+    # ):
+    #     pass
+    
+    def save_output_topk(self, output):
+        output_df = {}
+        output_df['_'.join(self.config['request_features'])] = \
+            self.infer_df.apply(lambda row : '_'.join([str(row[feat]) for feat in self.config['request_features']]),
+                                axis='columns')
+        output_df[self.feature_config['fiid']] = output.tolist()
+        output_df = pd.DataFrame(output_df)
+        output_df.to_feather(self.config['output_save_path'])

evaluation/nexus/Nexus/abc/inference/reranker/AbsReranker.py

@@ -0,0 +1,188 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Union, List, Tuple, Dict, Literal, Optional
+
+import multiprocessing as mp
+from multiprocessing import Queue
+
+import math
+import gc
+import torch
+import numpy as np
+from tqdm import tqdm, trange
+from transformers import is_torch_npu_available
+
+logger = logging.getLogger(__name__)
+
+
+class AbsReranker(ABC):
+    """
+    Base class for Reranker.
+    Extend this class and implement :meth:`compute_score_single_gpu` for custom rerankers.
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs: Any,
+    ):
+        pass
+
+    def stop_self_pool(self):
+        if self.pool is not None:
+            self.stop_multi_process_pool(self.pool)
+            self.pool = None
+        try:
+            self.model.to('cpu')
+            torch.cuda.empty_cache()
+        except:
+            pass
+        gc.collect()
+
+    @staticmethod
+    def get_target_devices(devices: Union[str, int, List[str], List[int]]) -> List[str]:
+        """
+
+        Args:
+            devices (Union[str, int, List[str], List[int]]): Specified devices, can be `str`, `int`, list of `str`, or list of `int`.
+
+        Raises:
+            ValueError: Devices should be a string or an integer or a list of strings or a list of integers.
+
+        Returns:
+            List[str]: A list of target devices in format
+        """
+        if devices is None:
+            if torch.cuda.is_available():
+                return [f"cuda:{i}" for i in range(torch.cuda.device_count())]
+            elif is_torch_npu_available():
+                return [f"npu:{i}" for i in range(torch.npu.device_count())]
+            elif torch.backends.mps.is_available():
+                return ["mps"]
+            else:
+                return ["cpu"]
+        elif isinstance(devices, str):
+            return [devices]
+        elif isinstance(devices, int):
+            return [f"cuda:{devices}"]
+        elif isinstance(devices, list):
+            if isinstance(devices[0], str):
+                return devices
+            elif isinstance(devices[0], int):
+                return [f"cuda:{device}" for device in devices]
+            else:
+                raise ValueError("devices should be a string or an integer or a list of strings or a list of integers.")
+        else:
+            raise ValueError("devices should be a string or an integer or a list of strings or a list of integers.")
+
+    @abstractmethod
+    def compute_score(
+        self,
+        *args,
+        **kwargs
+    ):
+        """Compute score for each sentence pair
+
+        Returns:
+            numpy.ndarray: scores of all the sentence pairs.
+        """
+        pass
+
+    def __del__(self):
+        self.stop_self_pool()
+
+    @abstractmethod
+    def compute_score_single_gpu(
+        self,
+        sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
+        batch_size: int = 256,
+        query_max_length: Optional[int] = None,
+        max_length: int = 512,
+        normalize: bool = False,
+        device: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        """
+        This method should compute the scores of sentence_pair and return scores.
+        """
+        pass
+
+    # copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
+    def start_multi_process_pool(self) -> Dict[Literal["input", "output", "processes"], Any]:
+        """
+        Starts a multi-process pool to process the encoding with several independent processes
+        via :meth:`SentenceTransformer.encode_multi_process <sentence_transformers.SentenceTransformer.encode_multi_process>`.
+
+        This method is recommended if you want to encode on multiple GPUs or CPUs. It is advised
+        to start only one process per GPU. This method works together with encode_multi_process
+        and stop_multi_process_pool.
+
+        Returns:
+            Dict[str, Any]: A dictionary with the target processes, an input queue, and an output queue.
+        """
+        logger.info("Start multi-process pool on devices: {}".format(", ".join(map(str, self.target_devices))))
+
+        self.model.to("cpu")
+        self.model.share_memory()
+        ctx = mp.get_context("spawn")
+        input_queue = ctx.Queue()
+        output_queue = ctx.Queue()
+        processes = []
+
+        for device_id in tqdm(self.target_devices, desc='initial target device'):
+            p = ctx.Process(
+                target=AbsReranker._encode_multi_process_worker,
+                args=(device_id, self, input_queue, output_queue),
+                daemon=True,
+            )
+            p.start()
+            processes.append(p)
+
+        return {"input": input_queue, "output": output_queue, "processes": processes}
+
+    # copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
+
+    # copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
+    @staticmethod
+    def _encode_multi_process_worker(
+            target_device: str, model: 'AbsReranker', input_queue: Queue, results_queue: Queue
+    ) -> None:
+        """
+        Internal working process to encode sentences in multi-process setup
+        """
+        while True:
+            try:
+                chunk_id, sentences, kwargs = (
+                    input_queue.get()
+                )
+                embeddings = model.compute_score_single_gpu(
+                    sentences,
+                    device=target_device,
+                    **kwargs
+                )
+
+                results_queue.put([chunk_id, embeddings])
+            except:
+                break
+
+    # copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
+    @staticmethod
+    def stop_multi_process_pool(pool: Dict[Literal["input", "output", "processes"], Any]) -> None:
+        """
+        Stops all processes started with start_multi_process_pool.
+
+        Args:
+            pool (Dict[str, object]): A dictionary containing the input queue, output queue, and process list.
+
+        Returns:
+            None
+        """
+        for p in pool["processes"]:
+            p.terminate()
+
+        for p in pool["processes"]:
+            p.join()
+            p.close()
+
+        pool["input"].close()
+        pool["output"].close()

evaluation/nexus/Nexus/abc/inference/reranker/__init__.py

@@ -0,0 +1 @@
+from .AbsReranker import AbsReranker

evaluation/nexus/Nexus/abc/training/arguments.py

@@ -0,0 +1,36 @@
+from dataclasses import dataclass, field
+
+from Nexus.abc.arguments import AbsArguments
+from transformers import TrainingArguments
+
+@dataclass
+class AbsDataArguments(AbsArguments):
+    train_data: str = field(
+        default=None,
+        metadata={"help": "Path to the data directory."}
+    )
+    eval_data: str = field(
+        default=None,
+        metadata={"help": "Path to the data directory."}
+    )
+    data_cache_dir: str = field(
+        default=None,
+        metadata={"help": "Path to the cache directory."}
+    )
+
+
+@dataclass
+class AbsModelArguments(AbsArguments):
+    model_name_or_path: str = field(
+        default=None,
+        metadata={"help": "Name or path to the model."}
+    )
+    model_cache_dir: str = field(
+        default=None,
+        metadata={"help": "Path to the cache directory."}
+    )
+
+
+@dataclass
+class AbsTrainingArguments(AbsArguments, TrainingArguments):
+    pass

evaluation/nexus/Nexus/abc/training/dataset.py

@@ -0,0 +1,8 @@
+from torch.utils.data import Dataset
+
+from .arguments import AbsDataArguments
+
+
+class AbsDataset(Dataset):
+    def __init__(self, args: AbsDataArguments):
+        self.args = args

evaluation/nexus/Nexus/abc/training/embedder/AbsArguments.py

@@ -0,0 +1,46 @@
+import os
+from typing import Optional
+from dataclasses import dataclass, field
+from Nexus.abc.training.arguments import AbsDataArguments,  AbsModelArguments, AbsTrainingArguments
+from transformers import TrainingArguments
+
+
+@dataclass
+class AbsEmbedderModelArguments(AbsModelArguments):
+    """
+    Abstract class for model arguments.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "The model checkpoint for initialization."}
+    )
+
+
+@dataclass
+class AbsEmbedderDataArguments(AbsDataArguments):
+    """
+    Abstract class for data arguments.
+    """
+    train_data: str = field(
+        default=None, metadata={
+            "help": "One or more paths to training data. `query: str`, `pos: List[str]`, `neg: List[str]` are required in the training data.",
+            "nargs": "+"
+        }
+    )
+    cache_path: Optional[str] = field(
+        default=None, metadata={"help": "Where do you want to store the cached data"}
+    )
+
+
+    def __post_init__(self):
+        if not isinstance(self.train_data, list):
+            self.train_data = self.train_data.split()
+            
+        for train_dir in self.train_data:
+            if not os.path.exists(train_dir):
+                raise FileNotFoundError(f"cannot find file: {train_dir}, please set a true path")
+
+
+@dataclass
+class AbsEmbedderTrainingArguments(AbsTrainingArguments):
+    pass

evaluation/nexus/Nexus/abc/training/embedder/AbsDataset.py

@@ -0,0 +1,34 @@
+import logging
+from dataclasses import dataclass
+from abc import abstractmethod
+
+
+from Nexus.abc.training.dataset import AbsDataset
+
+logger = logging.getLogger(__name__)
+
+
+class AbsEmbedderTrainDataset(AbsDataset):
+    """Abstract class for training dataset.
+
+    Args:
+        args (AbsEmbedderDataArguments): Data arguments.
+    """
+    
+    pass
+
+@dataclass
+class AbsEmbedderCollator():
+    """
+    The abstract embedder collator.
+    """
+    @abstractmethod
+    def __call__(self, features):
+        return super().__call__(features)
+
+@dataclass
+class CallbackOutput:
+    save_checkpoint: str = None
+    stop_training: bool = False
+
+        

evaluation/nexus/Nexus/abc/training/embedder/AbsModeling.py

@@ -0,0 +1,85 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+import torch.distributed as dist
+from transformers import AutoTokenizer
+
+import logging
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, List, Union
+from Nexus.abc.training.modeling import AbsModelOutput, AbsEmbedder
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EmbedderOutput(AbsModelOutput):
+    """
+    Output information returned by the model.
+    """
+    pass
+
+class AbsEmbedderModel(AbsEmbedder, nn.Module):
+    """Abstract class of embedding model for training.
+
+    Args:
+        base_model: The base model to train on.
+        tokenizer (AutoTokenizer, optional): The tokenizer to use. Defaults to ``None``.
+        negatives_cross_device (bool, optional): If True, will compute cross devices negative loss. Defaults to ``False``.
+        temperature (float, optional): Temperature to control the scale of scores. Defaults to ``1.0``.
+        sub_batch_size (int, optional): Sub-batch size during encoding. If negative, will not split to sub-batch.
+            Defaults to ``-1``.
+        kd_loss_type (str, optional): Type of knowledge distillation loss. Defaults to ``"kl_div"``.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.init_modules()
+
+    def init_modules(self):
+        self.loss_function=self.get_loss_function()
+        self.score_function=self.get_score_function()
+        
+    @abstractmethod
+    def get_loss_function(self):
+        pass
+    
+    @abstractmethod
+    def get_score_function(self):
+        pass
+
+    @abstractmethod
+    def compute_loss(self, batch, *args, **kwargs):
+        """Abstract method compute the loss.
+
+        Args:
+            scores (torch.Tensor): Computed score.
+            target (torch.Tensor): The target value.
+        """
+        pass
+
+    @abstractmethod
+    def compute_score(self, *args, **kwargs):
+        """Abstract method to compute the score.
+        """
+        pass
+
+    @abstractmethod
+    def save(self, output_dir: str):
+        """Abstract method to save the model.
+
+        Args:
+            output_dir (str): Directory for saving the model.
+        """
+        pass
+
+    def forward(self, batch, *args, **kwargs):
+        """The computation performed at every call.
+
+        Args:
+            Datas after datacollator
+
+        Returns:
+            EmbedderOutput: Output of the forward call of model.
+        """
+        return self.compute_loss(batch,*args, **kwargs)

evaluation/nexus/Nexus/abc/training/embedder/AbsRunner.py

@@ -0,0 +1,79 @@
+import os
+import logging
+from pathlib import Path
+from typing import Tuple, Union
+from abc import ABC, abstractmethod
+from transformers import set_seed, PreTrainedTokenizer
+
+
+from .AbsArguments import (
+    AbsEmbedderModelArguments,
+    AbsEmbedderDataArguments,
+    AbsEmbedderTrainingArguments
+)
+from .AbsTrainer import AbsEmbedderTrainer
+from .AbsModeling import AbsEmbedderModel
+from .AbsDataset import (
+    AbsEmbedderTrainDataset, AbsEmbedderCollator)
+from Nexus.abc.training.runner import AbsRunner
+
+logger = logging.getLogger(__name__)
+
+
+class AbsEmbedderRunner(AbsRunner):
+    """Abstract class to run embedding model fine-tuning.
+
+    Args:
+        model_args (AbsEmbedderModelArguments): Model arguments
+        data_args (AbsEmbedderDataArguments): Data arguments.
+        training_args (AbsEmbedderTrainingArguments): Training arguments.
+    """
+    def __init__(
+        self,
+        model_args: AbsEmbedderModelArguments,
+        data_args: AbsEmbedderDataArguments,
+        training_args: AbsEmbedderTrainingArguments
+    ):
+        self.model_args = model_args
+        self.data_args = data_args
+        self.training_args = training_args
+
+    @abstractmethod
+    def load_model(self, *args, **kwargs) -> AbsEmbedderModel:
+        pass
+        
+    @abstractmethod
+    def load_trainer(self) -> AbsEmbedderTrainer:
+        """Abstract method to load the trainer.
+
+        Returns:
+            AbsEmbedderTrainer: The loaded trainer instance.
+        """
+        pass
+
+    @abstractmethod
+    def load_dataset(self, *args, **kwargs) -> AbsEmbedderTrainDataset:
+        pass
+    
+    @abstractmethod
+    def load_data_collator(self) -> AbsEmbedderCollator:
+        """Loads the appropriate data collator.
+
+        Returns:
+            AbsEmbedderCollator: Loaded data collator.
+        """
+        pass
+
+    def run(self):
+        """
+        Executes the training process.
+        """
+        Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
+
+        # Training
+        self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
+        print('Training completed.')
+        self.trainer.save_model()
+        print('Training completed. save done')
+
+

evaluation/nexus/Nexus/abc/training/embedder/AbsTrainer.py

@@ -0,0 +1,39 @@
+import logging
+import torch
+from typing import Optional
+from abc import ABC, abstractmethod
+from transformers.trainer import Trainer
+from Nexus.abc.training.trainer import AbsTrainer
+from Nexus.abc.training.dataset import AbsDataset
+logger = logging.getLogger(__name__)
+
+
+class AbsEmbedderTrainer(AbsTrainer):
+    """
+    Abstract class for the trainer of embedder.
+    """
+    @abstractmethod
+    def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        pass
+    
+    def compute_loss(self, model, inputs, return_outputs=False ,*args, **kwargs):
+        """
+        How the loss is computed by Trainer. By default, all models return the loss in the first element.
+
+        Subclass and override for custom behavior.
+        
+        Args:
+            model (AbsEmbedderModel): The model being trained.
+            inputs (dict): A dictionary of input tensors to be passed to the model.
+            return_outputs (bool, optional): If ``True``, returns both the loss and the model's outputs. Otherwise,
+                returns only the loss.
+        
+        Returns:
+            Union[torch.Tensor, tuple(torch.Tensor, EmbedderOutput)]: The computed loss. If ``return_outputs`` is ``True``, 
+                also returns the model's outputs in a tuple ``(loss, outputs)``.
+        """
+        outputs = model(**inputs)
+        loss = outputs.loss
+
+        return (loss, outputs) if return_outputs else loss
+    

evaluation/nexus/Nexus/abc/training/embedder/__init__.py

@@ -0,0 +1,25 @@
+from .AbsArguments import AbsEmbedderModelArguments, AbsEmbedderDataArguments,AbsEmbedderTrainingArguments
+from .AbsDataset import (
+    AbsEmbedderTrainDataset,
+    AbsEmbedderCollator,
+    CallbackOutput
+)
+from .AbsModeling import (
+    EmbedderOutput,
+    AbsEmbedderModel
+)
+from .AbsRunner import AbsEmbedderRunner
+from .AbsTrainer import AbsEmbedderTrainer
+
+__all__ = [
+    'AbsEmbedderModelArguments',
+    'AbsEmbedderDataArguments',
+    'AbsEmbedderTrainingArguments',
+    'AbsEmbedderTrainDataset',
+    'AbsEmbedderCollator',
+    'EmbedderOutput',
+    'AbsEmbedderModel',
+    'AbsEmbedderRunner',
+    'AbsEmbedderTrainer',
+    'CallbackOutput'
+]

evaluation/nexus/Nexus/abc/training/modeling.py

@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, asdict
+
+from torch import nn
+
+@dataclass
+class AbsModelOutput:
+    def to_dict(self):
+        return asdict(self)
+
+
+class AbsModel(ABC, nn.Module):
+
+    @abstractmethod
+    def compute_loss(self, *args, **kwargs):
+        pass
+
+
+class AbsEmbedder(ABC, nn.Module):
+    
+    @abstractmethod
+    def compute_score(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def encode_query(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def encode_info(self, *args, **kwargs):
+        pass
+
+
+class AbsReranker(ABC, nn.Module):
+    
+    @abstractmethod
+    def compute_score(self, *args, **kwargs):
+        pass

evaluation/nexus/Nexus/abc/training/reranker/AbsArguments.py

@@ -0,0 +1,39 @@
+import os
+from typing import Optional
+from dataclasses import dataclass, field
+
+from transformers import TrainingArguments
+from Nexus.abc.training.arguments import AbsDataArguments,  AbsModelArguments, AbsTrainingArguments
+
+@dataclass
+class AbsRerankerModelArguments(AbsModelArguments):
+    """
+    Abstract class for reranker model arguments.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "The model checkpoint for initialization."}
+    )
+
+
+
+
+@dataclass
+class AbsRerankerDataArguments(AbsDataArguments):
+    """
+    Abstract class for reranker data arguments.
+    """
+    train_data: str = field(
+        default=None, metadata={
+            "help": "One or more paths to training data. `query: str`, `pos: List[str]`, `neg: List[str]` are required in the training data.",
+            "nargs": "+"
+        }
+    )
+    cache_path: Optional[str] = field(
+        default=None, metadata={"help": "Where do you want to store the cached data"}
+    )
+
+
+@dataclass
+class AbsRerankerTrainingArguments(AbsTrainingArguments):
+    pass

evaluation/nexus/Nexus/abc/training/reranker/AbsDataset.py

@@ -0,0 +1,25 @@
+import logging
+from dataclasses import dataclass
+from Nexus.abc.training.dataset import AbsDataset
+
+logger = logging.getLogger(__name__)
+
+
+class AbsRerankerTrainDataset(AbsDataset):
+    """Abstract class for reranker training dataset.
+
+    Args:
+        args (AbsRerankerDataArguments): Data arguments.
+        tokenizer (PreTrainedTokenizer): Tokenizer to use.
+    """
+    pass
+
+@dataclass
+class AbsRerankerCollator():
+    """
+    The abstract reranker collator.
+    """
+    def __call__(self, features):
+        return super().__call__(features)
+
+# remove Abs LLM reranker class

evaluation/nexus/Nexus/abc/training/reranker/AbsModeling.py

@@ -0,0 +1,77 @@
+import torch
+from torch import nn, Tensor
+from transformers import AutoTokenizer
+from transformers.file_utils import ModelOutput
+
+import logging
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, List, Union
+
+from Nexus.abc.training.modeling import AbsModelOutput, AbsModel, AbsReranker
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RerankerOutput(AbsModelOutput, ModelOutput):
+    scores: Optional[Tensor] = None
+    embedding: Optional[Tensor] = None
+    
+    def to_dict(self):
+        return {k: v for k, v in self.__dict__.items()}
+
+
+class AbsRerankerModel(AbsReranker, nn.Module):
+    """Abstract class of embedding model for training.
+
+    Args:
+        base_model: The base model to train on.
+        tokenizer (AutoTokenizer, optional): The tokenizer to use. Defaults to ``None``.
+        train_batch_size (int, optional): Batch size used for training. Defaults to ``4``.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.init_modules()
+
+    def init_modules(self):
+        self.loss_function=self.get_loss_function()
+        self.score_function=self.get_score_function()
+        
+    @abstractmethod
+    def get_loss_function(self):
+        pass
+    
+    @abstractmethod
+    def get_score_function(self):
+        pass
+
+
+    def forward(self, batch, *args, **kwargs):
+        """The computation performed at every call.
+        """
+        return self.compute_loss(batch, *args, **kwargs)
+
+    @abstractmethod
+    def compute_score(self, *args, **kwargs):
+        return super().compute_score(*args, **kwargs)
+
+    @abstractmethod
+    def compute_loss(self, *args, **kwargs):
+        """Compute the loss.
+        """
+        pass
+    
+    def save(self, output_dir: str):
+        """Save the model.
+
+        Args:
+            output_dir (str): Directory for saving the model.
+        """
+        pass
+
+    def save_pretrained(self, *args, **kwargs):
+        """
+        save model (and tokenizer if has)
+        """
+        pass

evaluation/nexus/Nexus/abc/training/reranker/AbsRunner.py

@@ -0,0 +1,75 @@
+import os
+import logging
+from pathlib import Path
+from typing import Tuple
+from abc import ABC, abstractmethod
+from transformers import set_seed, PreTrainedTokenizer
+
+
+from .AbsArguments import (
+    AbsRerankerModelArguments,
+    AbsRerankerDataArguments,
+    AbsRerankerTrainingArguments
+)
+from .AbsTrainer import AbsRerankerTrainer
+from .AbsModeling import AbsRerankerModel
+from .AbsDataset import (
+    AbsRerankerTrainDataset, AbsRerankerCollator
+)
+
+from Nexus.abc.training.runner import AbsRunner
+
+
+logger = logging.getLogger(__name__)
+
+
+class AbsRerankerRunner(AbsRunner):
+    """Abstract class to run reranker model fine-tuning.
+
+    Args:
+        model_args (AbsRerankerModelArguments): Model arguments
+        data_args (AbsRerankerDataArguments): Data arguments.
+        training_args (AbsRerankerTrainingArguments): Training arguments.
+    """
+    def __init__(
+        self,
+        model_args: AbsRerankerModelArguments,
+        data_args: AbsRerankerDataArguments,
+        training_args: AbsRerankerTrainingArguments
+    ):
+        self.model_args = model_args
+        self.data_args = data_args
+        self.training_args = training_args
+
+
+    @abstractmethod
+    def load_model(self, *args, **kwargs) -> AbsRerankerModel:
+        pass
+
+    @abstractmethod
+    def load_trainer(self) -> AbsRerankerTrainer:
+        """Abstract method to load the trainer.
+
+        Returns:
+            AbsRerankerTrainer: The loaded trainer instance.
+        """
+        pass
+
+    @abstractmethod
+    def load_dataset(self, *args, **kwargs):
+        pass
+    
+    def load_data_collator(self) -> AbsRerankerCollator:
+        """Loads the appropriate data collator.
+        """
+        pass
+
+    def run(self):
+        """
+        Executes the training process.
+        """
+        Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
+
+        # Training
+        self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
+        self.trainer.save_model()

evaluation/nexus/Nexus/abc/training/reranker/AbsTrainer.py

@@ -0,0 +1,41 @@
+import logging
+from typing import Optional
+from abc import ABC, abstractmethod
+from transformers.trainer import Trainer
+from Nexus.abc.training.trainer import AbsTrainer
+from Nexus.abc.training.dataset import AbsDataset
+
+logger = logging.getLogger(__name__)
+
+class AbsRerankerTrainer(AbsTrainer):
+    """
+    Abstract class for the trainer of reranker.
+    """
+    @abstractmethod
+    def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        pass
+
+
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        """
+        How the loss is computed by Trainer. By default, all models return the loss in the first element.
+
+        Subclass and override for custom behavior.
+        
+        Args:
+            model (AbsRerankerModel): The model being trained.
+            inputs (dict): A dictionary of input tensors to be passed to the model.
+            return_outputs (bool, optional): If ``True``, returns both the loss and the model's outputs. Otherwise,
+                returns only the loss. Defaults to ``False``.
+        
+        Returns:
+            Union[torch.Tensor, tuple(torch.Tensor, RerankerOutput)]: The computed loss. If ``return_outputs`` is ``True``, 
+                also returns the model's outputs in a tuple ``(loss, outputs)``.
+        """
+        inputs_dict={
+            'batch':inputs
+        }
+        outputs = model(**inputs_dict)
+        loss = outputs.loss
+
+        return (loss, outputs) if return_outputs else loss

evaluation/nexus/Nexus/abc/training/reranker/__init__.py

@@ -0,0 +1,19 @@
+from .AbsArguments import AbsRerankerDataArguments, AbsRerankerModelArguments, AbsRerankerTrainingArguments
+from .AbsDataset import (
+    AbsRerankerTrainDataset, AbsRerankerCollator
+)
+from .AbsModeling import AbsRerankerModel, RerankerOutput
+from .AbsTrainer import AbsRerankerTrainer
+from .AbsRunner import AbsRerankerRunner
+
+__all__ = [
+    "AbsRerankerDataArguments",
+    "AbsRerankerModelArguments",
+    "AbsRerankerTrainingArguments",
+    "AbsRerankerTrainDataset",
+    "AbsRerankerCollator",
+    "AbsRerankerModel",
+    "RerankerOutput",
+    "AbsRerankerTrainer",
+    "AbsRerankerRunner",
+]

evaluation/nexus/Nexus/abc/training/runner.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+
+from .dataset import AbsDataset
+from .modeling import AbsModel
+from .trainer import AbsTrainer
+from .arguments import AbsTrainingArguments, AbsDataArguments, AbsModelArguments
+
+
+class AbsRunner(ABC):
+    def __init__(
+        self,
+        model_args: AbsModelArguments,
+        data_args: AbsDataArguments,
+        training_args: AbsTrainingArguments,
+    ):
+        self.model_args = model_args
+        self.data_args = data_args
+        self.training_args = training_args
+
+    @abstractmethod
+    def load_model(self, *args, **kwargs) -> AbsModel:
+        pass
+
+    @abstractmethod
+    def load_dataset(self, *args, **kwargs) -> AbsDataset:
+        pass
+
+    @abstractmethod
+    def load_trainer(self, *args, **kwargs) -> AbsTrainer:
+        pass
+
+    @abstractmethod
+    def run(self, *args, **kwargs):
+        pass

evaluation/nexus/Nexus/abc/training/trainer.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+import inspect
+from typing import Optional
+
+import torch
+
+from .dataset import AbsDataset
+from .modeling import AbsModel
+from .arguments import AbsTrainingArguments
+from transformers import Trainer
+
+
+def _patch_accelerate_unwrap_model_signature():
+    try:
+        from accelerate import Accelerator
+    except ImportError:
+        return
+
+    signature = inspect.signature(Accelerator.unwrap_model)
+    if "keep_torch_compile" in signature.parameters:
+        return
+    if getattr(Accelerator.unwrap_model, "_nexus_keep_torch_compile_patch", False):
+        return
+
+    original_unwrap_model = Accelerator.unwrap_model
+
+    def compatible_unwrap_model(self, model, keep_fp32_wrapper: bool = True, keep_torch_compile: bool = False):
+        return original_unwrap_model(self, model, keep_fp32_wrapper=keep_fp32_wrapper)
+
+    compatible_unwrap_model._nexus_keep_torch_compile_patch = True
+    Accelerator.unwrap_model = compatible_unwrap_model
+
+
+_patch_accelerate_unwrap_model_signature()
+
+class AbsTrainer(Trainer):
+    @abstractmethod
+    def _save(self, output_dir: Optional[str] = None, *args, **kwargs):
+        return super()._save(output_dir, *args, **kwargs)
+
+    @torch.no_grad()
+    @abstractmethod
+    def evaluate(self, eval_dataset: AbsDataset, *args, **kwargs):
+        return super().evaluate(eval_dataset, **args, **kwargs)

evaluation/nexus/Nexus/data_mining/__init__.py

@@ -0,0 +1,2 @@
+"""Data mining utilities for Nexus."""
+

evaluation/nexus/Nexus/data_mining/multimodal_hard_negative/__init__.py

@@ -0,0 +1,6 @@
+"""Multimodal hard negative mining for Nexus training JSONL files."""
+
+from .mining import run_mining
+
+__all__ = ["run_mining"]
+