--- 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 library_name: transformers --- # 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