README.md

---
license: apache-2.0
language:
- en
pipeline_tag: text-generation
tags:
- pebble
- language-model
- small-language-model
- pytorch
- safetensors
- custom-code
- mamba2
- hybrid
- chat
- sft
base_model:
- basically-ai/Pebble-10M
---

# Pebble-10M-Chat
![Banner](banner.png)

Pebble-10M-Chat is a compact, hybrid autoregressive chat language model. It combines the efficiency of state-space models with the proven performance of attention layers, optimized using a custom Muon + AdamW optimizer split.

## Model Details

- **Architecture:** Hybrid Mamba2 / Transformer
- **Block Pattern:** 3 Mamba2 blocks : 1 Attention block (repeating)
- **Parameters:** \~10,000,000 (10M)
- **Hidden Dimension:** 384
- **Layers:** 8 (6 Mamba2, 2 Attention)
- **Vocab Size:** 2,048 (Custom Byte-Level BPE)
- **Context Length:** 512
- **Pretraining Tokens:** \~25,000,000,000 (\~25 Billion)
- **Optimizer:** Muon (for 2D hidden weights) + AdamW (for embeddings, norms, and scalars)
- **Precision:** fp32 master weights with bf16 autocast

## Dataset Sources

The base model was pretrained on a 25B token subset of the following datasets:

| Dataset       | Token Allocation | Share |
|---------------|------------------|-------|
| FineWeb-Edu   | 7.50 billion     | 30%   |
| DCLM          | 5.00 billion     | 20%   |
| Cosmopedia-v2 | 3.75 billion     | 15%   |
| FineMath-4+   | 3.75 billion     | 15%   |
| FinePhrase    | 3.00 billion     | 12%   |
| NPset         | 2.00 billion     | 8%    |

## Benchmarks

Pebble-10M-Chat was evaluated on several commonsense and arithmetic benchmarks.

| Benchmark       | Accuracy | Random Baseline |
|-----------------|----------|-----------------|
| PIQA            | 51.41%   | 50.00%          |
| ARC-Easy        | 26.09%   | 25.00%          |
| ARC-Challenge   | 20.22%   | 25.00%          |
| HellaSwag       | 25.30%   | 25.00%          |
| ArithMark-2.0   | 27.28%   | 25.00%          |
| ArithMark-3.0   | 27.50%   | 25.00%          |

### Evaluation Notes

- PIQA, ARC-Easy, ARC-Challenge, and HellaSwag were evaluated on their respective test splits.
- ArithMark-2.0 was evaluated on its train split due to the lack of a suitable test split.
- ArithMark-3.0 was evaluated on its train split due to the lack of a suitable test split.
- Results were obtained using zero-shot multiple-choice evaluation.
- The model was additionally fine-tuned using supervised fine-tuning (SFT).

## SFT Attribution

The 250,000,000 SFT tokens used for Pebble-10M-Chat were provided by [smol-smoltalk](https://huggingface.co/datasets/HuggingFaceTB/smol-smoltalk).

## Usage

To run the model for text generation, you will need to install the required dependencies. The included Mamba2 implementation relies on CUDA/Triton kernels and is intended to run on a CUDA-enabled GPU. Ampere-class GPUs or newer are recommended.

> **Note:** The model uses custom architecture code, so you must pass \`trust_remote_code=True\` when loading both the tokenizer and the model.

```bash
pip install transformers huggingface_hub torch
pip install causal-conv1d mamba-ssm
```

Here is a simple Python script to load the model and generate text interactively:

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "basically-ai/Pebble-10M-Chat"

def main():
    print("Loading Pebble-10M-Chat...")
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        dtype=torch.float32,
    ).to("cuda")
    model.eval()

    print(f"Model loaded successfully! VRAM usage: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
    print("Type 'quit' or 'exit' to stop.\n")

    while True:
        prompt = input("You: ")
        if prompt.lower() in ["quit", "exit"]:
            break

        # Tokenize the prompt
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

        # Generate text
        print("Pebble: ", end="", flush=True)
        with torch.inference_mode():
            outputs = model.generate(
                **inputs,
                max_new_tokens=100,       # How many tokens to generate
                do_sample=True,           # Use sampling (more creative)
                temperature=0.7,          # Controls randomness
                top_k=50,                 # Consider top 50 tokens
                top_p=0.95,               # Nucleus sampling
                repetition_penalty=1.2,   # Prevent repeating words
            )

        # Decode and print (skip the prompt part)
        generated_text = tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True,
        )
        print(generated_text)
        print()

if __name__ == "__main__":
    main()
```

## License

Apache 2.0