Upload 14 files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* 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
+banner.png filter=lfs diff=lfs merge=lfs -text
+benchmarks.png filter=lfs diff=lfs merge=lfs -text
+score_progression.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,31 @@
+BananaMind Community License 1.0
+
+Copyright (c) 2026 Banaxi-Tech
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the “Software”), to deal
+in the Software without restriction except as stated in Section 1 below,
+including without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to permit
+persons to whom the Software is furnished to do so, subject to the following
+conditions:
+
+Section 1 — Commercial License Requirement.
+If the Software (or any derivative work thereof) is used in a commercial
+product or service that has more than 200,000 monthly active users, or more
+than 500,000 US dollars (or the equivalent in other currencies) in monthly
+revenue, you shall contact banaxitech@gmail.com and purchase a commercial
+license from Banaxi-Tech. The permissions granted above terminate with respect
+to that product or service if a commercial license is not obtained within 30
+days of first exceeding either threshold.
+
+Section 2 — Attribution.
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,270 @@
+---
+license: other
+license_name: bananamind-community-license-1.0
+license_link: LICENSE
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+base_model: BananaMind/BananaMind-2-Pro-Preview
+datasets:
+- HuggingFaceTB/smol-smoltalk
+tags:
+- causal-lm
+- chat
+- instruction-tuned
+- full-finetune
+- small-language-model
+- bananamind
+- bananamind2
+- bananamind2-pro
+- preview-checkpoint
+- smoltalk
+- digit-tokenizer
+- pytorch
+- safetensors
+- custom-code
+- trust-remote-code
+- custom-architecture
+---
+
+![BananaMind 2 Pro](banner.png)
+
+# BananaMind-2-Pro-Preview-Chat
+
+BananaMind-2-Pro-Preview-Chat is the instruction-tuned version of [BananaMind-2-Pro-Preview](https://huggingface.co/BananaMind/BananaMind-2-Pro-Preview). It was fully fine-tuned for one epoch on [HuggingFaceTB/smol-smoltalk](https://huggingface.co/datasets/HuggingFaceTB/smol-smoltalk), with loss applied only to assistant content and the assistant-ending EOS token.
+
+The model has **138,971,520 parameters**, a **3,072-token context window**, and a custom **32,768-token digit-aware byte-level BPE tokenizer**. It supports system prompts, multi-turn conversations, grouped-query attention, QK normalization, RoPE, tied embeddings, and KV-cached generation.
+
+This chat model is based on the 96K Pro Preview checkpoint, which was pretrained on **51,904,512,000 tokens**. It is a preview-derived instruction model rather than a fine-tune of the completed 100B-token BananaMind 2 Pro run.
+
+![BananaMind 2 Pro Preview Chat benchmark results](benchmarks.png)
+
+## Model Details
+
+| Field | Value |
+|---|---:|
+| Parameters | 138,971,520 |
+| Base model | `BananaMind/BananaMind-2-Pro-Preview` |
+| Architecture | BananaMind2Pro decoder-only Transformer |
+| Layers | 24 |
+| Hidden size | 640 |
+| Intermediate size | 1,920 |
+| Attention heads | 8 |
+| KV heads | 4 |
+| Head dimension | 80 |
+| Attention style | Grouped-query attention with QK norm |
+| MLP | SwiGLU |
+| Position embeddings | RoPE, theta 100,000 |
+| Normalization | RMSNorm, epsilon 1e-6 |
+| Vocabulary size | 32,768 |
+| Context length | 3,072 |
+| Embeddings | Tied input/output embeddings |
+| Generation cache | KV cache supported |
+| Weight format | safetensors |
+| Training type | Full-parameter supervised fine-tuning |
+
+## Benchmarks
+
+### BananaMind Instruct Bench 1.1
+
+Self-reported results from the official [BananaMind Instruct Bench 1.1](https://huggingface.co/datasets/BananaMind/BananaMind-Instruct-Bench-1.1) script. The benchmark contains 300 deterministic, difficulty- and category-weighted instruction tasks. The BananaMind models used their native chat templates; Supra 1.5 used the benchmark's Alpaca-style fallback. Evaluation used greedy decoding, `repetition_penalty=1.1`, and seed 42.
+
+| Model | Overall Elo | General | Multi-turn | System Prompts | Context Recall | Code |
+|---|---:|---:|---:|---:|---:|---:|
+| **BananaMind-2-Pro-Preview-Chat** | **888** | **720** | **911** | **793** | **1,194** | **1,337** |
+| [BananaMind-2-Medium-Chat](https://huggingface.co/BananaMind/BananaMind-2-Medium-Chat) | 787 | 534 | 872 | 666 | 1,085 | 1,291 |
+| [BananaMind-2-Mini-Chat](https://huggingface.co/BananaMind/BananaMind-2-Mini-Chat) | 654 | 389 | 812 | 743 | 924 | 667 |
+| [Supra 1.5 50M Instruct](https://huggingface.co/SupraLabs/Supra-1.5-50M-Instruct-exp) | 647 | 520 | 804 | 590 | 860 | 780 |
+| [BananaMind-2-Nano-Chat](https://huggingface.co/BananaMind/BananaMind-2-Nano-Chat) | 643 | 353 | 796 | 600 | 978 | 885 |
+
+**Detailed final result**
+
+| Scope | Examples | Passed | Weighted Score | Elo |
+|---|---:|---:|---:|---:|
+| **Overall** | **300** | **98** | **27.09%** | **888** |
+| General | 120 | 33 | 21.84% | 720 |
+| Multi-turn | 75 | 26 | 26.32% | 911 |
+| System Prompts | 60 | 9 | 11.84% | 793 |
+| Context Recall | 30 | 19 | 50.53% | 1,194 |
+| Code | 15 | 11 | 65.26% | 1,337 |
+
+| Difficulty | Examples | Passed | Weighted Score | Elo |
+|---|---:|---:|---:|---:|
+| Easy | 100 | 59 | 60.58% | 1,014 |
+| Medium | 100 | 31 | 32.59% | 900 |
+| Hard | 100 | 8 | 8.53% | 756 |
+
+The complete final run used CUDA, bfloat16, the native chat template, and no sampling. Scores can vary with benchmark revision, Transformers version, dtype, hardware, and generation settings.
+
+### Fine-tuning Progression
+
+![BananaMind 2 Pro Preview Chat score progression](score_progression.png)
+
+| Optimizer step | Packed tokens | Overall Elo | Passed | Weighted Score |
+|---:|---:|---:|---:|---:|
+| 500 | 49,152,000 | 811 | 77/300 | 19.98% |
+| 1,000 | 98,304,000 | 835 | 85/300 | 22.04% |
+| 1,500 | 147,456,000 | 858 | 91/300 | 24.19% |
+| 2,000 | 196,608,000 | 866 | 94/300 | 24.99% |
+| 2,500 | 245,760,000 | 874 | 94/300 | 25.69% |
+| 3,000 | 294,912,000 | 875 | 95/300 | 25.84% |
+| 3,500 | 344,064,000 | 855 | 88/300 | 23.88% |
+| **4,014 (final)** | **394,522,611** | **888** | **98/300** | **27.09%** |
+
+The final export is the strongest measured checkpoint overall. Step 3,000 had a slightly higher code-category Elo of 1,345, but the final model passed more code tasks and scored higher overall.
+
+## Instruction Tuning
+
+| Field | Value |
+|---|---:|
+| Dataset | `HuggingFaceTB/smol-smoltalk` |
+| Dataset split | `train` |
+| Dataset rows | 460,341 |
+| Epochs | 1 |
+| Final optimizer step | 4,014 |
+| Packed tokens processed | 394,522,611 |
+| Supervised assistant tokens | 301,938,934 |
+| Sequence length | 3,072 |
+| Micro batch | 4 |
+| Gradient accumulation | 8 |
+| Effective batch | 32 sequences |
+| Peak learning rate | 1e-4 |
+| Warmup | 100 steps |
+| LR schedule | Constant after warmup |
+| Optimizer | AdamW |
+| Betas | 0.9, 0.95 |
+| Weight decay | 0.1 |
+| Gradient clipping | 1.0 |
+| Compilation | PyTorch compile enabled |
+| Seed | 1337 |
+
+System and user messages were retained as context but masked from the loss. Only assistant content and the assistant-ending EOS token contributed to the objective. All model parameters were trainable; no adapters or LoRA modules were used.
+
+## Chat Template
+
+The tokenizer includes a Jinja template for `system`, `user`, and `assistant` messages:
+
+```text
+<BOS><|system|>
+{system message}
+<|user|>
+{user message}
+<|assistant|>
+{assistant response}<EOS>
+```
+
+The role markers are plain text encoded by the existing tokenizer. Fine-tuning did not add or resize any tokens, and the input and output embeddings remain tied.
+
+## Usage
+
+This repository contains custom Transformers architecture code and must be loaded with `trust_remote_code=True`.
+
+```bash
+pip install -U torch transformers safetensors
+```
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "BananaMind/BananaMind-2-Pro-Preview-Chat"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = (
+    torch.bfloat16
+    if device == "cuda" and torch.cuda.is_bf16_supported()
+    else torch.float32
+)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    dtype=dtype,
+).to(device).eval()
+
+messages = [
+    {
+        "role": "system",
+        "content": "You are a concise and helpful assistant.",
+    },
+    {
+        "role": "user",
+        "content": "Write a Python function that squares a number.",
+    },
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    return_tensors="pt",
+    return_dict=True,
+)
+inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
+
+with torch.inference_mode():
+    output = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        do_sample=False,
+        repetition_penalty=1.1,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+        use_cache=True,
+    )
+
+new_tokens = output[0, inputs["input_ids"].shape[1]:]
+print(tokenizer.decode(new_tokens, skip_special_tokens=True))
+```
+
+For multi-turn chat, append the generated assistant response and the next user message to `messages`, then render the full conversation with the chat template again.
+
+## Suggested Generation Settings
+
+For stable responses:
+
+- `do_sample=False`
+- `repetition_penalty=1.1`
+- `max_new_tokens=128` to `256`
+- `use_cache=True`
+
+For more varied responses:
+
+- `do_sample=True`
+- `temperature=0.6` to `0.8`
+- `top_p=0.9`
+- `repetition_penalty=1.1`
+- `max_new_tokens=128` to `256`
+
+Keep a finite generation limit. The combined prompt and generated output must fit within the 3,072-token context window.
+
+## Repository Files
+
+| File | Description |
+|---|---|
+| `config.json` | Transformers model configuration |
+| `model.safetensors` | Fine-tuned model weights |
+| `tokenizer.json` | Custom 32,768-token tokenizer |
+| `tokenizer_config.json` | Tokenizer metadata and context length |
+| `chat_template.jinja` | System/user/assistant chat template |
+| `generation_config.json` | Generation configuration |
+| `configuration_bananamind2pro.py` | Custom Transformers config class |
+| `modeling_bananamind2pro.py` | Custom Transformers causal LM class |
+| `sft_metadata.json` | Fine-tuning provenance and token counts |
+| `LICENSE` | BananaMind Community License 1.0 |
+| `banner.png` | Model-card banner |
+| `benchmarks.png` | Instruct Bench model comparison |
+| `score_progression.png` | Instruct Bench Elo progression across fine-tuning checkpoints |
+
+## Intended Use and Limitations
+
+BananaMind-2-Pro-Preview-Chat is intended for small-model research, local chat experiments, educational demonstrations, code-generation experiments, instruction-tuning studies, and compact-model comparisons.
+
+The model has not received dedicated safety alignment. It can hallucinate facts, fail arithmetic or logical tasks, produce insecure or invalid code, misunderstand instructions, and generate biased, repetitive, or otherwise undesirable text. Do not rely on it for medical, legal, financial, safety-critical, or other high-stakes decisions.
+
+## License
+
+This repository is released under the [BananaMind Community License 1.0](LICENSE). Commercial products or services exceeding either threshold in Section 1 require a separate commercial license from Banaxi-Tech.

chat_template.jinja

@@ -0,0 +1,10 @@
+{{ bos_token }}{% for message in messages %}{% if message['role'] == 'assistant' %}<|assistant|>
+{{ message['content'] }}{{ eos_token }}
+{% elif message['role'] == 'system' %}<|system|>
+{{ message['content'] }}
+{% elif message['role'] == 'user' %}<|user|>
+{{ message['content'] }}
+{% else %}<|{{ message['role'] }}|>
+{{ message['content'] }}
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|assistant|>
+{% endif %}

config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "BananaMind2ProForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_bananamind2pro.BananaMind2ProConfig",
+    "AutoModelForCausalLM": "modeling_bananamind2pro.BananaMind2ProForCausalLM"
+  },
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "head_dim": 80,
+  "hidden_size": 640,
+  "intermediate_size": 1920,
+  "max_position_embeddings": 3072,
+  "model_type": "bananamind2_pro",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 100000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.7.0",
+  "unk_token_id": 3,
+  "use_cache": true,
+  "vocab_size": 32768,
+  "z_loss_coeff": 0.0,
+  "_name_or_path": "BananaMind/BananaMind-2-Pro-Preview-Chat"
+}

configuration_bananamind2pro.py

@@ -0,0 +1,48 @@
+"""BananaMind 2 Pro configuration."""
+from transformers import PretrainedConfig
+
+
+class BananaMind2ProConfig(PretrainedConfig):
+    model_type = "bananamind2_pro"
+
+    def __init__(
+        self,
+        vocab_size=32768,
+        hidden_size=640,
+        num_hidden_layers=24,
+        num_attention_heads=8,
+        num_key_value_heads=4,
+        head_dim=80,
+        intermediate_size=1920,
+        max_position_embeddings=3072,
+        rope_theta=100000.0,
+        rms_norm_eps=1e-6,
+        tie_word_embeddings=True,
+        use_cache=True,
+        z_loss_coeff=0.0,
+        bos_token_id=1,
+        eos_token_id=2,
+        pad_token_id=0,
+        unk_token_id=3,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.z_loss_coeff = z_loss_coeff
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            unk_token_id=unk_token_id,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "5.7.0",
+  "use_cache": true
+}

model.safetensors

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

modeling_bananamind2pro.py

@@ -0,0 +1,307 @@
+"""BananaMind 2 Pro implementation for Hugging Face Transformers."""
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_bananamind2pro import BananaMind2ProConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x_float = x.float()
+        rms = torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x_float * rms * self.weight.float()).type_as(x)
+
+
+def build_rope_inv_freq(head_dim, theta=100000.0):
+    return 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+
+
+def precompute_freqs_cis(head_dim, seq_len, theta=100000.0):
+    freqs = build_rope_inv_freq(head_dim, theta)
+    positions = torch.arange(seq_len, dtype=torch.float32)
+    freqs = torch.outer(positions, freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+
+def apply_rotary_emb(q, k, freqs_cis):
+    q_complex = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2))
+    k_complex = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(0)
+    q_out = torch.view_as_real(q_complex * freqs_cis).flatten(-2)
+    k_out = torch.view_as_real(k_complex * freqs_cis).flatten(-2)
+    return q_out.type_as(q), k_out.type_as(k)
+
+
+class BananaMind2ProAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.n_rep = self.n_head // self.n_kv_heads
+
+        self.q_proj = nn.Linear(config.hidden_size, self.n_head * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.n_head * self.head_dim, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        x,
+        freqs_cis,
+        attention_mask=None,
+        past_key_values=None,
+        use_cache=False,
+    ):
+        batch_size, seq_len, _ = x.shape
+        q = self.q_proj(x).view(
+            batch_size,
+            seq_len,
+            self.n_head,
+            self.head_dim,
+        ).transpose(1, 2)
+        k = self.k_proj(x).view(
+            batch_size,
+            seq_len,
+            self.n_kv_heads,
+            self.head_dim,
+        ).transpose(1, 2)
+        v = self.v_proj(x).view(
+            batch_size,
+            seq_len,
+            self.n_kv_heads,
+            self.head_dim,
+        ).transpose(1, 2)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q, k = apply_rotary_emb(q, k, freqs_cis)
+
+        past_length = 0
+        if use_cache and past_key_values is not None:
+            past_length = past_key_values.get_seq_length(self.layer_idx)
+            k, v = past_key_values.update(k, v, self.layer_idx)
+
+        kv_len = k.size(-2)
+        k = k.unsqueeze(2).expand(
+            batch_size,
+            self.n_kv_heads,
+            self.n_rep,
+            kv_len,
+            self.head_dim,
+        ).reshape(batch_size, self.n_head, kv_len, self.head_dim)
+        v = v.unsqueeze(2).expand(
+            batch_size,
+            self.n_kv_heads,
+            self.n_rep,
+            kv_len,
+            self.head_dim,
+        ).reshape(batch_size, self.n_head, kv_len, self.head_dim)
+
+        attn_mask = None
+        is_causal = past_length == 0 and attention_mask is None
+        if not is_causal:
+            query_positions = past_length + torch.arange(seq_len, device=x.device)
+            key_positions = torch.arange(kv_len, device=x.device)
+            causal = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1)
+            attn_mask = causal[None, None, :, :]
+            if attention_mask is not None:
+                key_padding = attention_mask.to(torch.bool)
+                if key_padding.size(-1) < kv_len:
+                    cached_padding = torch.ones(
+                        key_padding.size(0),
+                        kv_len - key_padding.size(-1),
+                        dtype=torch.bool,
+                        device=key_padding.device,
+                    )
+                    key_padding = torch.cat((cached_padding, key_padding), dim=-1)
+                else:
+                    key_padding = key_padding[:, -kv_len:]
+                attn_mask = attn_mask & key_padding[:, None, None, :]
+            is_causal = False
+
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+        )
+        y = y.transpose(1, 2).contiguous().view(
+            batch_size,
+            seq_len,
+            self.n_head * self.head_dim,
+        )
+        return self.o_proj(y)
+
+
+class BananaMind2ProSwiGLUMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.w_gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w_up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w_down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+class BananaMind2ProBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.ln_1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attn = BananaMind2ProAttention(config, layer_idx)
+        self.ln_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = BananaMind2ProSwiGLUMLP(config)
+
+    def forward(
+        self,
+        x,
+        freqs_cis,
+        attention_mask=None,
+        past_key_values=None,
+        use_cache=False,
+    ):
+        x = x + self.attn(
+            self.ln_1(x),
+            freqs_cis,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+        )
+        return x + self.mlp(self.ln_2(x))
+
+
+class BananaMind2ProPreTrainedModel(PreTrainedModel):
+    config_class = BananaMind2ProConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        std = 0.02
+        if hasattr(module, "NANOGPT_SCALE_INIT"):
+            std *= 2 * self.config.num_hidden_layers ** -0.5
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+
+class BananaMind2ProForCausalLM(BananaMind2ProPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.transformer = nn.ModuleDict(
+            {
+                "wte": nn.Embedding(config.vocab_size, config.hidden_size),
+                "h": nn.ModuleList(
+                    [
+                        BananaMind2ProBlock(config, layer_idx)
+                        for layer_idx in range(config.num_hidden_layers)
+                    ]
+                ),
+                "ln_f": RMSNorm(config.hidden_size, eps=config.rms_norm_eps),
+            }
+        )
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.transformer["wte"].weight
+        self._embd_scale = math.sqrt(config.hidden_size)
+        self._freqs_cis_cache = None
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.transformer["wte"]
+
+    def set_input_embeddings(self, value):
+        self.transformer["wte"] = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def _get_freqs_cis(self, seq_len, device):
+        cache = self._freqs_cis_cache
+        if cache is None or cache.device != device or cache.size(0) < seq_len:
+            cache = precompute_freqs_cis(
+                self.config.head_dim,
+                seq_len,
+                self.config.rope_theta,
+            ).to(device)
+            self._freqs_cis_cache = cache
+        return cache[:seq_len]
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        labels=None,
+        past_key_values: Optional[Cache] = None,
+        use_cache=None,
+        **kwargs,
+    ):
+        _, seq_len = input_ids.shape
+        if use_cache is None:
+            use_cache = self.config.use_cache and labels is None
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        past_length = past_key_values.get_seq_length() if use_cache else 0
+        total_length = past_length + seq_len
+        if total_length > self.config.max_position_embeddings:
+            raise ValueError(
+                f"Sequence length {total_length} exceeds the configured maximum "
+                f"of {self.config.max_position_embeddings}"
+            )
+
+        x = self.transformer["wte"](input_ids) * self._embd_scale
+        freqs_cis = self._get_freqs_cis(total_length, input_ids.device)[past_length:]
+
+        for block in self.transformer["h"]:
+            x = block(
+                x,
+                freqs_cis,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+
+        x = self.transformer["ln_f"](x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.float().reshape(-1, shift_logits.size(-1)),
+                shift_labels.reshape(-1),
+            )
+            if self.config.z_loss_coeff:
+                loss = loss + self.config.z_loss_coeff * logits.float().pow(2).mean()
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values if use_cache else None,
+        )

sft_metadata.json

@@ -0,0 +1,12 @@
+{
+  "base_model": "BananaMind-2-Pro-Preview",
+  "dataset": "HuggingFaceTB/smol-smoltalk",
+  "dataset_config": "default",
+  "dataset_split": "train",
+  "training_type": "full_parameter_supervised_finetuning",
+  "loss_mask": "assistant_content_and_eos_only",
+  "global_step": 4014,
+  "tokens_seen": 394522611,
+  "supervised_tokens_seen": 301938934,
+  "repo_id": "BananaMind/BananaMind-2-Pro-Preview-Chat"
+}

tokenizer_config.json

@@ -0,0 +1,11 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "eos_token": "<|eos|>",
+  "is_local": true,
+  "local_files_only": true,
+  "model_max_length": 3072,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}