Update README.md

README.md

@@ -33,7 +33,7 @@ Laguna XS 2.1-FP8 is a 33B total parameter Mixture-of-Experts model with 3B acti
 - **Mixed SWA and global attention layout**: Laguna XS 2.1 uses sigmoid gating with per-layer rotary scales, enabling mixed SWA (Sliding Window Attention) and global attention layers in a 3:1 ratio (across 40 total layers)
 - **KV cache in FP8**: KV cache quantized to FP8, reducing memory per token
 - **Native reasoning support**: Interleaved thinking between tool calls with support for enabling and disabling thinking per-request
-- **Local-ready**: At 33B total parameters and 3B activated, Laguna XS 2.1 is compact enough to run on a Mac with 36 GB of RAM. Available on [Ollama](https://ollama.com/library/laguna-xs-2.1) and [llama.cpp](https://github.com/ggml-org/llama.cpp/pull/25165) (BF16 and Q4\_K\_M only)
+- **Local-ready**: At 33B total parameters and 3B activated, Laguna XS 2.1 is compact enough to run on a Mac with 36 GB of RAM. Available on [Ollama](https://ollama.com/library/laguna-xs.2.1) and [llama.cpp](https://github.com/ggml-org/llama.cpp/pull/25165) (BF16 and Q4\_K\_M only)
 - **OpenMDW-1.1 license**: Use and modify the model and associated materials freely for commercial and non-commercial purposes ([learn more about OpenMDW](https://openmdw.ai/))
 
 ---
@@ -107,9 +107,7 @@ Laguna XS 2.1-FP8 is supported in vLLM, SGLang and Transformers, and TRT-LLM tha
 
 #### vLLM
 
-The full vLLM recipe is on the main [Laguna XS 2.1 model card](https://huggingface.co/poolside/Laguna-XS-2.1) and on the [vLLM recipes page](https://recipes.vllm.ai/poolside/Laguna-XS-2.1). Quantization is detected automatically from `quantization_config` in this checkpoint, so the same command works with `poolside/Laguna-XS-2.1-FP8` substituted for the model ID. Set `VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER=0` when serving with vLLM.
-> [!IMPORTANT]
-> Tool calling requires a vLLM build with [vllm-project/vllm#47311](https://github.com/vllm-project/vllm/pull/47311); on older builds pass `--tool-call-parser glm47` (the GLM 4.7 parser) instead of `poolside_v1`.
+The full vLLM recipe is on the main [Laguna XS 2.1 model card](https://huggingface.co/poolside/Laguna-XS-2.1) and on the [vLLM recipes page](https://recipes.vllm.ai/poolside/Laguna-XS-2.1). Quantization is detected automatically from `quantization_config` in this checkpoint, so the same command works with `poolside/Laguna-XS-2.1-FP8` substituted for the model ID. No extra flags required.
 
 > [!NOTE]
 > The FP8-quantized KV cache requires vLLM >= 0.22.0. Earlier versions produce scrambled output on non-Hopper GPUs because of a per-layer attention-head count bug, fixed in [vllm#42650](https://github.com/vllm-project/vllm/pull/42650). On older vLLM, disable the FP8 KV cache by adding `--kv-cache-dtype-skip-layers $(seq 0 39)`.
@@ -134,10 +132,13 @@ llm = LLM(model="poolside/Laguna-XS-2.1-FP8", trust_remote_code=True)
 
 #### Ollama
 
-Available on the [Ollama library](https://ollama.com/library/laguna-xs-2.1).
+Visit [Ollama's model library](https://ollama.com/library/laguna-xs.2.1) to pull to your local machine.
 
-> [!NOTE]
-> **macOS (Metal) users:** Chat (`ollama run` / `/api/chat`) works as expected on Linux/CUDA. On macOS/Metal it may currently return empty output; the root cause is not yet fully understood and we're investigating it with the Ollama team. On a Mac, use a Linux/CUDA host, or the `/api/generate` endpoint with `"raw": true`.
+> **macOS (Metal) note:** Chat (`ollama run` / `/api/chat`) works on Linux/CUDA —
+> coherent output with reasoning and tool-calling intact. On macOS/Metal it may
+> currently return empty output; the root cause is not yet fully understood and
+> we're investigating it with the Ollama team. On a Mac, use a Linux/CUDA host, or
+> the `/api/generate` endpoint with `"raw": true`.
 
 ## Controlling reasoning
 

chat_template.jinja

@@ -1,93 +1,132 @@
-{#- Iteration on laguna_glm_thinking_v8/chat_template.jinja -#}
-{#- No formatting instructions -#}
+{#- Copied from laguna_glm_thinking_v4/chat_template.jinja -#}
+{#- Removes prefix that references <think> token, and replaces message.reasoning_content reference with message.reasoning -#}
 {{- "〈|EOS|〉" -}}
 {%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set render_assistant_messages_raw = render_assistant_messages_raw | default(false) -%}
 {%- set add_generation_prompt = add_generation_prompt | default(false) -%}
 
 {#- ───── header (system message) ───── -#}
-{#- A caller-supplied system message with empty content opts out of the default below, producing no <system> block — used to train without a system message. -#}
-{%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%}
+{%- set system_message = "" -%}
 {%- if messages and messages[0].role == "system" -%}
   {%- set system_message = messages[0].content -%}
-  {%- set messages = messages[1:] -%}
 {%- endif -%}
 
-{%- set has_sys = system_message and system_message.strip() -%}
-{%- if has_sys or tools or enable_thinking -%}
-  {{- "<system>" -}}
+{%- if (system_message and system_message.strip()) or tools -%}
+  {{- "<system>\n" -}}
 
-  {%- if has_sys -%}
+  {%- if system_message and system_message.strip() -%}
+    {{- "\n" -}}
     {{- system_message.rstrip() -}}
-    {%- if tools -%}{{- "\n\n" -}}{%- endif -%}
   {%- endif -%}
 
   {%- if tools -%}
-    {{- "### Tools\n\n" -}}
-    {{- "You may call functions to assist with the user query.\n" -}}
-    {{- "All available function signatures are listed below:\n" -}}
-    {{- "<available_tools>\n" -}}
+    {{- "\n\n### Tools\n\n" -}}
+    {%- set ns = namespace(tool_string="You may call functions to assist with the user query.\n"
+          ~ "All available function signatures are listed below:\n"
+          ~ "<available_tools>\n") -%}
     {%- for tool in tools -%}
-      {{- (tool | tojson) ~ "\n" -}}
+      {%- set ns.tool_string = ns.tool_string ~ (tool | tojson) ~ "\n" -%}
     {%- endfor -%}
-    {{- "</available_tools>" -}}
+    {%- if enable_thinking -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "Wrap your thinking in '<think>', '</think>' tags, followed by a function call. For each function call, return an unescaped XML-like object with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<think> your thoughts here </think>\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- else -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "For each function call, return an unescaped XML-like object " ~
+            "with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- endif -%}
+    {{- tool_string -}}
   {%- endif -%}
 
-  {{- "</system>\n" -}}
+  {{- "\n</system>\n" -}}
 {%- endif -%}
 
 {#- ───── main loop ───── -#}
 {%- for message in messages -%}
   {%- set content = message.content if message.content is string else "" -%}
   {%- if message.role == "user" -%}
-    {{- "<user>" + content + "</user>\n" -}}
+    {{- "<user>\n" + content + "\n</user>\n" -}}
   {%- elif message.role == "assistant" -%}
     {%- generation -%}
-      {{- "<assistant>" -}}
-      {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content -#}
-      {%- set reasoning_content = '' -%}
-      {%- if message.reasoning is string -%}
-        {%- set reasoning_content = message.reasoning -%}
-      {%- elif message.reasoning_content is string -%}
-        {%- set reasoning_content = message.reasoning_content -%}
-      {%- endif -%}
-      {#- Display reasoning content for all messages if enable_thinking -#}
-      {%- if enable_thinking -%}
-        {{- '<think>' + reasoning_content + '</think>' -}}
-      {%- else -%}
-        {{- '</think>' -}}
-      {%- endif -%}
-      {#- Display main content (trailing newline only when no tool_calls follow) -#}
-      {%- if content -%}
+      {{- "<assistant>\n" -}}
+      {%- if render_assistant_messages_raw -%}
+        {#- Raw mode: prepend the generation prompt token, then dump content verbatim. -#}
+        {#- The generation prompt is <think> when enable_thinking, </think> otherwise. -#}
+        {#- Only prepend if content doesn't already start with it. -#}
+        {%- if enable_thinking -%}
+          {%- if not content.startswith('<think>') -%}
+            {{- '<think>' -}}
+          {%- endif -%}
+        {%- else -%}
+          {%- if not content.startswith('</think>') -%}
+            {{- '</think>' -}}
+          {%- endif -%}
+        {%- endif -%}
         {{- content -}}
-      {%- endif -%}
-      {%- if message.tool_calls -%}
-        {%- for tool_call in message.tool_calls -%}
-          {%- set function_data = tool_call.function -%}
-          {{- '<tool_call>' + function_data.name -}}
-          {%- set _args = function_data.arguments -%}
-          {%- for k, v in _args.items() -%}
-            {{- "<arg_key>" ~ k ~ "</arg_key>" -}}
-            {{- "<arg_value>" -}}{{- v | tojson(ensure_ascii=False) if v is not string else v -}}{{- "</arg_value>" -}}
+        {#- Append closing tag if content doesn't already end with it. -#}
+        {%- if not content.endswith('</assistant>\n') and not content.endswith('</assistant>') -%}
+          {{- '\n</assistant>' -}}
+        {%- endif -%}
+        {{- "\n" -}}
+      {%- else -%}
+        {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content, or from <think> tags -#}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning is string %}
+          {%- set reasoning_content = message.reasoning %}
+        {%- elif message.reasoning_content is string %}
+          {%- set reasoning_content = message.reasoning_content %}
+        {%- endif %}
+        {#- Always strip <think> tags from content if present to avoid duplication -#}
+        {%- if '</think>' in content %}
+          {%- if not reasoning_content %}
+            {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+          {%- endif %}
+          {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+        {%- endif %}
+        {#- Display reasoning content for all messages -#}
+        {%- if reasoning_content -%}
+          {{- '<think>\n' + reasoning_content.strip() + '\n</think>\n' -}}
+        {%- else -%}
+          {{- '</think>\n' -}}
+        {%- endif -%}
+        {#- Display main content -#}
+        {%- if content.strip() -%}
+          {{- content.strip() ~ "\n" -}}
+        {%- endif -%}
+        {%- if message.tool_calls -%}
+          {%- for tool_call in message.tool_calls -%}
+            {%- set function_data = tool_call.function -%}
+            {{- '<tool_call>' + function_data.name }}
+            {% set _args = function_data.arguments %}
+            {%- for k, v in _args.items() -%}
+              {{- "<arg_key>" ~ k ~ "</arg_key>\n" -}}
+              {{- "<arg_value>"}}{{ v | tojson(ensure_ascii=False) if v is not string else v }}{{ "</arg_value>\n" -}}
+            {%- endfor -%}
+            {{- "</tool_call>\n" -}}
           {%- endfor -%}
-          {{- "</tool_call>" -}}
-        {%- endfor -%}
+        {%- endif -%}
+        {{- "</assistant>\n" -}}
       {%- endif -%}
-      {{- "</assistant>\n" -}}
     {%- endgeneration -%}
   {%- elif message.role == "tool" -%}
-    {{- "<tool_response>" + content + "</tool_response>\n" -}}
-  {%- elif message.role == "system" -%}
-    {#- Render additional system messages (the first one, if any, is handled separately in the header and was sliced off above) -#}
-    {{- "<system>" + content + "</system>\n" -}}
+    {{- "<tool_response>\n" + content + "\n</tool_response>\n" -}}
+  {%- elif message.role == "system" and loop.index0 != 0 -%}
+    {#- Render additional system messages (skip the first one which is handled separately in the header) -#}
+    {{- "<system>\n" + content + "\n</system>\n" -}}
   {%- endif -%}
 {%- endfor -%}
 {#- ───── generation prompt ───── -#}
 {%- if add_generation_prompt -%}
-  {{- "<assistant>" -}}
+  {{- "<assistant>\n" -}}
   {#- ───── Include reasoning mode directive ───── -#}
-  {%- if enable_thinking -%}
-    {{- '<think>' -}}
-  {%- else -%}
+  {%- if not enable_thinking %}
     {{- '</think>' -}}
-  {%- endif -%}
-{%- endif -%}
+  {%- else %}
+    {{- '<think>' -}}
+  {%- endif %}
+{%- endif -%}

config.json

@@ -47,7 +47,7 @@
       "original_max_position_embeddings": 8192,
       "beta_slow": 1.0,
       "beta_fast": 64.0,
-      "attention_factor": 1.3465735902799727,
+      "attention_factor": 1.0,
       "partial_rotary_factor": 0.5
     },
     "sliding_attention": {
@@ -335,4 +335,4 @@
     },
     "version": "0.14.1.dev11+gf2ee47b"
   }
-}
+}

configuration_laguna.py

@@ -186,12 +186,6 @@ class LagunaConfig(PreTrainedConfig):
         if rope_parameters is None:
             rope_parameters = {"rope_type": "default", "rope_theta": 500000.0}
 
-        # config.json stores SWA rope nested in rope_parameters["sliding_attention"]
-        # and carries no top-level swa_rope_parameters. Derive it here, else the
-        # sliding-window layers silently reuse the full-attention rope.
-        if swa_rope_parameters is None and isinstance(rope_parameters, dict):
-            swa_rope_parameters = rope_parameters.get("sliding_attention")
-
         # If ``partial_rotary_factor`` is set at the top level, inject it into any
         # rope dict that does not already carry one so the rotary embedding picks
         # it up consistently for both full-attention and SWA layers.

generation_config.json

@@ -10,7 +10,6 @@
  "temperature": 1.0,
  "top_p": 1.0,
  "min_p": 0.0,
- "top_k": 20,
  "speculative_config": {
   "method": "dflash",
   "source": "huggingface",

modeling_laguna.py

@@ -633,8 +633,8 @@ class LagunaModel(LagunaPreTrainedModel):
         use_cache: bool | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> MoeModelOutputWithPast:
-        from transformers.cache_utils import DynamicCache
-        from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+        from ...cache_utils import DynamicCache
+        from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
@@ -877,20 +877,3 @@ class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
 
 
 __all__ = ["LagunaForCausalLM", "LagunaModel", "LagunaPreTrainedModel"]
-
-
-# --- Added: register the native Laguna checkpoint-conversion for trust_remote_code loads.
-# transformers >=5.12 skips checkpoint-conversion mappings for custom (remote) code
-# unless explicitly registered, which broke loading the shipped per-expert MoE weights.
-try:
-    from transformers.conversion_mapping import (
-        get_checkpoint_conversion_mapping as _lg_get,
-        register_checkpoint_conversion_mapping as _lg_reg,
-        USER_REGISTERED_MAPPINGS as _lg_user,
-    )
-    if "laguna" not in _lg_user:
-        _lg_m = _lg_get("laguna")
-        if _lg_m is not None:
-            _lg_reg("laguna", _lg_m, overwrite=True)
-except Exception:
-    pass