Updated cookbook

README.md

@@ -32,7 +32,7 @@ Laguna XS 2.1 is a 33B total parameter Mixture-of-Experts model with 3B activate
 - **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). High-quality FP8, NVFP4 and INT4 quantized variants available ([see the collection](https://huggingface.co/collections/poolside/laguna-xs-21))
+- **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). High-quality FP8, NVFP4 and INT4 quantized variants available ([see the collection](https://huggingface.co/collections/poolside/laguna-xs-21))
 - **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/))
 
 ---
@@ -118,8 +118,8 @@ pool acp setup --editor zed|jetbrains
 Use pool with Ollama with one-command setup:
 
 ```shell
-ollama pull laguna-xs-2.1
-ollama launch pool --model laguna-xs-2.1
+ollama pull laguna-xs.2.1
+ollama launch pool --model laguna-xs.2.1
 ```
 
 #### Feedback and issues
@@ -137,9 +137,6 @@ Serve Laguna XS 2.1 locally with vLLM and query it from any OpenAI-compatible cl
 > [!NOTE]
 > Laguna XS 2.1 support is available in vLLM 0.21.0 and later ([vllm-project/vllm#41129](https://github.com/vllm-project/vllm/pull/41129)).
 
-> [!IMPORTANT]
-> Use a vLLM build that includes [vllm-project/vllm#47311](https://github.com/vllm-project/vllm/pull/47311). Without that fix, the `poolside_v1` tool parser silently drops Laguna XS 2.1 tool calls that have no newline after the function name, and the raw tool-call markup leaks into the response. On an older build, pass `--tool-call-parser glm47` (the GLM 4.7 parser) instead.
-
 ```shell
 pip install 'vllm>=0.21.0'
 
@@ -268,12 +265,12 @@ The same recipe works for the [FP8](https://huggingface.co/poolside/Laguna-XS-2.
 > [!NOTE]
 > **Optional: speculative decoding with DFlash.** The [DFlash speculator](https://huggingface.co/poolside/Laguna-XS-2.1-DFlash) can be paired with Laguna XS 2.1 for lower latency. TRT-LLM support is in progress in [NVIDIA/TensorRT-LLM#15666](https://github.com/NVIDIA/TensorRT-LLM/pull/15666).
 
-#### llama.cpp
+#### Llama.cpp
 
 > [!NOTE]
 > Requires building llama.cpp from the upstream PR that adds Laguna XS 2.1 support until it lands ([ggml-org/llama.cpp#25165](https://github.com/ggml-org/llama.cpp/pull/25165)).
 
-Official GGUF conversions (BF16 and Q4\_K\_M) are available at [poolside/Laguna-XS-2.1-GGUF](https://huggingface.co/poolside/Laguna-XS-2.1-GGUF).
+GGUF weights (BF16 and Q4\_K\_M) are available at [poolside/Laguna-XS-2.1-GGUF](https://huggingface.co/poolside/Laguna-XS-2.1-GGUF).
 
 ```shell
 # Build llama.cpp from the PR branch
@@ -288,20 +285,19 @@ huggingface-cli download poolside/Laguna-XS-2.1-GGUF Laguna-XS-2.1-Q4_K_M.gguf -
 
 #### Ollama
 
-Available on the [Ollama library](https://ollama.com/library/laguna-xs-2.1).
+Available on the Ollama library:
 
-    ollama run laguna-xs-2.1          # default — Q4_K_M (imatrix)
-    ollama run laguna-xs-2.1:q8_0     # higher precision
-    ollama run laguna-xs-2.1:bf16     # full precision
+    ollama run laguna-xs.2.1          # default — Q4_K_M (imatrix)
+    ollama run laguna-xs.2.1:q8_0     # higher precision
+    ollama run laguna-xs.2.1:bf16     # full precision
 
 Reasoning and tool-calling work out of the box via the built-in `laguna` template.
 
-> [!NOTE]
-> **macOS (Metal) users:** `ollama run` may currently return empty output on Apple Silicon (a Metal f16 overflow in the MoE down-projection; fix pending in [ggml-org/llama.cpp#25389](https://github.com/ggml-org/llama.cpp/pull/25389)). Until it ships, use a Linux/CUDA host or the `/api/generate` endpoint with `"raw": true`.
-
-#### Atomic Chat
-
-Laguna XS 2.1 is also available in [Atomic Chat](https://atomic.chat/), a desktop app for running local models with a simple chat UI. To try it, download Atomic Chat, open the app, and choose Laguna XS 2.1 from the recommended models.
+> **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": {
@@ -226,4 +226,4 @@
     64,
     64
   ]
-}
+}

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

tokenizer_config.json

@@ -571,6 +571,5 @@
   "pad_token": "〈|PAD|〉",
   "sep_token": "〈|SEP|〉",
   "tokenizer_class": "PreTrainedTokenizerFast",
-  "unk_token": "〈|UNK|〉",
-  "chat_template": "{#- Iteration on laguna_glm_thinking_v8/chat_template.jinja -#}\n{#- No formatting instructions -#}\n{{- \"〈|EOS|〉\" -}}\n{%- set enable_thinking = enable_thinking | default(false) -%}\n{%- set add_generation_prompt = add_generation_prompt | default(false) -%}\n\n{#- ───── header (system message) ───── -#}\n{#- 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. -#}\n{%- 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.\" -%}\n{%- if messages and messages[0].role == \"system\" -%}\n  {%- set system_message = messages[0].content -%}\n  {%- set messages = messages[1:] -%}\n{%- endif -%}\n\n{%- set has_sys = system_message and system_message.strip() -%}\n{%- if has_sys or tools or enable_thinking -%}\n  {{- \"<system>\" -}}\n\n  {%- if has_sys -%}\n    {{- system_message.rstrip() -}}\n    {%- if tools -%}{{- \"\\n\\n\" -}}{%- endif -%}\n  {%- endif -%}\n\n  {%- if tools -%}\n    {{- \"### Tools\\n\\n\" -}}\n    {{- \"You may call functions to assist with the user query.\\n\" -}}\n    {{- \"All available function signatures are listed below:\\n\" -}}\n    {{- \"<available_tools>\\n\" -}}\n    {%- for tool in tools -%}\n      {{- (tool | tojson) ~ \"\\n\" -}}\n    {%- endfor -%}\n    {{- \"</available_tools>\" -}}\n  {%- endif -%}\n\n  {{- \"</system>\\n\" -}}\n{%- endif -%}\n\n{#- ───── main loop ───── -#}\n{%- for message in messages -%}\n  {%- set content = message.content if message.content is string else \"\" -%}\n  {%- if message.role == \"user\" -%}\n    {{- \"<user>\" + content + \"</user>\\n\" -}}\n  {%- elif message.role == \"assistant\" -%}\n    {%- generation -%}\n      {{- \"<assistant>\" -}}\n      {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content -#}\n      {%- set reasoning_content = '' -%}\n      {%- if message.reasoning is string -%}\n        {%- set reasoning_content = message.reasoning -%}\n      {%- elif message.reasoning_content is string -%}\n        {%- set reasoning_content = message.reasoning_content -%}\n      {%- endif -%}\n      {#- Display reasoning content for all messages if enable_thinking -#}\n      {%- if enable_thinking -%}\n        {{- '<think>' + reasoning_content + '</think>' -}}\n      {%- else -%}\n        {{- '</think>' -}}\n      {%- endif -%}\n      {#- Display main content (trailing newline only when no tool_calls follow) -#}\n      {%- if content -%}\n        {{- content -}}\n      {%- endif -%}\n      {%- if message.tool_calls -%}\n        {%- for tool_call in message.tool_calls -%}\n          {%- set function_data = tool_call.function -%}\n          {{- '<tool_call>' + function_data.name -}}\n          {%- set _args = function_data.arguments -%}\n          {%- for k, v in _args.items() -%}\n            {{- \"<arg_key>\" ~ k ~ \"</arg_key>\" -}}\n            {{- \"<arg_value>\" -}}{{- v | tojson(ensure_ascii=False) if v is not string else v -}}{{- \"</arg_value>\" -}}\n          {%- endfor -%}\n          {{- \"</tool_call>\" -}}\n        {%- endfor -%}\n      {%- endif -%}\n      {{- \"</assistant>\\n\" -}}\n    {%- endgeneration -%}\n  {%- elif message.role == \"tool\" -%}\n    {{- \"<tool_response>\" + content + \"</tool_response>\\n\" -}}\n  {%- elif message.role == \"system\" -%}\n    {#- Render additional system messages (the first one, if any, is handled separately in the header and was sliced off above) -#}\n    {{- \"<system>\" + content + \"</system>\\n\" -}}\n  {%- endif -%}\n{%- endfor -%}\n{#- ───── generation prompt ───── -#}\n{%- if add_generation_prompt -%}\n  {{- \"<assistant>\" -}}\n  {#- ───── Include reasoning mode directive ───── -#}\n  {%- if enable_thinking -%}\n    {{- '<think>' -}}\n  {%- else -%}\n    {{- '</think>' -}}\n  {%- endif -%}\n{%- endif -%}"
-}
+  "unk_token": "〈|UNK|〉"
+}