Update README.md

README.md

@@ -1,3 +1,147 @@
 ---
 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            | 58.43%   | 50.00%          |
+| ARC-Easy        | 37.29%   | 25.00%          |
+| ARC-Challenge   | 18.60%   | 25.00%          |
+| HellaSwag       | 26.81%   | 25.00%          |
+| ArithMark-2.0   | 27.64%   | 25.00%          |
+| ArithMark-3.0   | 32.80%   | 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**.
+
+## 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