Super-squash branch 'main' using huggingface_hub

Co-authored-by: MaziyarPanahi <MaziyarPanahi@users.noreply.huggingface.co>
Co-authored-by: Crystalcareai <Crystalcareai@users.noreply.huggingface.co>

.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

README.md

@@ -0,0 +1,311 @@
+---
+license: apache-2.0
+language:
+- en
+- es
+- fr
+- de
+- it
+- pt
+- ru
+- ar
+- hi
+- ko
+- zh
+library_name: transformers
+base_model:
+- arcee-ai/Trinity-Large-Base
+arxiv:
+- 2602.17004
+tags:
+- reasoning
+- agentic
+- tool-calling
+- thinking
+---
+<!-- markdownlint-disable first-line-h1 -->
+<!-- markdownlint-disable html -->
+<!-- markdownlint-disable no-duplicate-header -->
+
+<div align="center">
+  <picture>
+    <img
+      src="https://cdn-uploads.huggingface.co/production/uploads/6435718aaaef013d1aec3b8b/i-v1KyAMOW_mgVGeic9WJ.png"
+      alt="Arcee Trinity Large Thinking"
+      style="max-width: 100%; height: auto;"
+    >
+  </picture>
+</div>
+<hr>
+
+# Trinity-Large-Thinking
+
+## Introduction
+
+Trinity-Large-Thinking is a reasoning-optimized variant of Arcee AI's Trinity-Large family — a 398B-parameter sparse Mixture-of-Experts (MoE) model with approximately 13B active parameters per token. Built on Trinity-Large-Base and post-trained with extended chain-of-thought reasoning and agentic RL, Trinity-Large-Thinking delivers state-of-the-art performance on agentic benchmarks while maintaining strong general capabilities.
+
+Trinity-Large-Thinking generates explicit reasoning traces wrapped in `<think>...</think>` blocks before producing its final response. This thinking process is critical to the model's performance — **thinking tokens must be kept in context** for multi-turn conversations and agentic loops to function correctly.
+
+Try it at [chat.arcee.ai](http://chat.arcee.ai/)
+
+More details on the training of Trinity Large are available in the [technical report](https://arxiv.org/abs/2602.17004).
+
+## Key Highlights
+
+- **Agentic-first design**: Purpose-built for tool calling, multi-step planning, and agent workflows
+- **State-of-the-art agentic performance**: 94.7% on τ²-Bench, 91.9% on PinchBench, 98.2% on LiveCodeBench
+- **Native reasoning traces**: Extended chain-of-thought via `<think>...</think>` blocks
+- **Compatible with major agent frameworks**: Works out of the box with [OpenClaw](https://github.com/openclaw) and [Hermes Agent](https://github.com/NousResearch/hermes-agent)
+- **Ready to use on [OpenRouter](https://openrouter.ai/)**: No setup required — full reasoning and tool calling support via API
+
+## Model Variants
+
+The Trinity Large family consists of four checkpoints:
+
+- **Trinity-Large-Thinking** (this release): Reasoning-optimized, agentic post-training with extended chain-of-thought
+- **[Trinity-Large-Preview](https://huggingface.co/arcee-ai/Trinity-Large-Preview)**: Lightly post-trained, chat-ready instruct model (no reasoning_content).
+- **[Trinity-Large-TrueBase](https://huggingface.co/arcee-ai/Trinity-Large-TrueBase)**: 10T-token pre-anneal pretraining checkpoint
+- **[Trinity-Large-Base](https://huggingface.co/arcee-ai/Trinity-Large-Base)**: Full 17T-token pretrained foundation model with mid-training anneals
+
+## Architecture
+
+Trinity-Large-Thinking shares the same sparse MoE architecture as Trinity-Large-Preview.
+
+| Hyperparameter | Value |
+|:---|:---:|
+| Total parameters | ~398B |
+| Active parameters per token | ~13B |
+| Experts | 256 (1 shared) |
+| Active experts | 4 |
+| Routing strategy | 4-of-256 (1.56% sparsity) |
+| Dense layers | 6 |
+| Pretraining context length | 8,192 |
+| Context length after extension | 512k |
+| Architecture | Sparse MoE (AfmoeForCausalLM) |
+
+## Benchmarks
+
+| Benchmark | Trinity-Large-Thinking | Opus-4.6 | GLM-5 | MiniMax-M2.7 | Kimi-K2.5 |
+|---|---:|---:|---:|---:|---:|
+| IFBench | 52.3 | 53.1 | 72.3 | **75.7** | 70.2 |
+| GPQA-Diamond | 76.3 | **89.2** | 81.6 | 86.2 | 86.9 |
+| Tau2-Airline | **88.0** | 82.0 | 80.5 | 80.0 | 80.0 |
+| Tau2-Telecom | 94.7 | 92.1 | **98.2** | 84.8 | 95.9 |
+| PinchBench | 91.9 | **93.3** | 86.4 | 89.8 | 84.8 |
+| AIME25 | 96.3 | **99.8** | 93.3 | 80.0 | 96.3 |
+| BCFLv4 | 70.1 | **77.0** | 70.8 | 70.6 | 68.3 |
+| MMLU-Pro | 83.4 | **89.1** | 85.8 | 80.8 | 87.1 |
+| SWE-bench Verified* | 63.2 | **75.6** | 72.8 | 75.4 | 70.8 |
+
+ *All models evaluated in mini-swe-agent-v2
+
+## Thinking-in-Context: Important Usage Note
+
+Trinity-Large-Thinking produces reasoning traces inside `<think>...</think>` blocks before generating its final response.
+
+This means:
+
+1. **Multi-turn conversations**: When building chat applications, include the full assistant response (thinking + answer) in the conversation history for subsequent turns.
+2. **Agentic loops**: When using Trinity-Large-Thinking as the backbone of an agent (OpenClaw, Hermes Agent, or custom), ensure your tool-calling loop preserves `<think>` blocks in the message history between steps.
+3. **Context window management**: The 512k extended context window accommodates long reasoning chains across many agentic steps. If you must truncate history, prefer removing older turns entirely rather than stripping thinking tokens from recent turns.
+
+### How thinking works
+
+The model reasons internally before producing its response. When served via vLLM, the reasoning is separated into a dedicated `reasoning_content` field in the API response:
+
+    // API response structure
+    {
+      "message": {
+        "role": "assistant",
+        "reasoning_content": "The user wants flight information. I need to determine the date for next Tuesday, search for flights SFO → JFK, and filter by price < $300.",
+        "content": "\n",
+        "tool_calls": [{
+          "function": {
+            "name": "search_flights",
+            "arguments": "{\"origin\": \"SFO\", \"destination\": \"JFK\", \"date\": \"2026-04-07\", \"max_price\": 300}"
+          }
+        }]
+      }
+    }
+
+When building multi-turn agentic loops, include the `reasoning_content` back in the conversation history (re-wrapped in `<think>...</think>` tags within the assistant message) so the model retains its prior reasoning chain.
+
+## Training Configuration
+
+### Pretraining
+
+- Training tokens: 17 trillion
+- Data partner: [Datology](https://www.datologyai.com/)
+
+### Posttraining
+
+- Instruction tuning and agentic RL with extended chain-of-thought
+- Trained on tool-calling trajectories, multi-step agent tasks, and reasoning chains
+
+### Infrastructure
+
+- Hardware: 2,048 NVIDIA B300 GPUs
+- Parallelism: HSDP + Expert Parallelism
+- Compute partner: [Prime Intellect](https://www.primeintellect.ai/)
+
+## Usage
+
+### Running our model
+
+- [vLLM](#vllm) (recommended for agentic deployments)
+- [Transformers](#transformers)
+- [API](#api)
+
+### vLLM
+
+Supported in vLLM 0.11.1+. For agentic use with both reasoning and tool calling:
+
+    vllm serve arcee-ai/Trinity-Large-Thinking \
+      --dtype bfloat16 \
+      --enable-reasoning \
+      --reasoning-parser deepseek_r1 \
+      --enable-auto-tool-choice \
+      --tool-call-parser qwen3_coder
+
+This configuration:
+- `--reasoning-parser deepseek_r1` — Parses `<think>...</think>` reasoning blocks and exposes them via the `reasoning_content` field in the API response
+- `--tool-call-parser qwen3_coder` — Parses structured tool calls from the model output into the OpenAI-compatible `tool_calls` array
+
+**Extracting reasoning content from the API response:**
+
+```python
+from openai import OpenAI
+
+client = OpenAI(api_key="EMPTY", base_url="http://localhost:8000/v1")
+
+response = client.chat.completions.create(
+    model="arcee-ai/Trinity-Large-Thinking",
+    messages=[
+        {"role": "user", "content": "What's the weather like in Paris?"}
+    ],
+    tools=[ # your tool definitions here
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "description": "Get current weather for a location",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {"type": "string"}
+                    },
+                    "required": ["location"]
+                }
+            }
+        }
+    ],
+)
+
+# Access reasoning (thinking) content
+reasoning = response.choices[0].message.reasoning_content
+
+# Access final response or tool calls
+content = response.choices[0].message.content
+tool_calls = response.choices[0].message.tool_calls
+```
+
+**Note on thinking-in-context with vLLM**: When building multi-turn agentic loops, include both `reasoning_content` and `content` in the conversation history you send back to the model. The reasoning content should be re-wrapped in `<think>...</think>` tags within the assistant message.
+
+### Transformers
+
+Use the `main` transformers branch or pass `trust_remote_code=True` with a released version.
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+model_id = "arcee-ai/Trinity-Large-Thinking"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+messages = [
+    {"role": "user", "content": "Who are you?"},
+]
+
+input_ids = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    return_tensors="pt"
+).to(model.device)
+
+outputs = model.generate(
+    input_ids,
+    max_new_tokens=4096,
+    do_sample=True,
+    temperature=0.6,
+    top_k=50,
+    top_p=0.95
+)
+
+response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+print(response)
+```
+
+### API
+
+Available on OpenRouter:
+
+    curl -X POST "https://openrouter.ai/v1/chat/completions" \
+      -H "Authorization: Bearer $OPENROUTER_API_KEY" \
+      -H "Content-Type: application/json" \
+      -d '{
+        "model": "arcee-ai/trinity-large-thinking",
+        "messages": [
+          {
+            "role": "user",
+            "content": "What are some fun things to do in New York?"
+          }
+        ]
+      }'
+
+## Agentic Use Cases
+
+Trinity-Large-Thinking is optimized for deployment as the reasoning backbone of AI agent systems. It has been evaluated and performs excellently with:
+
+### OpenClaw
+
+Trinity-Large-Thinking works as a drop-in brain for OpenClaw agents. Its native tool-calling format is compatible with OpenClaw's execution loop, and the extended reasoning enables reliable multi-step task completion — from email triage to code generation to meeting scheduling. Our 91.9% PinchBench score reflects real-world OpenClaw task performance.
+
+### Hermes Agent
+
+Compatible with the Hermes Agent framework from Nous Research. Trinity-Large-Thinking's reasoning traces pair naturally with Hermes's skill-learning loop — the model's explicit chain-of-thought makes skill extraction more reliable, and its strong tool-calling capabilities integrate directly via the Hermes tool-use protocol.
+
+### Custom Agent Loops
+
+For custom implementations, the key integration pattern is:
+
+1. Send the user message with tool definitions
+2. Receive the response with `<think>` reasoning + tool calls
+3. Execute the tool calls
+4. Append the **full** assistant response (thinking + content + tool calls) and tool results to the message history
+5. Send the updated history back for the next step
+6. Repeat until the model produces a final response without tool calls
+
+## License
+
+Trinity-Large-Thinking is released under the Apache License, Version 2.0.
+
+## Citation
+
+If you use this model, please cite:
+
+    @misc{singh2026arceetrinity,
+      title        = {Arcee Trinity Large Technical Report},
+      author       = {Varun Singh and Lucas Krauss and Sami Jaghouar and Matej Sirovatka and Charles Goddard and Fares Obied and Jack Min Ong and Jannik Straube and Fern and Aria Harley and Conner Stewart and Colin Kealty and Maziyar Panahi and Simon Kirsten and Anushka Deshpande and Anneketh Vij and Arthur Bresnu and Pranav Veldurthi and Raghav Ravishankar and Hardik Bishnoi and DatologyAI Team and Arcee AI Team and Prime Intellect Team and Mark McQuade and Johannes Hagemann and Lucas Atkins},
+      year         = {2026},
+      eprint       = {2602.17004},
+      archivePrefix= {arXiv},
+      primaryClass = {cs.LG},
+      doi          = {10.48550/arXiv.2602.17004},
+      url          = {https://arxiv.org/abs/2602.17004}
+    }

__init__.py

@@ -0,0 +1,4 @@
+from .configuration_afmoe import AfmoeConfig
+from .modeling_afmoe     import AfmoeForCausalLM
+
+__all__ = ["AfmoeConfig", "AfmoeForCausalLM"]

chat_template.jinja

@@ -0,0 +1,159 @@
+<|begin_of_text|>{%- macro render_extra_keys(json_dict, handled_keys) -%}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{%- endmacro -%}
+
+{%- macro render_tool_call(raw_tool_call) -%}
+    {%- if raw_tool_call.function is defined and raw_tool_call.function is mapping %}
+        {%- set tool_call = raw_tool_call.function %}
+    {%- else %}
+        {%- set tool_call = raw_tool_call %}
+    {%- endif %}
+    {{- '<tool_call>\n<function=' + (tool_call.name | default('') | string) + '>\n' }}
+    {%- if tool_call.arguments is defined and tool_call.arguments is mapping %}
+        {%- for args_name, args_value in tool_call.arguments.items() %}
+            {{- '<parameter=' + (args_name | string) + '>\n' }}
+            {%- if args_value is mapping or (args_value is sequence and args_value is not string) %}
+                {{- args_value | tojson | safe }}
+            {%- else %}
+                {{- args_value | string }}
+            {%- endif %}
+            {{- '\n</parameter>\n' }}
+        {%- endfor %}
+    {%- endif %}
+    {{- '</function>\n</tool_call>' }}
+{%- endmacro -%}
+
+{%- set system_message = none %}
+{%- if messages and messages[0]["role"] == "system" %}
+    {%- set system_message = messages[0]["content"] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{%- set has_tools = tools is iterable and tools is not string and tools | length > 0 %}
+
+{%- if system_message is not none or has_tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if system_message is not none %}
+        {{- system_message }}
+    {%- else %}
+        {{- "You are Trinity Large, a helpful assistant developed by Arcee AI, that can interact with a computer to solve tasks." }}
+    {%- endif %}
+    {%- if has_tools %}
+        {{- "\n\n# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+        {%- for tool in tools %}
+            {%- if tool.function is defined and tool.function is mapping %}
+                {%- set tool = tool.function %}
+            {%- endif %}
+            {{- '\n<function>\n<name>' ~ (tool.name | default('') | string) ~ '</name>' }}
+            {%- if tool.description is defined and tool.description is not none %}
+                {{- '\n<description>' ~ (tool.description | string | trim) ~ '</description>' }}
+            {%- endif %}
+            {{- '\n<parameters>' }}
+            {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+                {%- for param_name, param_fields in tool.parameters.properties.items() %}
+                    {{- '\n<parameter>\n<name>' ~ (param_name | string) ~ '</name>' }}
+                    {%- if param_fields is mapping and param_fields.type is defined and param_fields.type is not none %}
+                        {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                    {%- endif %}
+                    {%- if param_fields is mapping and param_fields.description is defined and param_fields.description is not none %}
+                        {{- '\n<description>' ~ (param_fields.description | string | trim) ~ '</description>' }}
+                    {%- endif %}
+                    {%- if param_fields is mapping %}
+                        {%- set handled_keys = ['name', 'type', 'description'] %}
+                        {{- render_extra_keys(param_fields, handled_keys) }}
+                    {%- endif %}
+                    {{- '\n</parameter>' }}
+                {%- endfor %}
+            {%- endif %}
+            {%- if tool.parameters is defined %}
+                {%- set handled_keys = ['type', 'properties'] %}
+                {{- render_extra_keys(tool.parameters, handled_keys) }}
+            {%- endif %}
+            {{- '\n</parameters>' }}
+            {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+            {{- render_extra_keys(tool, handled_keys) }}
+            {{- '\n</function>' }}
+        {%- endfor %}
+        {{- "\n</tools>" }}
+        {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- endif %}
+
+{%- for message in loop_messages %}
+    {%- set role = message.role | default('') %}
+    {%- if role == "assistant" %}
+        {%- set content_str = '' if message.content is none else (message.content | string) %}
+        {%- set trimmed_content = content_str | trim %}
+
+        {%- set has_reasoning_content = message.reasoning_content is defined %}
+        {%- set has_reasoning = has_reasoning_content or (message.reasoning is defined) %}
+
+        {%- if has_reasoning_content %}
+            {%- set reasoning_value = message.reasoning_content %}
+        {%- elif message.reasoning is defined %}
+            {%- set reasoning_value = message.reasoning %}
+        {%- else %}
+            {%- set reasoning_value = none %}
+        {%- endif %}
+
+        {%- set has_tool_calls = message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls is not string and message.tool_calls | length > 0 %}
+
+        {{- '<|im_start|>assistant\n' }}
+        {%- if has_reasoning %}
+            {%- if reasoning_value %}
+                {{- '<think>' + (reasoning_value | string | trim) + '</think>' }}
+            {%- else %}
+                {{- '<think></think>' }}
+            {%- endif %}
+            {%- if trimmed_content %}
+                {{- '\n' + trimmed_content }}
+            {%- endif %}
+        {%- elif has_tool_calls %}
+            {%- if trimmed_content %}
+                {{- trimmed_content }}
+            {%- endif %}
+        {%- else %}
+            {{- content_str }}
+        {%- endif %}
+
+        {%- if has_tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- set separator = '\n' if ((loop.first and (has_reasoning or trimmed_content)) or (not loop.first)) else '' -%}
+                {{- separator + render_tool_call(tool_call) }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif role == "tool" or role == "observation" or role == "function" %}
+        {%- if loop.first or loop.previtem.role not in ["tool", "observation", "function"] %}
+            {{- '<|im_start|>user\n' }}
+        {%- endif %}
+        {{- '<tool_response>\n' }}
+        {{- '' if message.content is none else (message.content | string) }}
+        {{- '\n</tool_response>\n' }}
+        {%- if loop.last or loop.nextitem.role not in ["tool", "observation", "function"] %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + (role | string) }}
+        {{- '\n' + ('' if message.content is none else (message.content | string)) }}
+        {{- '<|im_end|>\n' }}
+    {%- endif %}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n<think>' }}
+{%- endif %}

config.json

@@ -0,0 +1,108 @@
+{
+  "architectures": [
+    "AfmoeForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_afmoe.AfmoeConfig",
+    "AutoModel": "modeling_afmoe.AfmoeModel",
+    "AutoModelForCausalLM": "modeling_afmoe.AfmoeForCausalLM"
+  },
+  "dtype": "bfloat16",
+  "global_attn_every_n_layers": 4,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 3072,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "load_balance_coeff": 0.00005,
+  "max_position_embeddings": 262144,
+  "model_type": "afmoe",
+  "moe_intermediate_size": 3072,
+  "mup_enabled": true,
+  "n_group": 1,
+  "num_attention_heads": 48,
+  "num_dense_layers": 6,
+  "num_expert_groups": 1,
+  "num_experts": 256,
+  "num_experts_per_tok": 4,
+  "num_hidden_layers": 60,
+  "num_key_value_heads": 8,
+  "num_limited_groups": 1,
+  "num_shared_experts": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000,
+  "route_norm": true,
+  "route_scale": 2.448,
+  "score_func": "sigmoid",
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "topk_group": 1,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "use_grouped_mm": true,
+  "vocab_size": 200192
+}

configuration_afmoe.py

@@ -0,0 +1,133 @@
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.configuration_utils import layer_type_validation
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+class AfmoeConfig(PretrainedConfig):
+    """
+    n_group (`int`, *optional*, defaults to 1):
+            Number of groups for routed experts.
+    topk_group (`int`, *optional*, defaults to 1):
+        Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups).
+    """
+    model_type = "afmoe"
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        num_hidden_layers: int = 32,
+        vocab_size: int = 200192,
+        hidden_size: int = 2048,
+        intermediate_size: int = 6144,
+        moe_intermediate_size=1408,
+        num_dense_layers=1,
+        num_attention_heads=16,
+        num_key_value_heads=None,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=16384,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        num_experts=64,
+        num_experts_per_tok=6,
+        num_shared_experts=2,
+        num_expert_groups=1,
+        num_limited_groups=1,
+        score_func="sigmoid",
+        route_norm=True,
+        route_scale=1.0,
+        global_attn_every_n_layers=4,
+        sliding_window=1024,
+        mup_enabled=False,
+        layer_types=None,
+        attention_dropout: float = 0.0,
+        n_group: int = 1,
+        topk_group: int = 1,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_dense_layers = num_dense_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+
+        # MoE specific
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_expert_groups = num_expert_groups
+        self.num_limited_groups = num_limited_groups
+        self.score_func = score_func
+        self.route_norm = route_norm
+        self.route_scale = route_scale
+
+
+        # Attention specific
+        self.attention_dropout = attention_dropout
+        self.global_attn_every_n_layers = global_attn_every_n_layers
+        self.sliding_window = sliding_window
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            self.layer_types = [
+                "sliding_attention" if bool((i + 1) % global_attn_every_n_layers) else "full_attention" for i in range(self.num_hidden_layers)
+            ]
+        layer_type_validation(self.layer_types)
+
+        # muP specific
+        self.mup_enabled = mup_enabled
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+
+
+        # Validate rope configs
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["AfmoeConfig"]

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 3,
+  "pad_token_id": 12,
+  "transformers_version": "4.57.3",
+  "temperature": 0.8,
+  "top_p": 0.8
+}

model-00001-of-00031.safetensors

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

model-00002-of-00031.safetensors

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

model-00003-of-00031.safetensors

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

model-00004-of-00031.safetensors

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

model-00005-of-00031.safetensors

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

model-00006-of-00031.safetensors

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

model-00007-of-00031.safetensors

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

model-00008-of-00031.safetensors

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

model-00009-of-00031.safetensors

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

model-00010-of-00031.safetensors

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

model-00011-of-00031.safetensors

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

model-00012-of-00031.safetensors

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

model-00013-of-00031.safetensors

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

model-00014-of-00031.safetensors

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

model-00015-of-00031.safetensors

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

model-00016-of-00031.safetensors

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

model-00017-of-00031.safetensors

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

model-00018-of-00031.safetensors

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

model-00019-of-00031.safetensors

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

model-00020-of-00031.safetensors

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

model-00021-of-00031.safetensors

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

model-00022-of-00031.safetensors

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

model-00023-of-00031.safetensors

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

model-00024-of-00031.safetensors

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

model-00025-of-00031.safetensors

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

model-00026-of-00031.safetensors

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

model-00027-of-00031.safetensors

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

model-00028-of-00031.safetensors

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

model-00029-of-00031.safetensors

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

model-00030-of-00031.safetensors

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

model-00031-of-00031.safetensors

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

modeling_afmoe.py

@@ -0,0 +1,680 @@
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import (
+    MoeCausalLMOutputWithPast,
+    MoeModelOutputWithPast,
+)
+from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_sliding_window_causal_mask,
+)
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.integrations import use_kernel_forward_from_hub
+
+
+try:
+    from .configuration_afmoe import AfmoeConfig
+except:
+    from configuration_afmoe import AfmoeConfig
+
+class AfmoeRotaryEmbedding(nn.Module):
+
+    def __init__(self, config: AfmoeConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    def _dynamic_frequency_update(self, position_ids, device):
+        """
+        dynamic RoPE layers should recompute `inv_freq` in the following situations:
+        1 - growing beyond the cached sequence length (allow scaling)
+        2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
+        """
+        seq_len = torch.max(position_ids) + 1
+        if seq_len > self.max_seq_len_cached:  # growth
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)
+            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation
+            self.max_seq_len_cached = seq_len
+
+        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
+            # This .to() is needed if the model has been moved to a device after being initialized (because
+            # the buffer is automatically moved, but not the original copy)
+            self.original_inv_freq = self.original_inv_freq.to(device)
+            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+            self.max_seq_len_cached = self.original_max_seq_len
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        if "dynamic" in self.rope_type:
+            self._dynamic_frequency_update(position_ids, device=x.device)
+
+        # Core RoPE block
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
+        device_type = x.device.type
+        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+
+        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
+        cos = cos * self.attention_scaling
+        sin = sin * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+@use_kernel_forward_from_hub("RMSNorm")
+class AfmoeRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float):
+        """
+        AfmoeRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.dropout(
+        attn_weights, p=dropout, training=module.training
+    )
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class AfmoeMLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size or config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class AfmoeTokenChoiceRouter(nn.Module):
+    """Token-choice top-K router for MoE routing."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.score_func = config.score_func
+        self.route_norm = config.route_norm
+        self.route_scale = config.route_scale
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)     
+
+    def forward(self, hidden_states, expert_bias: torch.Tensor | None):
+        _, _, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        scores = self.gate(hidden_states)
+
+        # Apply scoring function in float32 for stability
+        if self.score_func == "sigmoid":
+            scores = torch.sigmoid(scores.to(torch.float32))
+        else:
+            scores = F.softmax(scores.to(torch.float32), dim=-1)
+
+        if expert_bias is not None:
+            _, selected_experts = torch.topk(scores + expert_bias, k=self.top_k, dim=1)
+            top_scores = scores.gather(dim=1, index=selected_experts)
+        else:
+            top_scores, selected_experts = torch.topk(scores, k=self.top_k, dim=1)
+
+        # Normalize weights if using sigmoid
+        if self.score_func == "sigmoid" and self.route_norm:
+            denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20
+            top_scores = top_scores / denominator
+
+        top_scores = top_scores * self.route_scale
+        return top_scores, selected_experts
+
+class AfmoeMoE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.router = AfmoeTokenChoiceRouter(config)
+
+        self.shared_experts = None
+        if config.num_shared_experts > 0:
+            self.shared_experts = AfmoeMLP(
+                config, config.moe_intermediate_size * config.num_shared_experts
+            )
+        self.experts = nn.ModuleList(
+            [AfmoeMLP(
+                config, intermediate_size=config.moe_intermediate_size
+            ) for _ in range(config.num_experts)]
+        )
+        self.expert_bias = nn.Parameter(torch.zeros(config.num_experts, dtype=torch.float32), requires_grad=False)
+        
+
+    def forward(self, hidden_states):
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states_flat = hidden_states.view(-1, hidden_dim)
+
+        # Get routing decisions
+        top_scores, selected_experts = self.router(hidden_states, self.expert_bias)
+
+        # Process through shared experts
+        if self.shared_experts is not None:
+            shared_output = self.shared_experts(hidden_states_flat)
+        else:
+            shared_output = torch.zeros_like(hidden_states_flat)
+
+        # Reorder tokens by expert for efficient processing
+        token_indices_sorted = torch.argsort(selected_experts.view(-1), stable=True)
+        top_scores_sorted = top_scores.view(-1)[token_indices_sorted]
+        token_to_expert = selected_experts.view(-1)[token_indices_sorted]
+        token_indices_sorted = token_indices_sorted // self.config.num_experts_per_tok
+
+        # Gather input tokens
+        token_indices_expanded = token_indices_sorted.unsqueeze(-1).expand(
+            -1, hidden_dim
+        )
+        routed_input = torch.gather(
+            hidden_states_flat, dim=0, index=token_indices_expanded
+        )
+
+        routed_output = torch.zeros_like(routed_input)
+        for expert_id in range(self.config.num_experts):
+            mask = token_to_expert == expert_id
+            if mask.any():
+                expert_input = routed_input[mask]
+                expert_out = self.experts[expert_id](expert_input)
+                routed_output[mask] = expert_out
+          
+        routed_output = (
+            routed_output.to(torch.float32) * top_scores_sorted.unsqueeze(-1)
+        ).to(hidden_states.dtype)
+
+        # Scatter back to original positions
+        output = shared_output.scatter_add(
+            dim=0, index=token_indices_expanded, src=routed_output
+        )
+
+        return output.view(batch_size, seq_len, hidden_dim)
+
+
+class AfmoeAttention(nn.Module):
+    """Multi-headed attention with local/global pattern and gating."""
+
+    def __init__(self, config: AfmoeConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_local_attention = config.layer_types[layer_idx] == "sliding_attention"
+        self.sliding_window = config.sliding_window if self.is_local_attention else None
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, self.num_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, config.hidden_size, bias=False
+        )
+
+        self.q_norm = AfmoeRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = AfmoeRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+        self.gate_proj = nn.Linear(
+            config.hidden_size, self.num_heads * self.head_dim, bias=False
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:   
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape)
+        key_states = self.k_proj(hidden_states).view(hidden_shape)
+        value_states = self.v_proj(hidden_states).view(hidden_shape)
+        gate_states = self.gate_proj(hidden_states)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+        
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        if self.is_local_attention:
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            cache_kwargs = {"cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+
+        output, _ = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask=attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        output = output.view(*input_shape, -1).contiguous()
+        output = output * F.sigmoid(gate_states)
+        return self.o_proj(output)
+
+
+class AfmoeDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: AfmoeConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.layer_idx = layer_idx
+
+        self.self_attn = AfmoeAttention(config=config, layer_idx=layer_idx)
+        self.attention_type = config.layer_types[layer_idx]
+
+        # Dual normalization for attention
+        self.input_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Dual normalization for FFN
+        self.pre_mlp_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_mlp_layernorm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # MoE or dense FFN
+        self.moe_enabled = layer_idx >= config.num_dense_layers
+        if self.moe_enabled:
+            self.mlp = AfmoeMoE(config)
+        else:
+            self.mlp = AfmoeMLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+
+        # Self Attention with dual normalization
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        # FFN with dual normalization
+        residual = hidden_states
+        hidden_states = self.pre_mlp_layernorm(hidden_states)
+
+        if self.moe_enabled:
+            hidden_states = self.mlp(hidden_states)
+        else:
+            hidden_states = self.mlp(hidden_states)
+
+        hidden_states = self.post_mlp_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class AfmoePreTrainedModel(PreTrainedModel):
+    config_class = AfmoeConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["AfmoeDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _keep_in_fp32_modules = [
+        "input_layernorm",
+        "post_attention_layernorm",
+        "pre_mlp_layernorm",
+        "post_mlp_layernorm",
+        "q_norm",
+        "k_norm",
+        "norm",
+    ]
+    _supports_sdpa = True
+    _supports_attention_backend = True
+    supports_gradient_checkpointing = True
+
+
+class AfmoeModel(AfmoePreTrainedModel):
+    _no_split_modules = ["AfmoeDecoderLayer"]
+
+    def __init__(self, config: AfmoeConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                AfmoeDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = AfmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = AfmoeRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+            }
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+                "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
+            }
+
+        hidden_states = inputs_embeds
+
+        # Apply muP input scaling if enabled
+        if self.config.mup_enabled:
+            hidden_states = hidden_states * (self.config.hidden_size**0.5)
+
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = AfmoeModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        token_type_ids: Optional[torch.Tensor] = None,  # will be ignored
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+        outputs: MoeModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+
+__all__ = [
+    "AfmoeForCausalLM",
+    "AfmoeModel",
+    "AfmoePreTrainedModel",
+]

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<|begin_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|pad|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,271 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|begin_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|end_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<name>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "5": {
+      "content": "</name>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "6": {
+      "content": "<description>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "7": {
+      "content": "</description>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "8": {
+      "content": "<parameters>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "9": {
+      "content": "</parameters>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "10": {
+      "content": "<type>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "11": {
+      "content": "</type>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "12": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "14": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "15": {
+      "content": "<tools>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "16": {
+      "content": "</tools>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "17": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "18": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "19": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "20": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "21": {
+      "content": "<properties>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "22": {
+      "content": "</properties>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "23": {
+      "content": "<required>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "24": {
+      "content": "</required>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "25": {
+      "content": "<parameter>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "26": {
+      "content": "</parameter>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "27": {
+      "content": "<function>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "28": {
+      "content": "</function>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "29": {
+      "content": "<function=",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "30": {
+      "content": "<parameter=",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "31": {
+      "content": "<|reserved_special_18|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|begin_of_text|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": {},
+  "model_max_length": 65536,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "use_default_system_prompt": false
+}