BananaMind-2-Medium-Chat

BananaMind-2-Medium-Chat is the instruction-tuned version of BananaMind-2-Medium. It was fully fine-tuned on HuggingFaceTB/smol-smoltalk using assistant-only loss masking.

The model has 49,559,552 parameters, a 3,072-token context window, and a custom 12,288-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.

Model Details

Instruction Tuning

System and user messages were provided as context but masked from the loss. Only assistant content and the assistant-ending EOS token contributed to the training objective. All model parameters were trainable; no adapters or LoRA modules were used.

Chat Template

The tokenizer includes a Jinja chat template for system, user, and assistant messages. Its rendered structure is:

<BOS><|system|>
{system message}
<|user|>
{user message}
<|assistant|>
{assistant response}<EOS>

The role markers are plain text encoded with the existing tokenizer. No additional vocabulary entries were introduced during fine-tuning.

Usage

This model uses custom architecture code, so load it with trust_remote_code=True.

pip install -U transformers safetensors torch

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "BananaMind/BananaMind-2-Medium-Chat"

tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True,
)

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = (
    torch.bfloat16
    if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    else torch.float32
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_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,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_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 each generated assistant response to messages, add the next user message, and apply the chat template again.

Suggested Generation Settings

For stable responses:

• do_sample=False
• repetition_penalty=1.1
• max_new_tokens=128 to 256

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

Always keep a finite generation limit. Small models can enter repetitive continuations on difficult or contradictory prompts.

Benchmarks

Self-reported scores using lm_eval, the official ArithMark 2.0 script, and the internal BananaMind Instruction Bench. Scores may vary slightly by evaluation setup.

ARC Easy, ARC Challenge, PIQA, and HellaSwag use acc_norm,none.

BananaMind Instruction Bench (Internal)

* BananaMind Instruction Bench contains 100 examples: 40 General, 25 Multi-turn, 20 System Prompts, 10 Context Recall, and 5 Code. Supra 1.5 50M Instruct received conversation context through its ### Input field. Supra 1.0 50M Instruct used a copied checkpoint with 3.0x linear RoPE scaling, extending its context from 1,024 to 3,072 tokens without additional long-context training so it could fit the benchmark context.

Repository Files

Intended Use

This model is intended for small-model research, local chat experiments, instruction-tuning studies, educational demonstrations, and comparisons of compact language models.

License

Apache 2.0

Benchmark Average Formula

Average = (HellaSwag + ((ARC Easy + ARC Challenge) / 2) + PIQA + ArithMark 2.0) / 4