init this repo

.gitattributes

@@ -0,0 +1,40 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+*.axmodel filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/red-panda.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+build-output/
+tmp/
+*.safetensors

README.md

@@ -0,0 +1,97 @@
+---
+library_name: transformers
+license: mit
+base_model:
+  - Qwen/Qwen3-Embedding-0.6B
+tags:
+  - transformers
+  - sentence-transformers
+  - sentence-similarity
+  - feature-extraction
+  - text-embeddings-inference
+pipeline_tag: feature-extraction
+---
+
+# Qwen3-Embedding-0.6B
+
+This version of Qwen3-Embedding-0.6B has been converted to run on the Axera NPU using **w8a16** quantization.
+
+This model has been optimized with the following LoRA: 
+
+Compatible with Pulsar2 version: 4.1
+
+## Convert tools links:
+
+For those who are interested in model conversion, you can try to export axmodel through the original repo:
+https://huggingface.co/Qwen/Qwen3-Embedding-0.6B
+
+[Pulsar2 Link, How to Convert LLM from Huggingface to axmodel](https://pulsar2-docs.readthedocs.io/en/latest/appendix/build_llm.html)
+
+## Support Platform
+
+- AX650
+  - AX650N DEMO Board
+  - [M4N-Dock(爱芯派Pro)](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html)
+  - [M.2 Accelerator card](https://axcl-docs.readthedocs.io/zh-cn/latest/doc_guide_hardware.html)
+
+### Each subgraph is time-consuming
+```sh
+g1: 5.561 ms
+g2: 9.140 ms
+g3: 12.757 ms
+g4: 16.446 ms
+g5: 21.392 ms
+g6: 23.712 ms
+g7: 27.174 ms
+g8: 30.897 ms
+g9: 34.829 ms
+```
+|Chips | ttft | w8a16 |
+|--|--|--|--|
+|AX650| 155.708 ms | 6.42 tokens/sec|
+|AX650| 5093.42 ms | 0.19 tokens/sec|
+
+Longest time consumption: 181.908 ms
+Shortest time consumption: 5.561 ms
+LayerNum: 28
+
+## How to use
+
+Download all files from this repository to the device.
+
+**If you using AX650 Board**
+```
+root@ax650 ~/yongqiang/push_hugging_face/Qwen3-Embedding-0.6B # tree -L 1
+.
+├── config.json
+├── infer_axmodel.py
+├── qwen3_embedding_0.6b_axmodel
+├── qwen3_embedding_0.6b_tokenizer
+├── README.md
+└── utils
+
+3 directories, 3 files
+```
+
+#### Install transformer
+
+```
+# Requires transformers>=4.51.0
+pip install transformers==4.51.0
+```
+
+#### Inference with AX650 Host, such as M4N-Dock(爱芯派Pro) or AX650N DEMO Board
+
+```
+$ python3 infer_axmodel.py
+Model loaded successfully!
+slice_indices: [0]
+Slice prefill done: 0
+slice_indices: [0]
+Slice prefill done: 0
+slice_indices: [0]
+Slice prefill done: 0
+slice_indices: [0]
+Slice prefill done: 0
+[[0.7555467486381531, 0.1756950318813324], [0.4137178063392639, 0.4459586441516876]]
+```

infer_axmodel.py

@@ -0,0 +1,100 @@
+# Requires transformers>=4.51.0
+import os
+import numpy as np
+from axengine import InferenceSession
+from ml_dtypes import bfloat16
+from utils.infer_func import InferManager
+import argparse
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel, AutoConfig
+
+
+def last_token_pool(last_hidden_states: Tensor,
+                 attention_mask: Tensor) -> Tensor:
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+
+def get_detailed_instruct(task_description: str, query: str) -> str:
+    return f'Instruct: {task_description}\nQuery:{query}'
+
+
+if __name__ == "__main__":
+
+    """
+    python3 infer_axmodel.py
+    """
+
+    prompt = None
+    parser = argparse.ArgumentParser(description="Model configuration parameters")
+    parser.add_argument("--hf_model", type=str, default="./qwen3_embedding_0.6b_tokenizer/",
+                        help="Path to HuggingFace model")
+    parser.add_argument("--axmodel_path", type=str, default="./qwen3_embedding_0.6b_axmodel/",
+                        help="Path to save compiled axmodel of llama model")
+    args = parser.parse_args()
+
+    hf_model_path = args.hf_model
+    axmodel_path = args.axmodel_path
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    embeds = np.load(os.path.join(axmodel_path, "model.embed_tokens.weight.npy"))
+
+    # Each query must come with a one-sentence instruction that describes the task
+    task = 'Given a web search query, retrieve relevant passages that answer the query'
+
+    queries = [
+        get_detailed_instruct(task, 'What is the capital of China?'),
+        get_detailed_instruct(task, 'Explain gravity')
+    ]
+    # No need to add instruction for retrieval documents
+    documents = [
+        "The capital of China is Beijing.",
+        "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
+    ]
+    input_texts = queries + documents
+
+    tokenizer = AutoTokenizer.from_pretrained(hf_model_path, padding_side='left')
+    cfg = AutoConfig.from_pretrained(hf_model_path)
+
+    max_length = 8192
+
+    # Tokenize the input texts
+    batch_dict = tokenizer(
+        input_texts,
+        padding=True,
+        truncation=True,
+        max_length=max_length,
+        return_tensors="pt",
+    )
+    batch_dict.to(device)
+
+    input_ids = batch_dict['input_ids']
+    attention_mask = batch_dict['attention_mask']
+    inputs_embeds = np.take(embeds, input_ids.cpu().numpy(), axis=0)
+    prefill_data = inputs_embeds
+    prefill_data = prefill_data.astype(bfloat16)
+    token_ids = input_ids[0].cpu().numpy().tolist()
+    token_len = len(token_ids)
+
+    batch_num, seq_len, seq_dim = inputs_embeds.shape
+
+    imer = InferManager(cfg, axmodel_path)
+    last_hidden_state = np.zeros((batch_num, seq_len, seq_dim), dtype=bfloat16)
+    for batch_idx in range(batch_num):
+        last_hidden_state[batch_idx] = imer.prefill(tokenizer, token_ids, prefill_data[batch_idx], slice_len=128)
+
+    embeddings = last_token_pool(torch.from_numpy(last_hidden_state.astype(np.float32)), batch_dict['attention_mask'])
+    # normalize embeddings
+    embeddings = F.normalize(embeddings, p=2, dim=1)
+    scores = (embeddings[:2] @ embeddings[2:].T)
+    print(scores.tolist())
+    # Torch: [[0.7645568251609802, 0.14142508804798126], [0.13549736142158508, 0.5999549627304077]]
+    # Axmod: [[0.755547046661377, 0.17569507658481598], [0.4137181341648102, 0.4459584951400757]]

qwen3_embedding_0.6b_axmodel/model.embed_tokens.weight.bfloat16.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a55b140d86852835bd18d8200222a9f302340730f0670eb7e23a4895e5489033
+size 310618112

qwen3_embedding_0.6b_axmodel/model.embed_tokens.weight.float32.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7a027c062fb61cd505e046bc832345be155e1eb2fab629675cebe7973646c85
+size 621236224

qwen3_embedding_0.6b_axmodel/model.embed_tokens.weight.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bbdc47aee1b4cdb97a42a255306d4e0a1cb52f797bfdc32f94469eb0cd0744e
+size 621236352

qwen3_embedding_0.6b_axmodel/qwen3_p128_l0_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:025e17a92f3f19d58a36ef119294598073c4ccdc794aa9d4a2845a99b0c6b53d
+size 28019747

qwen3_embedding_0.6b_axmodel/qwen3_p128_l10_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95d809d6cc6889517b1aad7a4e62e51ffbf75580dda5ceafb667dbd5ac10ba6e
+size 28019779

qwen3_embedding_0.6b_axmodel/qwen3_p128_l11_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fb730805877eceea4aa037694bc2abb830fa960c6666d24b976d7ae35c058d0
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l12_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eef7266dff6af522a0a63095067a1c7823a9a1213e7bd498bcdb97f2814523ba
+size 28019427

qwen3_embedding_0.6b_axmodel/qwen3_p128_l13_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9322be28dad0729b006238364297659594e2193a97a37556fb06f63d3fec9fa0
+size 28019459

qwen3_embedding_0.6b_axmodel/qwen3_p128_l14_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5a8f239e4a4e793e0bdf86226c08c2089f5199118bd38d2be4957f9b7023dda
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l15_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8767d8ce2a3a20049f30da7982b3cb15799e6e0959b3deb484db15fef426010b
+size 28019619

qwen3_embedding_0.6b_axmodel/qwen3_p128_l16_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a80f39502798ca50e070ace41527fb176232ce315d6175ea65ec8e65467112fa
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l17_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:108d4abc6ceef53bf17786e0512d6c4e2de8d531cc14b1192fb9923ceeb2c10f
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l18_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8e07c47bd98d9062c3c5bc48fc7057f363cbb5b31c62540346814299b1843f4
+size 28019299

qwen3_embedding_0.6b_axmodel/qwen3_p128_l19_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a713464dcbd8187f27cecc4d6917acb8726ab9cffeaf6ed1e9c2bed6059cdd0
+size 28019075

qwen3_embedding_0.6b_axmodel/qwen3_p128_l1_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c37fefd3be1f8699961f1b5d6c164a74960ce90a2b219514add7ce960534a292
+size 28019363

qwen3_embedding_0.6b_axmodel/qwen3_p128_l20_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad7cd29e2ea7644f6ff39a179efb6b646743fc47a3c4049877db258b8072fbd1
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l21_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff616a72c372edc0068909aea76abd76b1c5121fa86e0052d93dfe42fdd247b7
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l22_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c77fb82587bbc3c682c28cae6bf3cc3ac290bd59ffd60b6f03c875970c3518b8
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l23_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c1d6057741c34a9b871b4ae2852c3e9616b2423abd9b1035a26c003d0158645
+size 28019491

qwen3_embedding_0.6b_axmodel/qwen3_p128_l24_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b14f1c52826b17c7eb54f6e2ab9c0a5c353cf34b0f90a252d2c3d13114d0a284
+size 28020259

qwen3_embedding_0.6b_axmodel/qwen3_p128_l25_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b81f7122c92a9b04a19581d4bbba8f8392037e79db1ae3ead467f2f6dcb9975
+size 28019395

qwen3_embedding_0.6b_axmodel/qwen3_p128_l26_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c5a70d619eebbd1b8a072902c807022c2cabd421018913fb47a44627132eb98
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l27_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26b447aa2df812d621e272c51115d94d46c8f8cb9bcce9281d86d2f9c2612211
+size 28020547

qwen3_embedding_0.6b_axmodel/qwen3_p128_l2_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:167e1b7111d86d7d0aba4879314fc59b93b2385242caddea37ae67c54b34cc4e
+size 28019107

qwen3_embedding_0.6b_axmodel/qwen3_p128_l3_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d929fe10751a50dae6fe2379149f912709b55f597c20686bcf2b4c812bb46d9
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l4_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:40dbedcb1f86e618d2f7cf68d44dfce85e120cdea9475ebb43f1d04852a536ed
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l5_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ee25b58ea5b0ae8a14a031c6f8bd87c826cf5419d9be3c542b9032cfb18b929
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l6_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f20cdb2296614658bc35171dda33c9ef84d71771620955dd30fe0b700f88cf0a
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l7_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8b81646a9906eef6eb69eda019cc8294db770dbaca915b7ea0e3a378f35f988
+size 28019555

qwen3_embedding_0.6b_axmodel/qwen3_p128_l8_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba4c0c875d63cafaff17f77378fa50a4ae4ad87e7936ffc3f15b06082d8bb68c
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_p128_l9_together.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b1d57cb66e5e04ddce9a5c8f7b308cc7995fedb4f92e9038260939d97d151cf
+size 28018723

qwen3_embedding_0.6b_axmodel/qwen3_post.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b4557f9ba5f22fce6e1584d4b1119e0bbd4dc0f0d241233d1e8ee7f39775757
+size 169711178

qwen3_embedding_0.6b_tokenizer/.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

qwen3_embedding_0.6b_tokenizer/1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+    "word_embedding_dimension": 1024,
+    "pooling_mode_cls_token": false,
+    "pooling_mode_mean_tokens": false,
+    "pooling_mode_max_tokens": false,
+    "pooling_mode_mean_sqrt_len_tokens": false,
+    "pooling_mode_weightedmean_tokens": false,
+    "pooling_mode_lasttoken": true,
+    "include_prompt": true
+}

qwen3_embedding_0.6b_tokenizer/README.md

@@ -0,0 +1,292 @@
+---
+license: apache-2.0
+base_model:
+- Qwen/Qwen3-0.6B-Base
+tags:
+- transformers
+- sentence-transformers
+- sentence-similarity
+- feature-extraction
+- text-embeddings-inference
+---
+# Qwen3-Embedding-0.6B
+
+<p align="center">
+    <img src="https://qianwen-res.oss-accelerate-overseas.aliyuncs.com/logo_qwen3.png" width="400"/>
+<p>
+
+## Highlights
+
+The Qwen3 Embedding model series is the latest proprietary model of the Qwen family, specifically designed for text embedding and ranking tasks. Building upon the dense foundational models of the Qwen3 series, it provides a comprehensive range of text embeddings and reranking models in various sizes (0.6B, 4B, and 8B). This series inherits the exceptional multilingual capabilities, long-text understanding, and reasoning skills of its foundational model. The Qwen3 Embedding series represents significant advancements in multiple text embedding and ranking tasks, including text retrieval, code retrieval, text classification, text clustering, and bitext mining.
+
+**Exceptional Versatility**: The embedding model has achieved state-of-the-art performance across a wide range of downstream application evaluations. The 8B size embedding model ranks **No.1** in the MTEB multilingual leaderboard (as of June 5, 2025, score **70.58**), while the reranking model excels in various text retrieval scenarios.
+
+**Comprehensive Flexibility**: The Qwen3 Embedding series offers a full spectrum of sizes (from 0.6B to 8B) for both embedding and reranking models, catering to diverse use cases that prioritize efficiency and effectiveness. Developers can seamlessly combine these two modules. Additionally, the embedding model allows for flexible vector definitions across all dimensions, and both embedding and reranking models support user-defined instructions to enhance performance for specific tasks, languages, or scenarios.
+
+**Multilingual Capability**: The Qwen3 Embedding series offer support for over 100 languages, thanks to the multilingual capabilites of Qwen3 models. This includes various programming languages, and provides robust multilingual, cross-lingual, and code retrieval capabilities.
+
+## Model Overview
+
+**Qwen3-Embedding-0.6B** has the following features:
+
+- Model Type: Text Embedding
+- Supported Languages: 100+ Languages
+- Number of Paramaters: 0.6B
+- Context Length: 32k
+- Embedding Dimension: Up to 1024, supports user-defined output dimensions ranging from 32 to 1024
+
+For more details, including benchmark evaluation, hardware requirements, and inference performance, please refer to our [blog](https://qwenlm.github.io/blog/qwen3-embedding/), [GitHub](https://github.com/QwenLM/Qwen3-Embedding).
+
+## Qwen3 Embedding Series Model list
+
+| Model Type       | Models               | Size | Layers | Sequence Length | Embedding Dimension | MRL Support | Instruction Aware |
+|------------------|----------------------|------|--------|-----------------|---------------------|-------------|----------------|
+| Text Embedding   | [Qwen3-Embedding-0.6B](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B) | 0.6B | 28     | 32K             | 1024                | Yes         | Yes            |
+| Text Embedding   | [Qwen3-Embedding-4B](https://huggingface.co/Qwen/Qwen3-Embedding-4B)   | 4B   | 36     | 32K             | 2560                | Yes         | Yes            |
+| Text Embedding   | [Qwen3-Embedding-8B](https://huggingface.co/Qwen/Qwen3-Embedding-8B)   | 8B   | 36     | 32K             | 4096                | Yes         | Yes            |
+| Text Reranking   | [Qwen3-Reranker-0.6B](https://huggingface.co/Qwen/Qwen3-Reranker-0.6B) | 0.6B | 28     | 32K             | -                   | -           | Yes            |
+| Text Reranking   | [Qwen3-Reranker-4B](https://huggingface.co/Qwen/Qwen3-Reranker-4B)   | 4B   | 36     | 32K             | -                   | -           | Yes            |
+| Text Reranking   | [Qwen3-Reranker-8B](https://huggingface.co/Qwen/Qwen3-Reranker-8B)   | 8B   | 36     | 32K             | -                   | -           | Yes            |
+
+> **Note**:
+> - `MRL Support` indicates whether the embedding model supports custom dimensions for the final embedding. 
+> - `Instruction Aware` notes whether the embedding or reranking model supports customizing the input instruction according to different tasks.
+> - Our evaluation indicates that, for most downstream tasks, using instructions (instruct) typically yields an improvement of 1% to 5% compared to not using them. Therefore, we recommend that developers create tailored instructions specific to their tasks and scenarios. In multilingual contexts, we also advise users to write their instructions in English, as most instructions utilized during the model training process were originally written in English.
+
+## Usage
+
+With Transformers versions earlier than 4.51.0, you may encounter the following error:
+```
+KeyError: 'qwen3'
+```
+
+### Sentence Transformers Usage
+
+```python
+# Requires transformers>=4.51.0
+# Requires sentence-transformers>=2.7.0
+
+from sentence_transformers import SentenceTransformer
+
+# Load the model
+model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B")
+
+# We recommend enabling flash_attention_2 for better acceleration and memory saving,
+# together with setting `padding_side` to "left":
+# model = SentenceTransformer(
+#     "Qwen/Qwen3-Embedding-0.6B",
+#     model_kwargs={"attn_implementation": "flash_attention_2", "device_map": "auto"},
+#     tokenizer_kwargs={"padding_side": "left"},
+# )
+
+# The queries and documents to embed
+queries = [
+    "What is the capital of China?",
+    "Explain gravity",
+]
+documents = [
+    "The capital of China is Beijing.",
+    "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
+]
+
+# Encode the queries and documents. Note that queries benefit from using a prompt
+# Here we use the prompt called "query" stored under `model.prompts`, but you can
+# also pass your own prompt via the `prompt` argument
+query_embeddings = model.encode(queries, prompt_name="query")
+document_embeddings = model.encode(documents)
+
+# Compute the (cosine) similarity between the query and document embeddings
+similarity = model.similarity(query_embeddings, document_embeddings)
+print(similarity)
+# tensor([[0.7646, 0.1414],
+#         [0.1355, 0.6000]])
+```
+
+### Transformers Usage
+
+```python
+# Requires transformers>=4.51.0
+
+import torch
+import torch.nn.functional as F
+
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel
+
+
+def last_token_pool(last_hidden_states: Tensor,
+                 attention_mask: Tensor) -> Tensor:
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+
+def get_detailed_instruct(task_description: str, query: str) -> str:
+    return f'Instruct: {task_description}\nQuery:{query}'
+
+# Each query must come with a one-sentence instruction that describes the task
+task = 'Given a web search query, retrieve relevant passages that answer the query'
+
+queries = [
+    get_detailed_instruct(task, 'What is the capital of China?'),
+    get_detailed_instruct(task, 'Explain gravity')
+]
+# No need to add instruction for retrieval documents
+documents = [
+    "The capital of China is Beijing.",
+    "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
+]
+input_texts = queries + documents
+
+tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B', padding_side='left')
+model = AutoModel.from_pretrained('Qwen/Qwen3-Embedding-0.6B')
+
+# We recommend enabling flash_attention_2 for better acceleration and memory saving.
+# model = AutoModel.from_pretrained('Qwen/Qwen3-Embedding-0.6B', attn_implementation="flash_attention_2", torch_dtype=torch.float16).cuda()
+
+max_length = 8192
+
+# Tokenize the input texts
+batch_dict = tokenizer(
+    input_texts,
+    padding=True,
+    truncation=True,
+    max_length=max_length,
+    return_tensors="pt",
+)
+batch_dict.to(model.device)
+outputs = model(**batch_dict)
+embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+
+# normalize embeddings
+embeddings = F.normalize(embeddings, p=2, dim=1)
+scores = (embeddings[:2] @ embeddings[2:].T)
+print(scores.tolist())
+# [[0.7645568251609802, 0.14142508804798126], [0.13549736142158508, 0.5999549627304077]]
+```
+
+### vLLM Usage
+
+```python
+# Requires vllm>=0.8.5
+import torch
+import vllm
+from vllm import LLM
+
+def get_detailed_instruct(task_description: str, query: str) -> str:
+    return f'Instruct: {task_description}\nQuery:{query}'
+
+# Each query must come with a one-sentence instruction that describes the task
+task = 'Given a web search query, retrieve relevant passages that answer the query'
+
+queries = [
+    get_detailed_instruct(task, 'What is the capital of China?'),
+    get_detailed_instruct(task, 'Explain gravity')
+]
+# No need to add instruction for retrieval documents
+documents = [
+    "The capital of China is Beijing.",
+    "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
+]
+input_texts = queries + documents
+
+model = LLM(model="Qwen/Qwen3-Embedding-0.6B", task="embed")
+
+outputs = model.embed(input_texts)
+embeddings = torch.tensor([o.outputs.embedding for o in outputs])
+scores = (embeddings[:2] @ embeddings[2:].T)
+print(scores.tolist())
+# [[0.7620252966880798, 0.14078938961029053], [0.1358368694782257, 0.6013815999031067]]
+```
+
+📌 **Tip**: We recommend that developers customize the `instruct` according to their specific scenarios, tasks, and languages. Our tests have shown that in most retrieval scenarios, not using an `instruct` on the query side can lead to a drop in retrieval performance by approximately 1% to 5%.
+
+### Text Embeddings Inference (TEI) Usage
+
+You can either run / deploy TEI on NVIDIA GPUs as:
+
+```bash
+docker run --gpus all -p 8080:80 -v hf_cache:/data --pull always ghcr.io/huggingface/text-embeddings-inference:cpu-1.7.2 --model-id Qwen/Qwen3-Embedding-0.6B --dtype float16
+```
+
+Or on CPU devices as:
+
+```bash
+docker run -p 8080:80 -v hf_cache:/data --pull always ghcr.io/huggingface/text-embeddings-inference:1.7.2 --model-id Qwen/Qwen3-Embedding-0.6B
+```
+
+And then, generate the embeddings sending a HTTP POST request as:
+
+```bash
+curl http://localhost:8080/embed \
+    -X POST \
+    -d '{"inputs": ["Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: What is the capital of China?", "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: Explain gravity"]}' \
+    -H "Content-Type: application/json"
+```
+
+## Evaluation
+
+### MTEB (Multilingual)
+
+| Model                            |  Size   |  Mean (Task)  | Mean (Type) | Bitxt Mining | Class. | Clust. | Inst. Retri. | Multi. Class. | Pair. Class. | Rerank | Retri. | STS  |
+|----------------------------------|:-------:|:-------------:|:-------------:|:--------------:|:--------:|:--------:|:--------------:|:---------------:|:--------------:|:--------:|:--------:|:------:|
+| NV-Embed-v2                      |   7B    |     56.29     | 49.58       | 57.84        | 57.29  | 40.80  | 1.04         | 18.63         | 78.94        | 63.82  | 56.72  | 71.10|
+| GritLM-7B                        |   7B    |     60.92     | 53.74       | 70.53        | 61.83  | 49.75  | 3.45         | 22.77         | 79.94        | 63.78  | 58.31  | 73.33|
+| BGE-M3                           |  0.6B   |     59.56     | 52.18       | 79.11        | 60.35  | 40.88  | -3.11        | 20.1          | 80.76        | 62.79  | 54.60  | 74.12|
+| multilingual-e5-large-instruct   |  0.6B   |     63.22     | 55.08       | 80.13        | 64.94  | 50.75  | -0.40        | 22.91         | 80.86        | 62.61  | 57.12  | 76.81|
+| gte-Qwen2-1.5B-instruct          |  1.5B   |     59.45     | 52.69       | 62.51        | 58.32  | 52.05  | 0.74         | 24.02         | 81.58        | 62.58  | 60.78  | 71.61|
+| gte-Qwen2-7b-Instruct            |   7B    |     62.51     | 55.93       | 73.92        | 61.55  | 52.77  | 4.94         | 25.48         | 85.13        | 65.55  | 60.08  | 73.98|
+| text-embedding-3-large           |    -    |     58.93     | 51.41       | 62.17        | 60.27  | 46.89  | -2.68        | 22.03         | 79.17        | 63.89  | 59.27  | 71.68|
+| Cohere-embed-multilingual-v3.0   |    -    |     61.12     | 53.23       | 70.50        | 62.95  | 46.89  | -1.89        | 22.74         | 79.88        | 64.07  | 59.16  | 74.80|
+| Gemini Embedding                 |    -    |     68.37     | 59.59       | 79.28        | 71.82  | 54.59  | 5.18         | **29.16**     | 83.63        | 65.58  | 67.71  | 79.40|
+| **Qwen3-Embedding-0.6B**         |  0.6B   |     64.33     | 56.00       | 72.22        | 66.83  | 52.33  | 5.09         | 24.59         | 80.83        | 61.41  | 64.64  | 76.17|
+| **Qwen3-Embedding-4B**           |   4B    |     69.45     | 60.86       | 79.36        | 72.33  | 57.15  | **11.56**    | 26.77         | 85.05        | 65.08  | 69.60  | 80.86|
+| **Qwen3-Embedding-8B**           |   8B    |   **70.58**   | **61.69**   | **80.89**    | **74.00** | **57.65** | 10.06      | 28.66         | **86.40**    | **65.63** | **70.88** | **81.08** |
+
+> **Note**: For compared models, the scores are retrieved from MTEB online [leaderboard](https://huggingface.co/spaces/mteb/leaderboard) on May 24th, 2025.
+
+### MTEB (Eng v2)
+
+| MTEB English / Models          |  Param.  | Mean(Task) | Mean(Type) | Class. | Clust. | Pair Class. | Rerank. | Retri. | STS   | Summ. |
+|--------------------------------|:--------:|:------------:|:------------:|:--------:|:--------:|:-------------:|:---------:|:--------:|:-------:|:-------:|
+| multilingual-e5-large-instruct |   0.6B   | 65.53      | 61.21      | 75.54  | 49.89  | 86.24       | 48.74   | 53.47  | 84.72 | 29.89 |
+| NV-Embed-v2                    |   7.8B   | 69.81      | 65.00      | 87.19  | 47.66  | 88.69       | 49.61   | 62.84  | 83.82 | 35.21 |
+| GritLM-7B                      |   7.2B   | 67.07      | 63.22      | 81.25  | 50.82  | 87.29       | 49.59   | 54.95  | 83.03 | 35.65 |
+| gte-Qwen2-1.5B-instruct        |   1.5B   | 67.20      | 63.26      | 85.84  | 53.54  | 87.52       | 49.25   | 50.25  | 82.51 | 33.94 |
+| stella_en_1.5B_v5              |   1.5B   | 69.43      | 65.32      | 89.38  | 57.06  | 88.02       | 50.19   | 52.42  | 83.27 | 36.91 |
+| gte-Qwen2-7B-instruct          |   7.6B   | 70.72      | 65.77      | 88.52  | 58.97  | 85.9        | 50.47   | 58.09  | 82.69 | 35.74 |
+| gemini-embedding-exp-03-07     |    -     | 73.3       | 67.67      | 90.05  | 59.39  | 87.7        | 48.59   | 64.35  | 85.29 | 38.28 |
+| **Qwen3-Embedding-0.6B**       |   0.6B   | 70.70      | 64.88      | 85.76  | 54.05  | 84.37       | 48.18   | 61.83  | 86.57 | 33.43 |
+| **Qwen3-Embedding-4B**         |    4B    | 74.60      | 68.10      | 89.84  | 57.51  | 87.01       | 50.76   | 68.46  | 88.72 | 34.39 |
+| **Qwen3-Embedding-8B**         |    8B    | 75.22      | 68.71      | 90.43  | 58.57  | 87.52       | 51.56   | 69.44  | 88.58 | 34.83 |
+
+### C-MTEB (MTEB Chinese)
+
+| C-MTEB           | Param. | Mean(Task) | Mean(Type) | Class. | Clust. | Pair Class. | Rerank. | Retr. | STS   |
+|------------------|--------|------------|------------|--------|--------|-------------|---------|-------|-------|
+| multilingual-e5-large-instruct | 0.6B   | 58.08      | 58.24      | 69.80  | 48.23  | 64.52       | 57.45   | 63.65 | 45.81 |
+| bge-multilingual-gemma2 | 9B     | 67.64      | 75.31      | 59.30  | 86.67  | 68.28       | 73.73   | 55.19 | -     |
+| gte-Qwen2-1.5B-instruct  | 1.5B   | 67.12      | 67.79      | 72.53  | 54.61  | 79.5        | 68.21   | 71.86 | 60.05 |
+| gte-Qwen2-7B-instruct    | 7.6B   | 71.62      | 72.19      | 75.77  | 66.06  | 81.16       | 69.24   | 75.70 | 65.20 |
+| ritrieve_zh_v1          | 0.3B   | 72.71      | 73.85      | 76.88  | 66.5   | 85.98       | 72.86   | 76.97 | 63.92 |
+| **Qwen3-Embedding-0.6B** | 0.6B   | 66.33      | 67.45      | 71.40  | 68.74  | 76.42       | 62.58   | 71.03 | 54.52 |
+| **Qwen3-Embedding-4B**   | 4B     | 72.27      | 73.51      | 75.46  | 77.89  | 83.34       | 66.05   | 77.03 | 61.26 |
+| **Qwen3-Embedding-8B**   | 8B     | 73.84      | 75.00      | 76.97  | 80.08  | 84.23       | 66.99   | 78.21 | 63.53 |
+
+
+## Citation
+
+If you find our work helpful, feel free to give us a cite.
+
+```
+@article{qwen3embedding,
+  title={Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models},
+  author={Zhang, Yanzhao and Li, Mingxin and Long, Dingkun and Zhang, Xin and Lin, Huan and Yang, Baosong and Xie, Pengjun and Yang, An and Liu, Dayiheng and Lin, Junyang and Huang, Fei and Zhou, Jingren},
+  journal={arXiv preprint arXiv:2506.05176},
+  year={2025}
+}
+```

qwen3_embedding_0.6b_tokenizer/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 28,
+  "model_type": "qwen3",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151669
+}

qwen3_embedding_0.6b_tokenizer/config_sentence_transformers.json

@@ -0,0 +1,8 @@
+{
+  "prompts": {
+    "query": "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:",
+    "document": ""
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine"
+}

qwen3_embedding_0.6b_tokenizer/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "max_new_tokens": 2048,
+  "transformers_version": "4.51.3"
+}

qwen3_embedding_0.6b_tokenizer/modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.models.Normalize"
+  }
+]

qwen3_embedding_0.6b_tokenizer/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:def76fb086971c7867b829c23a26261e38d9d74e02139253b38aeb9df8b4b50a
+size 11423705

qwen3_embedding_0.6b_tokenizer/tokenizer_config.json

@@ -0,0 +1,240 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set content = message.content %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in message.content %}\n                {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

utils/infer_func.py

@@ -0,0 +1,217 @@
+import torch
+import numpy as np
+from typing import List, Tuple
+from tqdm import tqdm
+from axengine import InferenceSession
+from ml_dtypes import bfloat16
+
+
+class InferManager:
+    def __init__(self, config, model_dir):
+
+        self.config = config
+        self.max_seq_len = 2559
+
+        self.sub_dim = config.hidden_size // config.num_attention_heads if not config.head_dim else config.head_dim
+        self.kv_dim = self.sub_dim * config.num_key_value_heads
+
+        self.k_caches = [
+            np.zeros((1, self.max_seq_len, self.kv_dim), dtype=bfloat16)
+            for _ in range(config.num_hidden_layers)
+        ]
+        self.v_caches = [
+            np.zeros((1, self.max_seq_len, self.kv_dim), dtype=bfloat16)
+            for _ in range(config.num_hidden_layers)
+        ]
+
+        self.decoder_sessions = []
+        for layer_idx in tqdm(range(config.num_hidden_layers), desc="Init InferenceSession"):
+            session = InferenceSession(
+                f"{model_dir}/qwen3_p128_l{layer_idx}_together.axmodel"
+            )
+            self.decoder_sessions.append(session)
+        self.post_process_session = InferenceSession(
+            f"{model_dir}/qwen3_post.axmodel"
+        )
+        print("Model loaded successfully!")
+
+    @staticmethod
+    def _top_p(probs: np.ndarray, p: float) -> np.ndarray:
+        sorted_indices = np.argsort(probs)
+        filtered = probs.copy()
+        cumulative = 0
+        for idx in sorted_indices[::-1]:
+            if cumulative >= p:
+                filtered[idx] = 0
+            cumulative += filtered[idx]
+        return filtered / cumulative
+
+    @staticmethod
+    def _softmax(logits: np.ndarray) -> np.ndarray:
+        logits = logits - logits.max()
+        exp_logits = np.exp(logits)
+        return (exp_logits / np.sum(exp_logits)).astype(np.float64)
+
+    def post_process(self, logits, top_k=1, top_p=0.9, temperature=0.6):
+        logits = logits.astype(np.float32).flatten()
+        candidate_indices = np.argpartition(logits, -top_k)[-top_k:]
+        candidate_logits = logits[candidate_indices] / temperature
+        candidate_probs = self._softmax(candidate_logits)
+        candidate_probs = self._top_p(candidate_probs, top_p)
+        candidate_probs = candidate_probs.astype(np.float64) / candidate_probs.sum()
+        chosen_idx = np.random.multinomial(1, candidate_probs).argmax()
+        next_token = candidate_indices[chosen_idx]
+        return next_token, candidate_indices, candidate_probs
+
+    def gen_slice_indices(self, token_len, prefill=128, expand=128):
+        remaining = max(0, token_len - prefill)
+        extra_blocks = (remaining + expand - 1) // expand
+        return list(range(extra_blocks + 1))
+
+    def prefill(
+        self,
+        tokenizer,
+        token_ids,
+        embed_data,
+        slice_len=128,
+    ):
+        """
+        Prefill step for chunked inference.
+        """
+        seq_len = len(token_ids)
+        slice_indices = [i for i in range(seq_len // slice_len + 1)]
+        print(f"slice_indices: {slice_indices}")
+        # total_prefill_len = (
+        #     slice_len * slice_indices[-1]
+        #     if slice_indices[-1] != 0
+        #     else slice_len
+        # )
+        total_prefill_len = slice_len * (slice_indices[-1] + 1)
+        # slice_indices = self.gen_slice_indices(seq_len)
+        # import pdb; pdb.set_trace()
+
+        if total_prefill_len > 0:
+            for slice_idx in slice_indices:
+                indices = np.arange(
+                    slice_idx * slice_len,
+                    (slice_idx + 1) * slice_len,
+                    dtype=np.uint32
+                ).reshape((1, slice_len))
+
+                mask = (
+                    np.zeros((1, slice_len, slice_len * (slice_idx + 1)))
+                    - 65536
+                )
+                data = np.zeros((1, slice_len, self.config.hidden_size)).astype(bfloat16)
+                for i, t in enumerate(
+                    range(
+                        slice_idx * slice_len,
+                        (slice_idx + 1) * slice_len,
+                    )
+                ):
+                    if t < len(token_ids):
+                        mask[:, i, : slice_idx * slice_len + i + 1] = 0
+                        data[:, i : i + 1, :] = (
+                            embed_data[t]
+                            .reshape((1, 1, self.config.hidden_size))
+                            .astype(bfloat16)
+                        )
+
+                remain_len = (
+                    seq_len - slice_idx * slice_len
+                    if slice_idx == slice_indices[-1]
+                    else slice_len
+                )
+                mask = mask.astype(bfloat16)
+                for layer_idx in range(self.config.num_hidden_layers):
+                    input_feed = {
+                        "K_cache": (
+                            self.k_caches[layer_idx][:, 0 : slice_len * slice_idx, :]
+                            if slice_idx
+                            else np.zeros((1, 1, self.config.hidden_size), dtype=bfloat16)
+                        ),
+                        "V_cache": (
+                            self.v_caches[layer_idx][:, 0 : slice_len * slice_idx, :]
+                            if slice_idx
+                            else np.zeros((1, 1, self.config.hidden_size), dtype=bfloat16)
+                        ),
+                        "indices": indices,
+                        "input": data,
+                        "mask": mask,
+                    }
+                    outputs = self.decoder_sessions[layer_idx].run(None, input_feed, shape_group=slice_idx + 1)
+                    self.k_caches[layer_idx][
+                        :,
+                        slice_idx * slice_len : slice_idx * slice_len + remain_len,
+                        :,
+                    ] = outputs[0][:, :remain_len, :]
+                    self.v_caches[layer_idx][
+                        :,
+                        slice_idx * slice_len : slice_idx * slice_len + remain_len,
+                        :,
+                    ] = outputs[1][:, :remain_len, :]
+                    data = outputs[2]
+
+                print("Slice prefill done:", slice_idx)
+            
+            return data[:, :remain_len, :]
+            # post_out = self.post_process_session.run(
+            #     None,
+            #     {
+            #         "input": data[
+            #             :, seq_len - (len(slice_indices) - 1) * slice_len - 1, None, :
+            #         ]
+            #     }
+            # )[0]
+            # next_token, possible_tokens, possible_probs = self.post_process(post_out)
+            # possible_decoded = [tokenizer.decode([t]) for t in possible_tokens]
+            # possible_probs_str = [str((t, p)) for t, p in zip(possible_decoded, possible_probs)]
+            # token_ids.append(next_token)
+            # return token_ids
+
+    def decode(
+        self,
+        tokenizer,
+        token_ids,
+        embed_matrix,
+        prefill_len=128,
+        slice_len=128
+    ):
+        # import pdb; pdb.set_trace()
+        print("answer >>", tokenizer.decode(token_ids[-1], skip_special_tokens=True), end='', flush=True)
+        self.max_seq_len = 2559
+        mask = np.zeros((1, 1, self.max_seq_len + 1), dtype=np.float32).astype(bfloat16)
+        mask[:, :, :self.max_seq_len] -= 65536
+        seq_len = len(token_ids) - 1
+        if prefill_len > 0:
+            mask[:, :, :seq_len] = 0
+        for step_idx in range(self.max_seq_len):
+            if prefill_len > 0 and step_idx < seq_len:
+                continue
+            # import pdb; pdb.set_trace()
+            cur_token = token_ids[step_idx]
+            indices = np.array([step_idx], np.uint32).reshape((1, 1))
+            data = embed_matrix[cur_token, :].reshape((1, 1, self.config.hidden_size)).astype(bfloat16)
+            for layer_idx in range(self.config.num_hidden_layers):
+                input_feed = {
+                    "K_cache": self.k_caches[layer_idx],
+                    "V_cache": self.v_caches[layer_idx],
+                    "indices": indices,
+                    "input": data,
+                    "mask": mask,
+                }
+                outputs = self.decoder_sessions[layer_idx].run(None, input_feed, shape_group=0)
+                self.k_caches[layer_idx][:, step_idx, :] = outputs[0][:, :, :]
+                self.v_caches[layer_idx][:, step_idx, :] = outputs[1][:, :, :]
+                data = outputs[2]
+            mask[..., step_idx] = 0
+            if step_idx < seq_len - 1:
+                continue
+            else:
+                post_out = self.post_process_session.run(None, {"input": data})[0]
+                next_token, possible_tokens, possible_probs = self.post_process(post_out)
+                token_ids.append(next_token)
+                if next_token == tokenizer.eos_token_id and next_token > seq_len:
+                    break
+            print(tokenizer.decode(next_token, skip_special_tokens=True), end='', flush=True)
+