Initialize this repo

.gitattributes

@@ -0,0 +1,37 @@
+*.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

README.md

@@ -0,0 +1,90 @@
+---
+license: bsd-3-clause
+language:
+  - en
+  - zh
+base_model:
+  - HuggingFaceTB/Gemma-3-1B-it
+pipeline_tag: Text Generation
+tags:
+  - HuggingFaceTB
+  - Gemma-3-1B-it
+---
+
+# Gemma-3-1B-it-Int8
+
+This version of Gemma-3-1B-it has been converted to run on the Axera NPU using **w8a16** quantization.
+
+Compatible with Pulsar2 version: 4.0
+
+## Convert tools links:
+
+For those who are interested in model conversion, you can try to export axmodel through the original repo:
+- https://huggingface.co/HuggingFaceTB/Gemma-3-1B-it
+
+- [Github for gemma-3-1B-it.axera](https://github.com/AXERA-TECH/gemma-3-1B-it.axera)
+
+- [Pulsar2 Link, How to Convert LLM from Huggingface to axmodel](https://pulsar2-docs.readthedocs.io/en/latest/appendix/build_llm.html)
+
+## Support Platform
+- AX650
+  - [M4N-Dock(爱芯派Pro)](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html)
+
+<!-- ## TODO Model infer time -->
+
+## How to use
+
+Download all files from this repository to the device.
+
+**Using AX650 Board**
+
+```bash
+ai@ai-bj ~/yongqiang/Gemma-3-1B-it $ tree -L 1
+.
+├── config.json
+├── gemma3_axmodel
+├── gemma3_tokenizer
+├── infer_axmodel.py
+├── README.md
+└── utils
+
+3 directories, 3 files
+```
+
+#### Inference with AX650 Host, such as M4N-Dock(爱芯派Pro) or AX650N DEMO Board
+
+**Text Generation**
+
+input text:
+
+```sh
+$ python3 infer_axmodel.py -q "请用中文介绍一下你自己."
+```
+
+log information:
+
+```bash
+[INFO] Compiler version: 5.0-patch1-dirty 93949955-dirty
+Init InferenceSession: 100%|██████████████████████████████████████████████████████████| 26/26 [00:21<00:00,  1.18it/s]
+[INFO] Using provider: AxEngineExecutionProvider
+[INFO] Model type: 2 (triple core)
+[INFO] Compiler version: 5.0-patch1-dirty 93949955-dirty
+Model loaded successfully!
+slice_indices: [0]
+Slice prefill done: 0
+answer >> 您好！我是一个大型语言模型，由 Google 训练。
+
+简单来说，我可以帮你做很多事情，比如：
+
+*   **回答你的问题：** 无论你问什么，我都会尽力用清晰、准确的语言来回答。
+*   **生成文本：** 比如写诗歌、故事、邮件、代码等等。
+*   **翻译语言：** 我可以将一种语言翻译成另一种语言。
+*   **总结文本：** 我可以帮你快速阅读一段文字，提取关键信息。
+*   **进行创意写作：** 我们可以一起头脑风暴，一起创作故事或文章。
+
+我还在不断学习和进步，所以我的能力也在不断提升。
+
+我是一个工具，可以帮助你，但不能代替人类的思考和判断。
+
+希望我能帮到你^@！ 你有什么想问的或者想让我做什么吗？ 😊
+```

gemma3_axmodel/gemma3_text_p128_l0_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l10_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l11_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l12_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l13_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l14_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l15_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l16_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l17_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l18_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l19_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l1_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l20_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l21_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l22_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l23_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l24_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l25_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l2_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l3_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l4_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l5_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l6_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l7_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l8_together.axmodel

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

gemma3_axmodel/gemma3_text_p128_l9_together.axmodel

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

gemma3_axmodel/gemma3_text_post.axmodel

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

gemma3_axmodel/model.embed_tokens.weight.bfloat16.bin

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

gemma3_axmodel/model.embed_tokens.weight.npy

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

gemma3_tokenizer/.gitattributes

@@ -0,0 +1,38 @@
+*.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
+model.safetensors filter=lfs diff=lfs merge=lfs -text
+tokenizer.model filter=lfs diff=lfs merge=lfs -text

gemma3_tokenizer/README.md

@@ -0,0 +1,514 @@
+---
+license: gemma
+library_name: transformers
+pipeline_tag: text-generation
+extra_gated_heading: Access Gemma on Hugging Face
+extra_gated_prompt: To access Gemma on Hugging Face, you’re required to review and
+  agree to Google’s usage license. To do this, please ensure you’re logged in to Hugging
+  Face and click below. Requests are processed immediately.
+extra_gated_button_content: Acknowledge license
+base_model: google/gemma-3-1b-pt
+---
+
+# Gemma 3 model card
+
+**Model Page**: [Gemma](https://ai.google.dev/gemma/docs/core)
+
+**Resources and Technical Documentation**:
+
+* [Gemma 3 Technical Report][g3-tech-report]
+* [Responsible Generative AI Toolkit][rai-toolkit]
+* [Gemma on Kaggle][kaggle-gemma]
+* [Gemma on Vertex Model Garden][vertex-mg-gemma3]
+
+**Terms of Use**: [Terms][terms]
+
+**Authors**: Google DeepMind
+
+## Model Information
+
+Summary description and brief definition of inputs and outputs.
+
+### Description
+
+Gemma is a family of lightweight, state-of-the-art open models from Google,
+built from the same research and technology used to create the Gemini models.
+Gemma 3 models are multimodal, handling text and image input and generating text
+output, with open weights for both pre-trained variants and instruction-tuned
+variants. Gemma 3 has a large, 128K context window, multilingual support in over
+140 languages, and is available in more sizes than previous versions. Gemma 3
+models are well-suited for a variety of text generation and image understanding
+tasks, including question answering, summarization, and reasoning. Their
+relatively small size makes it possible to deploy them in environments with
+limited resources such as laptops, desktops or your own cloud infrastructure,
+democratizing access to state of the art AI models and helping foster innovation
+for everyone.
+
+### Inputs and outputs
+
+-   **Input:**
+    -  Text string, such as a question, a prompt, or a document to be summarized
+    -  Images, normalized to 896 x 896 resolution and encoded to 256 tokens
+       each
+    -  Total input context of 128K tokens for the 4B, 12B, and 27B sizes, and
+       32K tokens for the 1B size
+
+-   **Output:**
+    -   Generated text in response to the input, such as an answer to a
+        question, analysis of image content, or a summary of a document
+    -   Total output context of 8192 tokens
+
+### Usage
+
+Below, there are some code snippets on how to get quickly started with running the model. First, install the Transformers library. Gemma 3 is supported starting from transformers 4.50.0. 
+
+```sh
+$ pip install -U transformers
+```
+
+Then, copy the snippet from the section that is relevant for your use case.
+
+#### Running with the `pipeline` API
+
+With instruction-tuned models, you need to use chat templates to process our inputs first. Then, you can pass it to the pipeline.
+
+```python
+from transformers import pipeline
+import torch
+
+pipe = pipeline("text-generation", model="google/gemma-3-1b-it", device="cuda", torch_dtype=torch.bfloat16)
+
+messages = [
+    [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": "You are a helpful assistant."},]
+        },
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": "Write a poem on Hugging Face, the company"},]
+        },
+    ],
+]
+
+output = pipe(messages, max_new_tokens=50)
+```
+
+#### Running the model on a single / multi GPU
+
+```python
+from transformers import AutoTokenizer, BitsAndBytesConfig, Gemma3ForCausalLM
+import torch
+
+model_id = "google/gemma-3-1b-it"
+
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
+model = Gemma3ForCausalLM.from_pretrained(
+    model_id, quantization_config=quantization_config
+).eval()
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+messages = [
+    [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": "You are a helpful assistant."},]
+        },
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": "Write a poem on Hugging Face, the company"},]
+        },
+    ],
+]
+inputs = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    tokenize=True,
+    return_dict=True,
+    return_tensors="pt",
+).to(model.device).to(torch.bfloat16)
+
+
+with torch.inference_mode():
+    outputs = model.generate(**inputs, max_new_tokens=64)
+
+outputs = tokenizer.batch_decode(outputs)
+```
+
+
+### Citation
+
+```none
+@article{gemma_2025,
+    title={Gemma 3},
+    url={https://goo.gle/Gemma3Report},
+    publisher={Kaggle},
+    author={Gemma Team},
+    year={2025}
+}
+```
+
+## Model Data
+
+Data used for model training and how the data was processed.
+
+### Training Dataset
+
+These models were trained on a dataset of text data that includes a wide variety
+of sources. The 27B model was trained with 14 trillion tokens, the 12B model was
+trained with 12 trillion tokens, 4B model was trained with 4 trillion tokens and
+1B with 2 trillion tokens. Here are the key components:
+
+-   Web Documents: A diverse collection of web text ensures the model is
+    exposed to a broad range of linguistic styles, topics, and vocabulary. The
+    training dataset includes content in over 140 languages.
+-   Code: Exposing the model to code helps it to learn the syntax and
+    patterns of programming languages, which improves its ability to generate
+    code and understand code-related questions.
+-   Mathematics: Training on mathematical text helps the model learn logical
+    reasoning, symbolic representation, and to address mathematical queries.
+-   Images: A wide range of images enables the model to perform image
+    analysis and visual data extraction tasks.
+
+The combination of these diverse data sources is crucial for training a powerful
+multimodal model that can handle a wide variety of different tasks and data
+formats.
+
+### Data Preprocessing
+
+Here are the key data cleaning and filtering methods applied to the training
+data:
+
+-   CSAM Filtering: Rigorous CSAM (Child Sexual Abuse Material) filtering
+    was applied at multiple stages in the data preparation process to ensure
+    the exclusion of harmful and illegal content.
+-   Sensitive Data Filtering: As part of making Gemma pre-trained models
+    safe and reliable, automated techniques were used to filter out certain
+    personal information and other sensitive data from training sets.
+-   Additional methods: Filtering based on content quality and safety in
+    line with [our policies][safety-policies].
+
+## Implementation Information
+
+Details about the model internals.
+
+### Hardware
+
+Gemma was trained using [Tensor Processing Unit (TPU)][tpu] hardware (TPUv4p,
+TPUv5p and TPUv5e). Training vision-language models (VLMS) requires significant
+computational power. TPUs, designed specifically for matrix operations common in
+machine learning, offer several advantages in this domain:
+
+-   Performance: TPUs are specifically designed to handle the massive
+    computations involved in training VLMs. They can speed up training
+    considerably compared to CPUs.
+-   Memory: TPUs often come with large amounts of high-bandwidth memory,
+    allowing for the handling of large models and batch sizes during training.
+    This can lead to better model quality.
+-   Scalability: TPU Pods (large clusters of TPUs) provide a scalable
+    solution for handling the growing complexity of large foundation models.
+    You can distribute training across multiple TPU devices for faster and more
+    efficient processing.
+-   Cost-effectiveness: In many scenarios, TPUs can provide a more
+    cost-effective solution for training large models compared to CPU-based
+    infrastructure, especially when considering the time and resources saved
+    due to faster training.
+-   These advantages are aligned with
+    [Google's commitments to operate sustainably][sustainability].
+
+### Software
+
+Training was done using [JAX][jax] and [ML Pathways][ml-pathways].
+
+JAX allows researchers to take advantage of the latest generation of hardware,
+including TPUs, for faster and more efficient training of large models. ML
+Pathways is Google's latest effort to build artificially intelligent systems
+capable of generalizing across multiple tasks. This is specially suitable for
+foundation models, including large language models like these ones.
+
+Together, JAX and ML Pathways are used as described in the
+[paper about the Gemini family of models][gemini-2-paper]; *"the 'single
+controller' programming model of Jax and Pathways allows a single Python
+process to orchestrate the entire training run, dramatically simplifying the
+development workflow."*
+
+## Evaluation
+
+Model evaluation metrics and results.
+
+### Benchmark Results
+
+These models were evaluated against a large collection of different datasets and
+metrics to cover different aspects of text generation:
+
+#### Reasoning and factuality
+
+| Benchmark                      | Metric         | Gemma 3 PT 1B  | Gemma 3 PT 4B | Gemma 3 PT 12B | Gemma 3 PT 27B |
+| ------------------------------ |----------------|:--------------:|:-------------:|:--------------:|:--------------:|
+| [HellaSwag][hellaswag]         | 10-shot        |      62.3      |      77.2     |      84.2      |      85.6      |
+| [BoolQ][boolq]                 | 0-shot         |      63.2      |      72.3     |      78.8      |      82.4      |
+| [PIQA][piqa]                   | 0-shot         |      73.8      |      79.6     |      81.8      |      83.3      |
+| [SocialIQA][socialiqa]         | 0-shot         |      48.9      |      51.9     |      53.4      |      54.9      |
+| [TriviaQA][triviaqa]           | 5-shot         |      39.8      |      65.8     |      78.2      |      85.5      |
+| [Natural Questions][naturalq]  | 5-shot         |      9.48      |      20.0     |      31.4      |      36.1      |
+| [ARC-c][arc]                   | 25-shot        |      38.4      |      56.2     |      68.9      |      70.6      |
+| [ARC-e][arc]                   | 0-shot         |      73.0      |      82.4     |      88.3      |      89.0      |
+| [WinoGrande][winogrande]       | 5-shot         |      58.2      |      64.7     |      74.3      |      78.8      |
+| [BIG-Bench Hard][bbh]          | few-shot       |      28.4      |      50.9     |      72.6      |      77.7      |
+| [DROP][drop]                   | 1-shot         |      42.4      |      60.1     |      72.2      |      77.2      |
+
+[hellaswag]: https://arxiv.org/abs/1905.07830
+[boolq]: https://arxiv.org/abs/1905.10044
+[piqa]: https://arxiv.org/abs/1911.11641
+[socialiqa]: https://arxiv.org/abs/1904.09728
+[triviaqa]: https://arxiv.org/abs/1705.03551
+[naturalq]: https://github.com/google-research-datasets/natural-questions
+[arc]: https://arxiv.org/abs/1911.01547
+[winogrande]: https://arxiv.org/abs/1907.10641
+[bbh]: https://paperswithcode.com/dataset/bbh
+[drop]: https://arxiv.org/abs/1903.00161
+
+#### STEM and code
+
+| Benchmark                      | Metric         | Gemma 3 PT 4B | Gemma 3 PT 12B | Gemma 3 PT 27B |
+| ------------------------------ |----------------|:-------------:|:--------------:|:--------------:|
+| [MMLU][mmlu]                   | 5-shot         |      59.6     |      74.5      |      78.6      |
+| [MMLU][mmlu] (Pro COT)         | 5-shot         |      29.2     |      45.3      |      52.2      |
+| [AGIEval][agieval]             | 3-5-shot       |      42.1     |      57.4      |      66.2      |
+| [MATH][math]                   | 4-shot         |      24.2     |      43.3      |      50.0      |
+| [GSM8K][gsm8k]                 | 8-shot         |      38.4     |      71.0      |      82.6      |
+| [GPQA][gpqa]                   | 5-shot         |      15.0     |      25.4      |      24.3      |
+| [MBPP][mbpp]                   | 3-shot         |      46.0     |      60.4      |      65.6      |
+| [HumanEval][humaneval]         | 0-shot         |      36.0     |      45.7      |      48.8      |
+
+[mmlu]: https://arxiv.org/abs/2009.03300
+[agieval]: https://arxiv.org/abs/2304.06364
+[math]: https://arxiv.org/abs/2103.03874
+[gsm8k]: https://arxiv.org/abs/2110.14168
+[gpqa]: https://arxiv.org/abs/2311.12022
+[mbpp]: https://arxiv.org/abs/2108.07732
+[humaneval]: https://arxiv.org/abs/2107.03374
+
+#### Multilingual
+
+| Benchmark                            | Gemma 3 PT 1B | Gemma 3 PT 4B | Gemma 3 PT 12B | Gemma 3 PT 27B |
+| ------------------------------------ |:-------------:|:-------------:|:--------------:|:--------------:|
+| [MGSM][mgsm]                         |      2.04     |      34.7     |      64.3     |      74.3     |
+| [Global-MMLU-Lite][global-mmlu-lite] |      24.9     |      57.0     |      69.4     |      75.7     |
+| [WMT24++][wmt24pp] (ChrF)            |      36.7     |      48.4     |      53.9     |      55.7     |
+| [FloRes][flores]                     |      29.5     |      39.2     |      46.0     |      48.8     |
+| [XQuAD][xquad] (all)                 |      43.9     |      68.0     |      74.5     |      76.8     |
+| [ECLeKTic][eclektic]                 |      4.69     |      11.0     |      17.2     |      24.4     |
+| [IndicGenBench][indicgenbench]       |      41.4     |      57.2     |      61.7     |      63.4     |
+
+[mgsm]: https://arxiv.org/abs/2210.03057
+[flores]: https://arxiv.org/abs/2106.03193
+[xquad]: https://arxiv.org/abs/1910.11856v3
+[global-mmlu-lite]: https://huggingface.co/datasets/CohereForAI/Global-MMLU-Lite
+[wmt24pp]: https://arxiv.org/abs/2502.12404v1
+[eclektic]: https://arxiv.org/abs/2502.21228
+[indicgenbench]: https://arxiv.org/abs/2404.16816
+
+#### Multimodal
+
+| Benchmark                      | Gemma 3 PT 4B | Gemma 3 PT 12B | Gemma 3 PT 27B |
+| ------------------------------ |:-------------:|:--------------:|:--------------:|
+| [COCOcap][coco-cap]            |      102      |      111       |      116       |
+| [DocVQA][docvqa] (val)         |      72.8     |      82.3      |      85.6      |
+| [InfoVQA][info-vqa] (val)      |      44.1     |      54.8      |      59.4      |
+| [MMMU][mmmu] (pt)              |      39.2     |      50.3      |      56.1      |
+| [TextVQA][textvqa] (val)       |      58.9     |      66.5      |      68.6      |
+| [RealWorldQA][realworldqa]     |      45.5     |      52.2      |      53.9      |
+| [ReMI][remi]                   |      27.3     |      38.5      |      44.8      |
+| [AI2D][ai2d]                   |      63.2     |      75.2      |      79.0      |
+| [ChartQA][chartqa]             |      63.6     |      74.7      |      76.3      |
+| [VQAv2][vqav2]                 |      63.9     |      71.2      |      72.9      |
+| [BLINK][blinkvqa]              |      38.0     |      35.9      |      39.6      |
+| [OKVQA][okvqa]                 |      51.0     |      58.7      |      60.2      |
+| [TallyQA][tallyqa]             |      42.5     |      51.8      |      54.3      |
+| [SpatialSense VQA][ss-vqa]     |      50.9     |      60.0      |      59.4      |
+| [CountBenchQA][countbenchqa]   |      26.1     |      17.8      |      68.0      |
+
+[coco-cap]: https://cocodataset.org/#home
+[docvqa]: https://www.docvqa.org/
+[info-vqa]: https://arxiv.org/abs/2104.12756
+[mmmu]: https://arxiv.org/abs/2311.16502
+[textvqa]: https://textvqa.org/
+[realworldqa]: https://paperswithcode.com/dataset/realworldqa
+[remi]: https://arxiv.org/html/2406.09175v1
+[ai2d]: https://allenai.org/data/diagrams
+[chartqa]: https://arxiv.org/abs/2203.10244
+[vqav2]: https://visualqa.org/index.html
+[blinkvqa]: https://arxiv.org/abs/2404.12390
+[okvqa]: https://okvqa.allenai.org/
+[tallyqa]: https://arxiv.org/abs/1810.12440
+[ss-vqa]: https://arxiv.org/abs/1908.02660
+[countbenchqa]: https://github.com/google-research/big_vision/blob/main/big_vision/datasets/countbenchqa/
+
+## Ethics and Safety
+
+Ethics and safety evaluation approach and results.
+
+### Evaluation Approach
+
+Our evaluation methods include structured evaluations and internal red-teaming
+testing of relevant content policies. Red-teaming was conducted by a number of
+different teams, each with different goals and human evaluation metrics. These
+models were evaluated against a number of different categories relevant to
+ethics and safety, including:
+
+-   **Child Safety**: Evaluation of text-to-text and image to text prompts
+    covering child safety policies, including child sexual abuse and
+    exploitation.
+-   **Content Safety:** Evaluation of text-to-text and image to text prompts
+    covering safety policies including, harassment, violence and gore, and hate
+    speech.
+-   **Representational Harms**: Evaluation of text-to-text and image to text
+    prompts covering safety policies including bias, stereotyping, and harmful
+    associations or inaccuracies.
+
+In addition to development level evaluations, we conduct "assurance
+evaluations" which are our 'arms-length' internal evaluations for responsibility
+governance decision making. They are conducted separately from the model
+development team, to inform decision making about release. High level findings
+are fed back to the model team, but prompt sets are held-out to prevent
+overfitting and preserve the results' ability to inform decision making.
+Assurance evaluation results are reported to our Responsibility & Safety Council
+as part of release review.
+
+### Evaluation Results
+
+For all areas of safety testing, we saw major improvements in the categories of
+child safety, content safety, and representational harms relative to previous
+Gemma models. All testing was conducted without safety filters to evaluate the
+model capabilities and behaviors. For both text-to-text and image-to-text, and
+across all model sizes, the model produced minimal policy violations, and showed
+significant improvements over previous Gemma models' performance with respect
+to ungrounded inferences. A limitation of our evaluations was they included only
+English language prompts.
+
+## Usage and Limitations
+
+These models have certain limitations that users should be aware of.
+
+### Intended Usage
+
+Open vision-language models (VLMs) models have a wide range of applications
+across various industries and domains. The following list of potential uses is
+not comprehensive. The purpose of this list is to provide contextual information
+about the possible use-cases that the model creators considered as part of model
+training and development.
+
+-   Content Creation and Communication
+    -   Text Generation: These models can be used to generate creative text
+        formats such as poems, scripts, code, marketing copy, and email drafts.
+    -   Chatbots and Conversational AI: Power conversational interfaces
+        for customer service, virtual assistants, or interactive applications.
+    -   Text Summarization: Generate concise summaries of a text corpus,
+        research papers, or reports.
+    -   Image Data Extraction: These models can be used to extract,
+        interpret, and summarize visual data for text communications.
+-   Research and Education
+    -   Natural Language Processing (NLP) and VLM Research: These
+        models can serve as a foundation for researchers to experiment with VLM
+        and NLP techniques, develop algorithms, and contribute to the
+        advancement of the field.
+    -   Language Learning Tools: Support interactive language learning
+        experiences, aiding in grammar correction or providing writing practice.
+    -   Knowledge Exploration: Assist researchers in exploring large
+        bodies of text by generating summaries or answering questions about
+        specific topics.
+
+### Limitations
+
+-   Training Data
+    -   The quality and diversity of the training data significantly
+        influence the model's capabilities. Biases or gaps in the training data
+        can lead to limitations in the model's responses.
+    -   The scope of the training dataset determines the subject areas
+        the model can handle effectively.
+-   Context and Task Complexity
+    -   Models are better at tasks that can be framed with clear
+        prompts and instructions. Open-ended or highly complex tasks might be
+        challenging.
+    -   A model's performance can be influenced by the amount of context
+        provided (longer context generally leads to better outputs, up to a
+        certain point).
+-   Language Ambiguity and Nuance
+    -   Natural language is inherently complex. Models might struggle
+        to grasp subtle nuances, sarcasm, or figurative language.
+-   Factual Accuracy
+    -   Models generate responses based on information they learned
+        from their training datasets, but they are not knowledge bases. They
+        may generate incorrect or outdated factual statements.
+-   Common Sense
+    -   Models rely on statistical patterns in language. They might
+        lack the ability to apply common sense reasoning in certain situations.
+
+### Ethical Considerations and Risks
+
+The development of vision-language models (VLMs) raises several ethical
+concerns. In creating an open model, we have carefully considered the following:
+
+-   Bias and Fairness
+    -   VLMs trained on large-scale, real-world text and image data can
+        reflect socio-cultural biases embedded in the training material. These
+        models underwent careful scrutiny, input data pre-processing described
+        and posterior evaluations reported in this card.
+-   Misinformation and Misuse
+    -   VLMs can be misused to generate text that is false, misleading,
+        or harmful.
+    -   Guidelines are provided for responsible use with the model, see the
+        [Responsible Generative AI Toolkit][rai-toolkit].
+-   Transparency and Accountability:
+    -   This model card summarizes details on the models' architecture,
+        capabilities, limitations, and evaluation processes.
+    -   A responsibly developed open model offers the opportunity to
+        share innovation by making VLM technology accessible to developers and
+        researchers across the AI ecosystem.
+
+Risks identified and mitigations:
+
+-   **Perpetuation of biases**: It's encouraged to perform continuous
+    monitoring (using evaluation metrics, human review) and the exploration of
+    de-biasing techniques during model training, fine-tuning, and other use
+    cases.
+-   **Generation of harmful content**: Mechanisms and guidelines for content
+    safety are essential. Developers are encouraged to exercise caution and
+    implement appropriate content safety safeguards based on their specific
+    product policies and application use cases.
+-   **Misuse for malicious purposes**: Technical limitations and developer
+    and end-user education can help mitigate against malicious applications of
+    VLMs. Educational resources and reporting mechanisms for users to flag
+    misuse are provided. Prohibited uses of Gemma models are outlined in the
+    [Gemma Prohibited Use Policy][prohibited-use].
+-   **Privacy violations**: Models were trained on data filtered for removal
+    of certain personal information and other sensitive data. Developers are
+    encouraged to adhere to privacy regulations with privacy-preserving
+    techniques.
+
+### Benefits
+
+At the time of release, this family of models provides high-performance open
+vision-language model implementations designed from the ground up for
+responsible AI development compared to similarly sized models.
+
+Using the benchmark evaluation metrics described in this document, these models
+have shown to provide superior performance to other, comparably-sized open model
+alternatives.
+
+[g3-tech-report]: https://goo.gle/Gemma3Report
+[rai-toolkit]: https://ai.google.dev/responsible
+[kaggle-gemma]: https://www.kaggle.com/models/google/gemma-3
+[vertex-mg-gemma3]: https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/gemma3
+[terms]: https://ai.google.dev/gemma/terms
+[safety-policies]: https://ai.google/static/documents/ai-responsibility-update-published-february-2025.pdf
+[prohibited-use]: https://ai.google.dev/gemma/prohibited_use_policy
+[tpu]: https://cloud.google.com/tpu/docs/intro-to-tpu
+[sustainability]: https://sustainability.google/operating-sustainably/
+[jax]: https://github.com/jax-ml/jax
+[ml-pathways]: https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/
+[sustainability]: https://sustainability.google/operating-sustainably/
+[gemini-2-paper]: https://arxiv.org/abs/2312.11805

gemma3_tokenizer/added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "<image_soft_token>": 262144
+}

gemma3_tokenizer/config.json

@@ -0,0 +1,37 @@
+{
+  "architectures": [
+    "Gemma3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_logit_softcapping": null,
+  "bos_token_id": 2,
+  "cache_implementation": "hybrid",
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "final_logit_softcapping": null,
+  "head_dim": 256,
+  "hidden_activation": "gelu_pytorch_tanh",
+  "hidden_size": 1152,
+  "initializer_range": 0.02,
+  "intermediate_size": 6912,
+  "max_position_embeddings": 32768,
+  "model_type": "gemma3_text",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 26,
+  "num_key_value_heads": 1,
+  "pad_token_id": 0,
+  "query_pre_attn_scalar": 256,
+  "rms_norm_eps": 1e-06,
+  "rope_local_base_freq": 10000,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": 512,
+  "sliding_window_pattern": 6,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.50.0.dev0",
+  "use_cache": true,
+  "vocab_size": 262144
+}

gemma3_tokenizer/generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 2,
+  "cache_implementation": "hybrid",
+  "do_sample": true,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "pad_token_id": 0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "4.50.0.dev0"
+}

gemma3_tokenizer/model.safetensors

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

gemma3_tokenizer/special_tokens_map.json

@@ -0,0 +1,33 @@
+{
+  "boi_token": "<start_of_image>",
+  "bos_token": {
+    "content": "<bos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eoi_token": "<end_of_image>",
+  "eos_token": {
+    "content": "<eos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "image_token": "<image_soft_token>",
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

gemma3_tokenizer/tokenizer.json

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

gemma3_tokenizer/tokenizer.model

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

infer_axmodel.py

@@ -0,0 +1,78 @@
+import torch
+from transformers import AutoConfig, AutoTokenizer
+import numpy as np
+import os
+from typing import List, Tuple
+from axengine import InferenceSession
+from ml_dtypes import bfloat16
+from utils.infer_func import InferManager
+import argparse
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Model configuration parameters")
+    parser.add_argument("--hf_model", type=str, default="./gemma3_tokenizer/",
+                        help="Path to HuggingFace model")
+    parser.add_argument("--axmodel_path", type=str, default="./gemma3_axmodel/",
+                        help="Path to save compiled axmodel of llama model")
+    parser.add_argument("-q", "--question", type=str, default="请用中文介绍一下你自己.",
+                        help="Your question that you want to ask the model.")
+    args = parser.parse_args()
+
+    hf_model_path = args.hf_model
+    axmodel_path = args.axmodel_path
+    prompt = args.question
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    embeds = np.load(os.path.join(axmodel_path, "model.embed_tokens.weight.npy"))
+
+    # load the tokenizer and the model
+    tokenizer = AutoTokenizer.from_pretrained(hf_model_path)
+    cfg = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=True)
+
+    eos_token_id = None
+    if isinstance(cfg.eos_token_id, list) and len(cfg.eos_token_id) > 1:
+        eos_token_id = cfg.eos_token_id 
+
+    ######################################################################
+    # Gemma-3
+    if "gemma3" in cfg.model_type.lower():
+        messages = [
+            [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": "You are a helpful assistant."},]
+                },
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": prompt},]
+                },
+            ],
+        ]
+
+        model_inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        )
+        # model_inputs = tokenizer([text], return_tensors="pt").to(device)
+        input_ids = model_inputs.input_ids
+
+    # Gemma-2
+    if "gemma2" in cfg.model_type.lower():
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    ######################################################################
+
+    token_ids = input_ids[0].cpu().numpy().tolist()
+    token_len = len(token_ids)
+    prefill_data = np.take(embeds, token_ids, axis=0)
+    prefill_data = prefill_data.astype(bfloat16)
+
+    imer = InferManager(cfg, axmodel_path)
+
+    token_ids = imer.prefill(tokenizer, token_ids, prefill_data, slice_len=128)
+    imer.decode(tokenizer, token_ids, embeds, slice_len=128, eos_token_id=eos_token_id)
+    print("\n")

utils/infer_func.py

@@ -0,0 +1,271 @@
+import torch
+import numpy as np
+from typing import List, Tuple
+from tqdm import tqdm
+from axengine import InferenceSession
+import os
+import re
+from ml_dtypes import bfloat16
+
+
+# Discover model files automatically from model_dir.
+# We expect files like: <prefix>_p128_l<idx>_together.axmodel and <prefix>_post.axmodel
+# we try to detect model prefix and layer files automatically
+def _find_axmodel_files(base_dir: str, expected_layers: int = None):
+    files = os.listdir(base_dir)
+    layer_pattern = re.compile(r"^(?P<prefix>.*)_p128_l(?P<idx>\d+)_together\.axmodel$")
+    post_pattern = re.compile(r"^(?P<prefix>.*)_post\.axmodel$")
+
+    # collect prefix -> [(idx, fname)]
+    prefix_map = {}
+    for fname in files:
+        m = layer_pattern.match(fname)
+        if m:
+            prefix = m.group("prefix")
+            idx = int(m.group("idx"))
+            prefix_map.setdefault(prefix, []).append((idx, fname))
+
+    if not prefix_map:
+        # fallback to hardcoded pattern if nothing detected
+        prefix = "gemma3_text"
+        layer_files = [(
+            i, f"{prefix}_p128_l{i}_together.axmodel"
+        ) for i in range(expected_layers or 0)]
+    else:
+        # choose the prefix with the most layers (most likely the correct one)
+        prefix = max(prefix_map.items(), key=lambda kv: len(kv[1]))[0]
+        # debug info
+        print(f"Detected prefixes: {list(prefix_map.keys())}, chosen: {prefix}, layers: {len(prefix_map[prefix])}")
+        layer_files = sorted(prefix_map[prefix], key=lambda it: it[0])
+
+    # find post process file
+    post_file = None
+    for fname in files:
+        m = post_pattern.match(fname)
+        if m and m.group("prefix") == prefix:
+            post_file = fname
+            break
+    if post_file is None:
+        candidate = os.path.join(base_dir, f"{prefix}_post.axmodel")
+        if os.path.exists(candidate):
+            post_file = f"{prefix}_post.axmodel"
+        else:
+            for fname in files:
+                if fname.endswith("_post.axmodel"):
+                    post_file = fname
+                    break
+
+    return layer_files, post_file, prefix
+
+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)
+        ]
+
+        layer_files, post_file, prefix = _find_axmodel_files(model_dir, config.num_hidden_layers)
+
+        self.decoder_sessions = []
+        for _, fname in tqdm(layer_files, desc="Init InferenceSession"):
+            session = InferenceSession(os.path.join(model_dir, fname))
+            self.decoder_sessions.append(session)
+
+        # post_file was returned by _find_axmodel_files; ensure it was found
+        if post_file is None:
+            raise FileNotFoundError("Cannot find post process .axmodel file in model_dir")
+        self.post_process_session = InferenceSession(os.path.join(model_dir, post_file))
+        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)
+
+        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,
+        eos_token_id=None, # 某些模型有多个 eos_token_id
+    ):
+        print("answer >>", tokenizer.decode(token_ids[-1], skip_special_tokens=True), end='', flush=True)
+        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
+            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)
+                if eos_token_id is not None and next_token in eos_token_id:
+                    break
+                elif next_token == tokenizer.eos_token_id:
+                    break
+                else:
+                    pass
+                token_ids.append(next_token)
+            print(tokenizer.decode(next_token, skip_special_tokens=True), end='', flush=True)
+